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.

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",

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)
+		})
+	}
+}