From 4c5d01aef3670eb269a9afe92db3a8b811185d33 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Tue, 3 Nov 2026 00:34:18 +0000 Subject: [PATCH] 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 | 7 +++++- server.go | 9 ++++++-- websocket.go | 47 ++++++++++++++++++++++++++++++++++++++ websocket_identity_test.go | 43 ++++++++++++++++++++++++++++++++++ 4 files changed, 103 insertions(+), 3 deletions(-) create mode 100644 websocket_identity_test.go diff --git a/api.go b/api.go index 01af368..64e54d0 100644 --- a/api.go +++ b/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", diff --git a/server.go b/server.go index 355cf80..69d3707 100644 --- a/server.go +++ b/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 } } diff --git a/websocket.go b/websocket.go index b409228..3e7baa5 100644 --- a/websocket.go +++ b/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:// diff --git a/websocket_identity_test.go b/websocket_identity_test.go new file mode 100644 index 0000000..8b71fca --- /dev/null +++ b/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) + }) + } +}