diff --git a/internal/worker/ratelimit.go b/internal/worker/ratelimit.go index 0a49ead..0cc29e9 100644 --- a/internal/worker/ratelimit.go +++ b/internal/worker/ratelimit.go @@ -2,6 +2,7 @@ package worker import ( + "net" "net/http" "sync" "time" @@ -206,14 +207,13 @@ func (pcrl *PerClientRateLimiter) Stats() map[string]any { } // PerClientRateLimitMiddleware creates middleware that applies per-client rate limiting. -// Uses X-Forwarded-For or RemoteAddr to identify clients. +// Identifies clients by the peer IP of the connection; forwarding headers are spoofable and ignored. func PerClientRateLimitMiddleware(limiter *PerClientRateLimiter) func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Get client identifier (prefer X-Real-IP from RealIP middleware) clientKey := r.RemoteAddr - if xff := r.Header.Get("X-Real-IP"); xff != "" { - clientKey = xff + if host, _, err := net.SplitHostPort(clientKey); err == nil { + clientKey = host } if !limiter.Allow(clientKey) { diff --git a/internal/worker/ratelimit_test.go b/internal/worker/ratelimit_test.go new file mode 100644 index 0000000..ad95b11 --- /dev/null +++ b/internal/worker/ratelimit_test.go @@ -0,0 +1,35 @@ +package worker + +import ( + "net/http" + "net/http/httptest" + "testing" +) + +func TestPerClientRateLimitMiddleware_KeysOnPeerIP(t *testing.T) { + h := PerClientRateLimitMiddleware(NewPerClientRateLimiter(0.001, 1))( + http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })) + + do := func(remote, realIP string) int { + req := httptest.NewRequest(http.MethodGet, "/", nil) + req.RemoteAddr = remote + if realIP != "" { + req.Header.Set("X-Real-IP", realIP) + } + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + return rec.Code + } + + if got := do("10.0.0.1:1111", "1.1.1.1"); got != http.StatusOK { + t.Fatalf("first request: got %d", got) + } + // Same peer, new port and forged header: must share the exhausted bucket. + if got := do("10.0.0.1:2222", "2.2.2.2"); got != http.StatusTooManyRequests { + t.Fatalf("spoofed header/new port bypassed limit: got %d", got) + } + // Different peer is unaffected. + if got := do("10.0.0.2:1111", ""); got != http.StatusOK { + t.Fatalf("other peer: got %d", got) + } +} diff --git a/internal/worker/service.go b/internal/worker/service.go index 2d6627b..bc6e95a 100644 --- a/internal/worker/service.go +++ b/internal/worker/service.go @@ -1226,7 +1226,7 @@ func (s *Service) setupMiddleware() { // Add gzip compression for responses >1KB (reduces bandwidth ~70% for JSON) s.router.Use(middleware.Compress(5)) // Level 5 = good balance of speed vs compression - // Apply per-client rate limiting (after RealIP so we get the real client IP) + // Apply per-client rate limiting (keyed on the socket peer address; forwarding headers are untrusted) if s.rateLimiter != nil { s.router.Use(PerClientRateLimitMiddleware(s.rateLimiter)) }