fix(ws): gate the identity sent in connection_init

The upgrade request passed the ban and rate-limit checks on its own headers. The backend then acted as the identity in the connection_init payload. A banned user could connect without credentials and send their token in the payload. The proxy now runs the same checks on the payload identity.

README.md

@@ -821,7 +821,7 @@ Enable rate limiting based on user roles using the `ROLE_RATE_LIMIT` environment
 - `ROLE_RATE_LIMIT`: Enable role-based rate limiting (default: `false`)
 
 **Features:**
-- **Burst control**: Optional burst limits for handling traffic spikes
+- **Burst control**: Optional `burst` caps requests in a short window (one tenth of the interval). `req` counts requests per `interval`. A `burst` at or above `req` has no effect.
 
 `ratelimit.json` is loaded once at startup; changing the file requires a restart.
 

api.go

@@ -143,8 +143,13 @@ func periodicallyReloadBannedUsers(ctx context.Context) {
 	}
 }
 
-func checkIfUserIsBanned(c fiber.Ctx, userID string) bool {
+func isUserBanned(userID string) bool {
 	_, found := bannedUsersIDs.Load(userID)
+	return found
+}
+
+func checkIfUserIsBanned(c fiber.Ctx, userID string) bool {
+	found := isUserBanned(userID)
 
 	cfg.Logger.Debug(&libpack_logger.LogMessage{
 		Message: "Checking if user is banned",

backend_health.go

@@ -118,7 +118,7 @@ func (bhm *BackendHealthManager) WaitForBackendReady(timeout time.Duration) erro
 		}
 
 		// Exponential backoff (1.5x growth, capped at maxDelay), no jitter
-		time.Sleep(currentDelay)
+		time.Sleep(min(currentDelay, time.Until(deadline)))
 		currentDelay = time.Duration(float64(currentDelay) * 1.5)
 		if currentDelay > maxDelay {
 			currentDelay = maxDelay

backend_health_deadline_test.go

@@ -0,0 +1,32 @@
+package main
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
+	"github.com/valyala/fasthttp"
+)
+
+// Pins theorems/Theorems/ReadyBackoff.lean: the startup wait never sleeps past its timeout.
+func TestWaitForBackendReady_HonoursTimeout(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	addr := ln.Addr().String()
+	_ = ln.Close() // closed port: every probe is refused immediately
+
+	bhm := NewBackendHealthManager(&fasthttp.Client{}, "http://"+addr, "", libpack_logger.New())
+	defer bhm.cancel()
+
+	const timeout = 3 * time.Second
+	start := time.Now()
+	if err := bhm.WaitForBackendReady(timeout); err == nil {
+		t.Fatal("want error for unreachable backend")
+	}
+	if elapsed := time.Since(start); elapsed > timeout+500*time.Millisecond {
+		t.Fatalf("waited %v, want <= %v", elapsed, timeout+500*time.Millisecond)
+	}
+}

cache/memory/memory.go

@@ -158,7 +158,8 @@ func (c *Cache) Set(key string, value []byte, ttl time.Duration) {
 		c.evictToFreeMemory(memoryToFree)
 	} else if atomic.LoadInt64(&c.entryCount) >= c.maxCacheSize {
 		// Fall back to count-based eviction for backward compatibility
-		c.evictOldest(int(c.maxCacheSize / 10)) // Evict 10% of entries
+		// 10% of entries, at least one: max/10 is 0 for caps below 10.
+		c.evictOldest(int(max(1, c.maxCacheSize/10)))
 	}
 
 	// A non-positive TTL is clamped to defaultTTL so this backend expires

cache/memory/small_cap_test.go

@@ -0,0 +1,23 @@
+package libpack_cache_memory
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// Pins theorems/Theorems/MemCacheCap.lean: every cap >= 1 bounds the entry count.
+func TestSet_SmallEntryCapIsEnforced(t *testing.T) {
+	for _, maxEntries := range []int64{1, 2, 5, 9, 10, 25} {
+		t.Run(fmt.Sprintf("cap_%d", maxEntries), func(t *testing.T) {
+			c := NewWithSize(time.Minute, DefaultMaxMemorySize, maxEntries)
+			defer c.Shutdown()
+			for i := 0; i < 100; i++ {
+				c.Set(fmt.Sprintf("k%d", i), []byte("v"), time.Minute)
+				assert.LessOrEqual(t, c.CountQueries(), maxEntries)
+			}
+		})
+	}
+}

events.go

@@ -12,6 +12,9 @@ import (
 const (
 	initialDelay    = 60 * time.Second
 	cleanupInterval = 1 * time.Hour
+
+	// minRetentionDays guards against a cutoff at or after NOW(), which would delete all event history.
+	minRetentionDays = 1
 )
 
 // Use parameterized queries to prevent SQL injection
@@ -46,6 +49,10 @@ func enableHasuraEventCleaner(ctx context.Context) error {
 	logger := cfg.Logger
 	cfgMutex.RUnlock()
 
+	if clearOlderThan < minRetentionDays {
+		return fmt.Errorf("HASURA_EVENT_CLEANER_OLDER_THAN must be at least %d day, got %d: event cleaner not started", minRetentionDays, clearOlderThan)
+	}
+
 	logger.Info(&libpack_logger.LogMessage{
 		Message: "Event cleaner enabled",
 		Pairs:   map[string]any{"interval_in_days": clearOlderThan},

events_retention_test.go

@@ -0,0 +1,26 @@
+package main
+
+import (
+	"context"
+	"testing"
+
+	libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
+	"github.com/stretchr/testify/assert"
+)
+
+// Pins theorems/Theorems/EventCleaner.lean: a retention below one day would delete all event history.
+func TestEnableHasuraEventCleaner_RejectsNonPositiveRetention(t *testing.T) {
+	for _, days := range []int{0, -1, -30} {
+		cfgMutex.Lock()
+		cfg = &config{Logger: libpack_logger.New()}
+		cfg.HasuraEventCleaner.Enable = true
+		cfg.HasuraEventCleaner.EventMetadataDb = "postgres://fake:[email protected]:1/fake"
+		cfg.HasuraEventCleaner.ClearOlderThan = days
+		cfgMutex.Unlock()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		err := enableHasuraEventCleaner(ctx)
+		cancel()
+		assert.Error(t, err, "days=%d", days)
+	}
+}

ratelimit.go

@@ -304,35 +304,19 @@ func rateLimitedRequest(userID, userRole string) bool {
 
 // checkRateLimit performs the actual rate limit check.
 //
-// Two independent gates apply, both must pass:
+// Two gates apply, both must pass:
 //
-//  1. Sustained: tickerRate (requests measured over the full Interval) must
+//  1. Sustained: the request count over the full Interval must stay at or
-//     stay at or below Req. This is the only gate when Burst <= 0, and its
+//     below Req (Req is per Interval, not per second).
-//     comparison is unchanged from before the burst fix.
+//  2. Burst (only when Burst > 0): the request count within a short trailing
-//  2. Burst (only when roleConfig.burstCounter != nil, i.e. Burst > 0): the
+//     window (see burstWindowDuration) must stay at or below Burst.
-//     raw request count within a short trailing window (see
-//     burstWindowDuration), much shorter than Interval, must stay at or
-//     below Burst.
 //
-// Because the burst window is short relative to Interval, a client can
+// The burst window lies inside the Interval, so its count never exceeds the
-// legally land up to Burst requests in a tight cluster - more than a
+// sustained count. Burst >= Req therefore never rejects anything the sustained
-// steady Req-paced stream would deliver in that same short span - as long
+// gate admits; Burst only takes effect when it is below Req.
-// as the sustained gate still holds over the full Interval. That is what
-// makes burst > req a real spike allowance instead of a no-op. Burst is a
-// short-window ceiling layered ON TOP of the sustained Req rate, not a
-// replacement for it. burst < req is an unusual config: the burst gate
-// becomes the tighter, binding constraint whenever requests cluster in
-// time, but the sustained cap for evenly-spread traffic is still Req
-// alone. The pre-fix code compared the SAME sustained-window tickerRate
-// against both Burst and Req, so when burst < req its Burst check tripped
-// first and Burst was effectively the sustained cap too. This gate's
-// Req-only sustained check can therefore allow MORE through than the
-// pre-fix code did for evenly-spread traffic under a burst < req config.
-// In both cases the burst gate never allows more than Burst requests
-// within its own window (fail-safe).
 func checkRateLimit(userID, userRole string, roleConfig RateLimitConfig, endpoint string) bool {
 	roleConfig.RateCounterTicker.Incr(1)
-	tickerRate := roleConfig.RateCounterTicker.GetRate()
+	tickerRate := roleConfig.RateCounterTicker.Get()
 
 	logDetails := map[string]any{
 		"user_role":   userRole,
@@ -348,7 +332,7 @@ func checkRateLimit(userID, userRole string, roleConfig RateLimitConfig, endpoin
 		Pairs:   map[string]any{"log_details": logDetails},
 	})
 
-	if tickerRate > float64(roleConfig.Req) {
+	if tickerRate > int64(roleConfig.Req) {
 		cfg.Logger.Debug(&libpack_logger.LogMessage{
 			Message: "Rate limit exceeded",
 			Pairs:   map[string]any{"log_details": logDetails},

ratelimit_test.go

@@ -372,59 +372,26 @@ func Test_checkRateLimit_BurstSemantics(t *testing.T) {
 		}
 	})
 
-	t.Run("b_burst_above_req_allows_short_spike_sustained_still_capped", func(t *testing.T) {
+	t.Run("b_burst_above_req_is_noop_sustained_caps_at_req", func(t *testing.T) {
 		const (
 			interval = 10 * time.Second
 			req      = 3
-			burst    = 6 // burst > req
+			burst    = 6 // burst >= req: dominated by the sustained gate
 		)
 		rc := newBurstTestRoleConfig(interval, req, burst)
 		fc := newFakeClock(time.Now())
 		rc.burstCounter.now = fc.Now
 
-		// A rapid cluster of `burst` requests must all be allowed, even
+		for i := 1; i <= req; i++ {
-		// though burst (6) > req (3): the sustained gate only bounds the
-		// total over the full 10s Interval (tickerRate <= 3 needs count <=
-		// 30), so a small cluster like this never trips it.
-		for i := 1; i <= burst; i++ {
 			if !checkRateLimit("u", "role-b", rc, "") {
-				t.Fatalf("request %d in first burst: got denied, want allowed (<= Burst=%d)", i, burst)
+				t.Fatalf("request %d: got denied, want allowed (<= Req=%d per Interval)", i, req)
 			}
 		}
 
-		// The (burst+1)th request in the SAME short window must be denied
+		// Clearing the burst window must not readmit: Req is per Interval.
-		// by the burst gate specifically - not by the sustained gate,
+		fc.Advance(2 * burstWindowDuration(interval))
-		// which alone would still allow it (tickerRate = 7/10 = 0.7 <= 3).
-		// This is the proof burst > req is not a no-op: it is the gate
-		// that actually fires here.
 		if checkRateLimit("u", "role-b", rc, "") {
-			t.Fatalf("request %d: got allowed, want denied by burst gate (> Burst=%d)", burst+1, burst)
+			t.Fatalf("request %d: got allowed, want denied by the sustained gate (Req=%d per Interval)", req+1, req)
-		}
-
-		// Roll the burst window forward (fake clock only - the real
-		// sustained ticker has not aged, since no real time passed) and
-		// repeat clusters of `burst` requests. Each cluster alone passes
-		// burst, but the cumulative sustained count keeps climbing until
-		// it exceeds req*interval.Seconds() = 30, at which point the
-		// sustained gate must deny regardless of how the burst window
-		// resets. This proves "sustained still capped at Req".
-		sustainedTripped := false
-		for batch := 0; batch < 10 && !sustainedTripped; batch++ {
-			fc.Advance(2 * burstWindowDuration(interval)) // clears the burst window
-			for i := 1; i <= burst; i++ {
-				allowed := checkRateLimit("u", "role-b", rc, "")
-				rate := rc.RateCounterTicker.GetRate()
-				if rate > float64(req) {
-					if allowed {
-						t.Fatalf("batch %d req %d: allowed with sustained rate %.2f > Req=%d", batch, i, rate, req)
-					}
-					sustainedTripped = true
-					break
-				}
-			}
-		}
-		if !sustainedTripped {
-			t.Fatalf("sustained gate never tripped after repeated bursts; expected it to cap total requests at Req over the Interval")
 		}
 	})
 

ratelimit_units_test.go

@@ -0,0 +1,44 @@
+package main
+
+import (
+	"testing"
+	"time"
+
+	libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
+)
+
+func setupRateLimitCfg() {
+	cfg = &config{}
+	parseConfig()
+	cfg.Logger = libpack_logger.New()
+}
+
+// Pins theorems/Theorems/RateLimitUnits.lean: Req counts requests per Interval, not per second.
+func TestCheckRateLimit_ReqIsPerInterval(t *testing.T) {
+	const req = 3
+	setupRateLimitCfg()
+	rc := newBurstTestRoleConfig(time.Minute, req, 0)
+	for i := 1; i <= req; i++ {
+		if !checkRateLimit("u", "r", rc, "") {
+			t.Fatalf("request %d denied, want allowed (Req=%d per minute)", i, req)
+		}
+	}
+	if checkRateLimit("u", "r", rc, "") {
+		t.Fatalf("request %d allowed, want denied: Req=%d per minute", req+1, req)
+	}
+}
+
+// With Burst >= Req the burst gate never rejects what the sustained gate admits.
+func TestCheckRateLimit_BurstAboveReqIsNoOp(t *testing.T) {
+	const req, burst = 3, 6
+	setupRateLimitCfg()
+	rc := newBurstTestRoleConfig(time.Minute, req, burst)
+	for i := 1; i <= req; i++ {
+		if !checkRateLimit("u", "r", rc, "") {
+			t.Fatalf("request %d denied", i)
+		}
+	}
+	if checkRateLimit("u", "r", rc, "") {
+		t.Fatalf("request %d allowed, want denied by the sustained gate", req+1)
+	}
+}

server.go

@@ -442,6 +442,11 @@ func processGraphQLRequest(c fiber.Ctx) error {
 // never an error; the request stays anonymous, exactly like before
 // verification existed.
 func extractUserInfo(c fiber.Ctx) (string, string, error) {
+	return extractUserInfoFromHeaders(func(key string) string { return c.Get(key) })
+}
+
+// extractUserInfoFromHeaders resolves user id and role from request headers via get.
+func extractUserInfoFromHeaders(get func(string) string) (string, string, error) {
 	extractedUserID := defaultValue
 	extractedRoleName := defaultValue
 
@@ -450,7 +455,7 @@ func extractUserInfo(c fiber.Ctx) (string, string, error) {
 	// Authorization header must still be checked even with no claim path
 	// set, otherwise JWT_VERIFY_SIGNATURE=true alone would never call
 	// Verify and the forged-token rejection below would never trigger.
-	if authorization := c.Get("Authorization"); authorization != "" &&
+	if authorization := get("Authorization"); authorization != "" &&
 		(len(cfg.Client.JWTUserClaimPath) > 0 || len(cfg.Client.JWTRoleClaimPath) > 0 || cfg.Client.JWTVerifier != nil) {
 		usr, role, err := extractClaimsFromJWTHeader(authorization)
 		if err != nil {
@@ -461,7 +466,7 @@ func extractUserInfo(c fiber.Ctx) (string, string, error) {
 
 	// Override role from header if configured
 	if cfg.Client.RoleFromHeader != "" {
-		if role := c.Get(cfg.Client.RoleFromHeader); role != "" {
+		if role := get(cfg.Client.RoleFromHeader); role != "" {
 			extractedRoleName = role
 		}
 	}

websocket.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net/http"
 	"sync"
@@ -231,6 +232,22 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web
 	// Try to extract headers from connection_init payload (for GraphQL WebSocket protocols)
 	enrichedHeaders := wsp.extractAuthFromPayload(message, headers)
 
+	// The upgrade request was gated on its own headers; the backend acts as the
+	// identity in enrichedHeaders, so gate that one too.
+	if authErr := authorizePayloadIdentity(headers, enrichedHeaders); authErr != nil {
+		wsp.errors.Add(1)
+		if wsp.logger != nil {
+			wsp.logger.Warning(&libpack_logger.LogMessage{
+				Message: "WebSocket connection_init identity rejected",
+				Pairs:   map[string]any{"connection_id": connectionID, "reason": authErr.Error()},
+			})
+		}
+		_ = clientConn.WriteMessage(websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.ClosePolicyViolation, authErr.Error()))
+		_ = clientConn.Close() // Best-effort cleanup
+		return
+	}
+
 	// Connect to backend WebSocket with enriched headers
 	backendConn, err := wsp.dialBackend(ctx, enrichedHeaders)
 	if err != nil {
@@ -586,6 +603,36 @@ func (wsp *WebSocketProxy) extractAuthFromPayload(message []byte, originalHeader
 	return enrichedHeaders
 }
 
+var (
+	errWSUnauthorized = errors.New("invalid or expired authorization token")
+	errWSBanned       = errors.New("user is banned")
+	errWSRateLimited  = errors.New("rate limit exceeded")
+)
+
+// authorizePayloadIdentity applies the ban and rate-limit checks of the upgrade
+// path to the identity that connection_init supplied. It does nothing when the
+// payload left Authorization and the role header unchanged, because the upgrade
+// request already passed those checks (and counted against the rate limit).
+func authorizePayloadIdentity(upgrade, enriched http.Header) error {
+	roleHeader := cfg.Client.RoleFromHeader
+	if upgrade.Get("Authorization") == enriched.Get("Authorization") &&
+		(roleHeader == "" || upgrade.Get(roleHeader) == enriched.Get(roleHeader)) {
+		return nil
+	}
+
+	userID, role, err := extractUserInfoFromHeaders(enriched.Get)
+	if err != nil {
+		return errWSUnauthorized
+	}
+	if isUserBanned(userID) {
+		return errWSBanned
+	}
+	if cfg.Client.RoleRateLimit && !rateLimitedRequest(userID, role) {
+		return errWSRateLimited
+	}
+	return nil
+}
+
 // dialBackend establishes a WebSocket connection to the backend
 func (wsp *WebSocketProxy) dialBackend(ctx context.Context, headers http.Header) (*gorillaws.Conn, error) {
 	// Convert http:// to ws:// or https:// to wss://

websocket_identity_test.go

@@ -0,0 +1,43 @@
+package main
+
+import (
+	"encoding/base64"
+	"net/http"
+	"testing"
+
+	libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
+	"github.com/stretchr/testify/assert"
+)
+
+func unsignedJWT(sub, role string) string {
+	enc := base64.RawURLEncoding.EncodeToString
+	return "Bearer " + enc([]byte(`{"alg":"none"}`)) + "." +
+		enc([]byte(`{"sub":"`+sub+`","role":"`+role+`"}`)) + "."
+}
+
+// Pins theorems/Theorems/WsIdentityGate.lean: the identity the backend acts as is the one that is gated.
+func TestAuthorizePayloadIdentity(t *testing.T) {
+	cfg = &config{Logger: libpack_logger.New()}
+	cfg.Client.JWTUserClaimPath = "sub"
+	cfg.Client.JWTRoleClaimPath = "role"
+	bannedUsersIDs.Store("banned-user", "test")
+	defer bannedUsersIDs.Delete("banned-user")
+
+	tests := []struct {
+		name     string
+		upgrade  http.Header
+		enriched http.Header
+		wantErr  error
+	}{
+		{"payload adds banned identity", http.Header{}, http.Header{"Authorization": {unsignedJWT("banned-user", "user")}}, errWSBanned},
+		{"payload replaces clean identity with banned", http.Header{"Authorization": {unsignedJWT("ok-user", "user")}}, http.Header{"Authorization": {unsignedJWT("banned-user", "user")}}, errWSBanned},
+		{"payload identity clean", http.Header{}, http.Header{"Authorization": {unsignedJWT("ok-user", "user")}}, nil},
+		{"payload does not change identity", http.Header{"Authorization": {unsignedJWT("ok-user", "user")}}, http.Header{"Authorization": {unsignedJWT("ok-user", "user")}}, nil},
+		{"no authorization anywhere", http.Header{}, http.Header{}, nil},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.ErrorIs(t, authorizePayloadIdentity(tt.upgrade, tt.enriched), tt.wantErr)
+		})
+	}
+}