fix(monitors): rate limit heartbeats per token and kind; test limiter key and fail body cap

server/internal/api/heartbeats.go

@@ -74,17 +74,41 @@ func MaskLogPath(path string) string {
 	return heartbeatPathPrefix + "***" + rest[end:]
 }
 
-// heartbeatAllowed admits one request per token per second. A cron job pinging
+// heartbeatLimitKey buckets requests by token, kind and wall-clock second. The
+// kind is part of the key so a fast job's success ping is not rejected for
+// landing in the same second as its /start, which would leave the run marked
+// started and page as never finished. The token is hashed so Redis never holds
+// the secret.
+func heartbeatLimitKey(token, kind string, now time.Time) string {
+	return "vantage:hbrl:" + services.HashHeartbeatToken(token) + ":" + kind + ":" + strconv.FormatInt(now.Unix(), 10)
+}
+
+// readHeartbeatBody keeps up to MaxHeartbeatBody bytes of a /fail body as the
+// incident message. Other kinds carry no message, so their bodies are ignored.
+func readHeartbeatBody(kind string, r io.Reader) string {
+	if kind != services.HeartbeatFail || r == nil {
+		return ""
+	}
+	b, _ := io.ReadAll(io.LimitReader(r, services.MaxHeartbeatBody))
+	return string(b)
+}
+
+func writeHeartbeatLimited(c *gin.Context) {
+	c.Header("Retry-After", "1")
+	c.String(http.StatusTooManyRequests, "too many requests")
+}
+
+// heartbeatAllowed admits one request per token and kind per second. A cron job pinging
 // in a loop should not become a write per request, and a leaked URL should not
 // be a way to hammer Mongo. Like the status page limiter it allows when Redis
 // is down: a missed ping pages someone. It runs inside the handler rather than
 // as middleware because the token may come from a header.
-func heartbeatAllowed(c *gin.Context, token string) bool {
+func heartbeatAllowed(c *gin.Context, token, kind string) bool {
 	rdb := auth.Redis()
 	if rdb == nil {
 		return true
 	}
-	key := "vantage:hbrl:" + services.HashHeartbeatToken(token) + ":" + strconv.FormatInt(time.Now().Unix(), 10)
+	key := heartbeatLimitKey(token, kind, time.Now())
 	count, err := rdb.Incr(c.Request.Context(), key).Result()
 	if err != nil {
 		return true
@@ -108,17 +132,11 @@ func handleHeartbeat(c *gin.Context) {
 		c.String(http.StatusNotFound, "not found")
 		return
 	}
-	if !heartbeatAllowed(c, token) {
-		c.Header("Retry-After", "1")
-		c.String(http.StatusTooManyRequests, "too many requests")
+	if !heartbeatAllowed(c, token, kind) {
+		writeHeartbeatLimited(c)
 		return
 	}
-	var body string
-	if kind == services.HeartbeatFail && c.Request.Body != nil {
-		b, _ := io.ReadAll(io.LimitReader(c.Request.Body, services.MaxHeartbeatBody))
-		body = string(b)
-	}
-	err := services.RecordHeartbeat(token, kind, body, time.Now())
+	err := services.RecordHeartbeat(token, kind, readHeartbeatBody(kind, c.Request.Body), time.Now())
 	if errors.Is(err, services.ErrHeartbeatNotFound) {
 		c.String(http.StatusNotFound, "not found")
 		return

server/internal/api/heartbeats_test.go

@@ -1,6 +1,10 @@
 package api
 
 import (
+	"strings"
+	"time"
+
+	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/services"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -68,3 +72,41 @@ func TestHeartbeatUnresolvableIs404(t *testing.T) {
 		t.Fatalf("code = %d, want 404", w.Code)
 	}
 }
+
+func TestHeartbeatLimitKey(t *testing.T) {
+	now := time.Unix(1700000000, 0)
+	start := heartbeatLimitKey("tok", services.HeartbeatStart, now)
+	ping := heartbeatLimitKey("tok", services.HeartbeatPing, now)
+	if start == ping {
+		t.Fatalf("start and ping share a limiter key: %s", start)
+	}
+	if ping != heartbeatLimitKey("tok", services.HeartbeatPing, now.Add(500*time.Millisecond)) {
+		t.Fatal("same kind in the same second should share a key")
+	}
+	if strings.Contains(ping, "tok") {
+		t.Fatalf("key leaks the plaintext token: %s", ping)
+	}
+}
+
+func TestReadHeartbeatBody(t *testing.T) {
+	long := strings.Repeat("x", services.MaxHeartbeatBody+500)
+	if got := readHeartbeatBody(services.HeartbeatFail, strings.NewReader(long)); len(got) != services.MaxHeartbeatBody {
+		t.Fatalf("fail body not capped: %d", len(got))
+	}
+	if got := readHeartbeatBody(services.HeartbeatPing, strings.NewReader("hi")); got != "" {
+		t.Fatalf("ping body should be ignored, got %q", got)
+	}
+	if got := readHeartbeatBody(services.HeartbeatFail, nil); got != "" {
+		t.Fatalf("nil body: %q", got)
+	}
+}
+
+func TestHeartbeatRateLimitedWritesRetryAfter(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	writeHeartbeatLimited(c)
+	if w.Code != http.StatusTooManyRequests || w.Header().Get("Retry-After") != "1" {
+		t.Fatalf("got %d retry-after %q", w.Code, w.Header().Get("Retry-After"))
+	}
+}

web/e2e/heartbeat.spec.ts

@@ -26,7 +26,6 @@ test("heartbeat monitor: URL ping, /fail, header ping", async ({ page, request }
     let got = await (await page.request.get(`${BASE_URL}/api/monitors/${monitor.monitor_id}`)).json();
     expect(got.state.status).toBe("up");
 
-    await new Promise((r) => setTimeout(r, 1100)); // per-token limit is 1/s
     const fail = await request.post(`${BASE_URL}/public/hb/${monitor.heartbeat_token}/fail`, { data: "disk full" });
     expect(fail.status()).toBe(200);