Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions internal/middleware/ratelimit.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
package middleware

import (
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"log/slog"
"net"
Expand Down Expand Up @@ -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
}
Comment on lines +44 to +47

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 The cancellation guard only covers context.Canceled, but context.DeadlineExceeded surfaces through the same path when a client-set deadline expires before the store responds. Both represent the client's request lifecycle ending, not a store malfunction. Without this, a deadline-exceeded error from store.Increment would still be logged at error level and return a 503 to a client that has already gone. Consider broadening the check to errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded).

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)
Expand Down
54 changes: 54 additions & 0 deletions internal/middleware/ratelimit_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,11 @@ package middleware
import (
"context"
"errors"
"log/slog"
"net/http"
"net/http/httptest"
"strconv"
"sync/atomic"
"testing"
"time"

Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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
}
1 change: 1 addition & 0 deletions internal/observability/attrs.go
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ const (
RateLimitOutcomeAllowed RateLimitOutcome = "allowed"
RateLimitOutcomeLimited RateLimitOutcome = "limited"
RateLimitOutcomeStoreError RateLimitOutcome = "store_error"
RateLimitOutcomeCanceled RateLimitOutcome = "canceled"
RateLimitOutcomeDisabled RateLimitOutcome = "disabled"
RateLimitOutcomeSkipped RateLimitOutcome = "skipped"
)
Expand Down
Loading