diff --git a/internal/middleware/ratelimit.go b/internal/middleware/ratelimit.go index 5a09e5c..698ffaf 100644 --- a/internal/middleware/ratelimit.go +++ b/internal/middleware/ratelimit.go @@ -1,8 +1,10 @@ package middleware import ( + "context" "crypto/sha256" "encoding/hex" + "errors" "fmt" "log/slog" "net" @@ -39,6 +41,10 @@ func RateLimit(store kv.Store, cfg config.RateLimit, logger *slog.Logger) Middle key := rateLimitKey(ip, cfg.Window, now) count, err := store.Increment(r.Context(), key, cfg.Window+time.Second) if err != nil { + if errors.Is(err, context.Canceled) { + trace.SpanFromContext(r.Context()).SetAttributes(observability.RateLimitAttrs{Outcome: observability.RateLimitOutcomeCanceled, Limit: cfg.Requests, Window: cfg.Window}.Attributes()...) + return + } trace.SpanFromContext(r.Context()).SetAttributes(observability.RateLimitAttrs{Outcome: observability.RateLimitOutcomeStoreError, Limit: cfg.Requests, Window: cfg.Window}.Attributes()...) logger.ErrorContext(r.Context(), "rate limit check failed", slog.Any("error", err), slog.String("client_ip", ip)) http.Error(w, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable) diff --git a/internal/middleware/ratelimit_test.go b/internal/middleware/ratelimit_test.go index 784a57d..3f9b721 100644 --- a/internal/middleware/ratelimit_test.go +++ b/internal/middleware/ratelimit_test.go @@ -3,9 +3,11 @@ package middleware import ( "context" "errors" + "log/slog" "net/http" "net/http/httptest" "strconv" + "sync/atomic" "testing" "time" @@ -78,6 +80,33 @@ func TestRateLimitFailsClosedOnStoreError(t *testing.T) { } } +func TestRateLimitDoesNotLogCanceledRequestsAsStoreErrors(t *testing.T) { + t.Parallel() + + var errorLogs atomic.Int64 + logger := slog.New(countingErrorHandler{count: &errorLogs}) + store := incrementFunc(func(ctx context.Context, _ string, _ time.Duration) (int64, error) { + return 0, ctx.Err() + }) + handler := RateLimit(store, config.RateLimit{Enabled: true, Requests: 2, Window: time.Minute}, logger)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.test/mcp", nil) + if err != nil { + t.Fatal(err) + } + req.RemoteAddr = "203.0.113.10:1234" + + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + if got := errorLogs.Load(); got != 0 { + t.Fatalf("error logs = %d, want 0", got) + } +} + func TestClientIPPrefersFlyHeaderAndNormalizes(t *testing.T) { t.Parallel() @@ -114,3 +143,28 @@ func (f incrementFunc) Set(context.Context, string, []byte, time.Duration) error func (f incrementFunc) Increment(ctx context.Context, key string, ttl time.Duration) (int64, error) { return f(ctx, key, ttl) } + +type countingErrorHandler struct { + count *atomic.Int64 +} + +var _ slog.Handler = countingErrorHandler{} + +func (h countingErrorHandler) Enabled(_ context.Context, level slog.Level) bool { + return level >= slog.LevelError +} + +func (h countingErrorHandler) Handle(_ context.Context, record slog.Record) error { + if record.Level >= slog.LevelError { + h.count.Add(1) + } + return nil +} + +func (h countingErrorHandler) WithAttrs([]slog.Attr) slog.Handler { + return h +} + +func (h countingErrorHandler) WithGroup(string) slog.Handler { + return h +} diff --git a/internal/observability/attrs.go b/internal/observability/attrs.go index 383a558..fc5b495 100644 --- a/internal/observability/attrs.go +++ b/internal/observability/attrs.go @@ -128,6 +128,7 @@ const ( RateLimitOutcomeAllowed RateLimitOutcome = "allowed" RateLimitOutcomeLimited RateLimitOutcome = "limited" RateLimitOutcomeStoreError RateLimitOutcome = "store_error" + RateLimitOutcomeCanceled RateLimitOutcome = "canceled" RateLimitOutcomeDisabled RateLimitOutcome = "disabled" RateLimitOutcomeSkipped RateLimitOutcome = "skipped" )