diff --git a/server/internal/api/heartbeats.go b/server/internal/api/heartbeats.go index c93b166..af38edb 100644 --- a/server/internal/api/heartbeats.go +++ b/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 diff --git a/server/internal/api/heartbeats_test.go b/server/internal/api/heartbeats_test.go index 84b5906..5bb0790 100644 --- a/server/internal/api/heartbeats_test.go +++ b/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")) + } +} diff --git a/web/e2e/heartbeat.spec.ts b/web/e2e/heartbeat.spec.ts index 735038b..5c88975 100644 --- a/web/e2e/heartbeat.spec.ts +++ b/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);