Files
graphql-monitoring-proxy/websocket.go
T
lukaszraczylo 4c5d01aef3 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.
2026-11-03 00:34:18 +00:00

715 lines
24 KiB
Go

package main
import (
"context"
"errors"
"fmt"
"net/http"
"sync"
"sync/atomic"
"time"
"github.com/goccy/go-json"
"github.com/gofiber/contrib/v3/websocket"
"github.com/gofiber/fiber/v3"
gorillaws "github.com/gorilla/websocket"
libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
libpack_monitoring "github.com/lukaszraczylo/graphql-monitoring-proxy/monitoring"
)
// WebSocketProxy handles WebSocket proxying for GraphQL subscriptions
type WebSocketProxy struct {
logger *libpack_logger.Logger
monitoring *libpack_monitoring.MetricsSetup
backendURL string
enabled bool
pingInterval time.Duration
pongTimeout time.Duration
maxMessageSize int64
// effectivePingInterval/effectivePongTimeout/keepaliveEnabled are the
// defaulted-and-clamped keepalive values, computed once in
// NewWebSocketProxy (see computeEffectiveKeepalive) and reused by every
// connection via effectiveKeepalive. Computing them once means any clamp
// warning is logged once per proxy instead of once per connection.
effectivePingInterval time.Duration
effectivePongTimeout time.Duration
keepaliveEnabled bool
// Statistics
activeConnections atomic.Int64
totalConnections atomic.Int64
messagesSent atomic.Int64
messagesReceived atomic.Int64
errors atomic.Int64
}
// WebSocketConfig holds WebSocket configuration
type WebSocketConfig struct {
Enabled bool
PingInterval time.Duration
PongTimeout time.Duration
MaxMessageSize int64
}
// NewWebSocketProxy creates a new WebSocket proxy
func NewWebSocketProxy(backendURL string, config WebSocketConfig, logger *libpack_logger.Logger, monitoring *libpack_monitoring.MetricsSetup) *WebSocketProxy {
if config.PingInterval == 0 {
config.PingInterval = 30 * time.Second
}
if config.PongTimeout == 0 {
config.PongTimeout = 60 * time.Second
}
if config.MaxMessageSize == 0 {
config.MaxMessageSize = 512 * 1024 // 512KB default
}
wsp := &WebSocketProxy{
logger: logger,
monitoring: monitoring,
backendURL: backendURL,
enabled: config.Enabled,
pingInterval: config.PingInterval,
pongTimeout: config.PongTimeout,
maxMessageSize: config.MaxMessageSize,
}
wsp.effectivePingInterval, wsp.effectivePongTimeout, wsp.keepaliveEnabled =
computeEffectiveKeepalive(wsp.pingInterval, wsp.pongTimeout, logger)
if logger != nil && config.Enabled {
logger.Info(&libpack_logger.LogMessage{
Message: "WebSocket proxy enabled",
Pairs: map[string]any{
"backend_url": backendURL,
"ping_interval": config.PingInterval,
"max_message_size": config.MaxMessageSize,
},
})
}
return wsp
}
// HandleWebSocket upgrades the connection and proxies WebSocket traffic
func (wsp *WebSocketProxy) HandleWebSocket(c fiber.Ctx) error {
if !wsp.enabled {
return fiber.NewError(fiber.StatusNotImplemented, "WebSocket support is disabled")
}
// Check if this is a WebSocket upgrade request
if !websocket.IsWebSocketUpgrade(c) {
return fiber.NewError(fiber.StatusUpgradeRequired, "WebSocket upgrade required")
}
// Capture headers from the upgrade request to forward to backend
headers := make(http.Header)
var subprotocols []string
for key, value := range c.Request().Header.All() {
keyStr := string(key)
// Capture subprotocol separately
if keyStr == "Sec-Websocket-Protocol" || keyStr == "Sec-WebSocket-Protocol" {
subprotocols = append(subprotocols, string(value))
}
// Forward important headers including WebSocket subprotocol
// Skip only connection-establishment headers that will be regenerated
if keyStr != "Connection" && keyStr != "Upgrade" &&
keyStr != "Sec-Websocket-Key" && keyStr != "Sec-Websocket-Version" &&
keyStr != "Sec-Websocket-Extensions" {
headers.Add(keyStr, string(value))
}
}
// Configure WebSocket with subprotocol support
config := websocket.Config{
Subprotocols: subprotocols,
}
return websocket.New(func(clientConn *websocket.Conn) {
// Use background context for long-lived WebSocket connections
// The original request context expires after the upgrade
wsp.handleConnection(context.Background(), clientConn, headers)
}, config)(c)
}
// computeEffectiveKeepalive derives the ping period and pong read-deadline
// timeout a WebSocket proxy should use, plus whether keepalive is enabled at
// all, from the (already-defaulted) configured pingInterval/pongTimeout. It
// runs once, from NewWebSocketProxy, and the result is stored on the proxy
// so effectiveKeepalive can hand it to every connection without
// recomputing - and without re-logging any clamp warning - per connection.
//
// Operators can opt out of keepalive entirely by configuring a NEGATIVE
// PingInterval or PongTimeout (WEBSOCKET_PING_INTERVAL / _PONG_TIMEOUT) - no
// ping goroutine runs and no read deadline is set, matching the
// pre-keepalive behaviour of blocking on ReadMessage until the peer or
// network closes the connection. A value of exactly zero does NOT disable
// keepalive here: NewWebSocketProxy rewrites a zero PingInterval/PongTimeout
// to the 30s/60s defaults before this function ever sees them (and main.go's
// env parsing defaults PING_INTERVAL/PONG_TIMEOUT to 30/60 for the same
// reason), so only a negative value can reach this function as a disable.
//
// When enabled, pongTimeout is clamped to at least 1x pingInterval. The
// ping goroutine's ticker actually fires at pingInterval/2 (two ping
// attempts per configured interval), so a peer already sees a ping well
// before one full pingInterval elapses; requiring pongTimeout >= pingInterval
// keeps that margin without forcing an operator who explicitly configures a
// tighter timeout to accept the old 2x floor (explicit override wins).
func computeEffectiveKeepalive(pingInterval, pongTimeout time.Duration, logger *libpack_logger.Logger) (effectivePingInterval, effectivePongTimeout time.Duration, enabled bool) {
if pingInterval <= 0 || pongTimeout <= 0 {
return 0, 0, false
}
if pongTimeout < pingInterval {
if logger != nil {
logger.Warn(&libpack_logger.LogMessage{
Message: "WebSocket pong_timeout too low for ping_interval; clamping to avoid reaping idle connections",
Pairs: map[string]any{
"configured_ping_interval": pingInterval,
"configured_pong_timeout": pongTimeout,
"effective_pong_timeout": pingInterval,
},
})
}
pongTimeout = pingInterval
}
return pingInterval, pongTimeout, true
}
// effectiveKeepalive returns the ping period and pong read-deadline timeout
// this proxy should use, plus whether keepalive is enabled at all. It is the
// single source of truth for these values: handleConnection calls it once
// per connection and passes the result to both the initial deadline setup
// and every deadline-refresh path, so they can never disagree. The values
// themselves are computed once per proxy, in NewWebSocketProxy - see
// computeEffectiveKeepalive for the opt-out and clamping rules.
func (wsp *WebSocketProxy) effectiveKeepalive() (pingInterval, pongTimeout time.Duration, enabled bool) {
return wsp.effectivePingInterval, wsp.effectivePongTimeout, wsp.keepaliveEnabled
}
// handleConnection manages a single WebSocket connection
func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *websocket.Conn, headers http.Header) {
connectionID := fmt.Sprintf("%p", clientConn)
startTime := time.Now()
wsp.activeConnections.Add(1)
wsp.totalConnections.Add(1)
defer wsp.activeConnections.Add(-1)
if wsp.logger != nil {
wsp.logger.Info(&libpack_logger.LogMessage{
Message: "WebSocket connection established",
Pairs: map[string]any{
"connection_id": connectionID,
"active_connections": wsp.activeConnections.Load(),
},
})
}
// Set message size limit
clientConn.SetReadLimit(wsp.maxMessageSize)
// Read first message to extract authentication from connection_init payload
// This bridges the gap between clients that send auth in payload vs Hasura expecting it in HTTP headers
messageType, message, err := clientConn.ReadMessage()
if err != nil {
wsp.errors.Add(1)
if wsp.logger != nil {
wsp.logger.Error(&libpack_logger.LogMessage{
Message: "Failed to read first message from client",
Pairs: map[string]any{
"connection_id": connectionID,
"error": err.Error(),
},
})
}
_ = clientConn.Close() // Best-effort cleanup
return
}
// 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 {
wsp.errors.Add(1)
if wsp.logger != nil {
wsp.logger.Error(&libpack_logger.LogMessage{
Message: "Failed to connect to backend WebSocket",
Pairs: map[string]any{
"connection_id": connectionID,
"error": err.Error(),
},
})
}
_ = clientConn.Close() // Best-effort cleanup
return
}
defer func() { _ = backendConn.Close() }() // Best-effort cleanup
// Forward the first message (connection_init) to backend
if err := backendConn.WriteMessage(messageType, message); err != nil {
wsp.errors.Add(1)
if wsp.logger != nil {
wsp.logger.Error(&libpack_logger.LogMessage{
Message: "Failed to forward connection_init to backend",
Pairs: map[string]any{
"connection_id": connectionID,
"error": err.Error(),
},
})
}
return
}
if wsp.logger != nil {
wsp.logger.Debug(&libpack_logger.LogMessage{
Message: "Backend WebSocket connection established",
Pairs: map[string]any{
"connection_id": connectionID,
"subprotocol": backendConn.Subprotocol(),
"has_authorization": headers.Get("Authorization") != "",
},
})
}
// Bound messages read from the backend as well, so a large or
// misbehaving backend cannot force us to forward an unbounded frame to the
// client (the client-side limit only bounds what the client can send us).
backendConn.SetReadLimit(wsp.maxMessageSize)
// Keepalive: wire the configured (previously unused) PingInterval /
// PongTimeout so a half-open peer (TCP still open but silent) is reaped
// instead of leaving both proxy goroutines (and this handler) blocked in
// ReadMessage forever. We ping both directions and enforce a read
// deadline on each side that is refreshed by real traffic or a pong, so
// an active connection is never closed while an idle one is cleaned up.
// See effectiveKeepalive for the opt-out and clamping rules; the values
// it returns are threaded through unchanged to every deadline-refresh
// path below so they can never disagree.
pingInterval, pongTimeout, keepaliveEnabled := wsp.effectiveKeepalive()
if keepaliveEnabled {
_ = clientConn.SetReadDeadline(time.Now().Add(pongTimeout))
clientConn.SetPongHandler(func(string) error {
return clientConn.SetReadDeadline(time.Now().Add(pongTimeout))
})
_ = backendConn.SetReadDeadline(time.Now().Add(pongTimeout))
backendConn.SetPongHandler(func(string) error {
return backendConn.SetReadDeadline(time.Now().Add(pongTimeout))
})
}
// Ping both peers periodically. WriteControl is safe to call concurrently
// with WriteMessage, so this goroutine does not contend with the two
// forwarding directions. Both libraries use the standard WebSocket
// opcode value for a ping (9).
//
// pingDone is closed when the goroutine returns. gofiber's contrib/v3
// websocket recycles clientConn into a sync.Pool (nilling its underlying
// conn) the instant this function returns, so the goroutine MUST be
// joined (cancelPing + <-pingDone below) before handleConnection
// returns - otherwise a still-running WriteControl call can nil-deref
// (crashing the process, since this goroutine has no recover) or write a
// stray ping onto an unrelated, newly-recycled connection.
pingCtx, cancelPing := context.WithCancel(ctx)
pingDone := make(chan struct{})
if keepaliveEnabled {
go func() {
defer close(pingDone)
ticker := time.NewTicker(pingInterval / 2)
defer ticker.Stop()
for {
select {
case <-pingCtx.Done():
return
case <-ticker.C:
writeDeadline := time.Now().Add(10 * time.Second)
_ = clientConn.WriteControl(gorillaws.PingMessage, nil, writeDeadline)
_ = backendConn.WriteControl(gorillaws.PingMessage, nil, writeDeadline)
}
}
}()
} else {
close(pingDone)
}
// Set up bidirectional proxying
var wg sync.WaitGroup
wg.Add(2)
// Client -> Backend
go func() {
defer wg.Done()
// When the client side finishes (e.g. the client disconnected), close
// the backend so the reverse goroutine's blocked read returns instead
// of leaking the goroutine and holding the backend connection open.
defer func() { _ = backendConn.Close() }()
wsp.proxyClientToBackend(ctx, clientConn, backendConn, connectionID, pongTimeout)
}()
// Backend -> Client
go func() {
defer wg.Done()
// Symmetric teardown: when the backend side finishes, close the
// client so the client-direction goroutine isn't left blocked either.
defer func() { _ = clientConn.Close() }()
wsp.proxyBackendToClient(ctx, backendConn, clientConn, connectionID, pongTimeout)
}()
// Wait for both directions to complete
wg.Wait()
// Stop the ping goroutine and wait for it to actually exit before this
// handler returns - see the pingDone comment above for why this join is
// required, not optional.
cancelPing()
<-pingDone
duration := time.Since(startTime)
if wsp.logger != nil {
wsp.logger.Info(&libpack_logger.LogMessage{
Message: "WebSocket connection closed",
Pairs: map[string]any{
"connection_id": connectionID,
"duration_seconds": duration.Seconds(),
"messages_sent": wsp.messagesSent.Load(),
"messages_received": wsp.messagesReceived.Load(),
},
})
}
if wsp.monitoring != nil {
wsp.monitoring.Update("graphql_proxy_websocket_connection_duration", nil, duration.Seconds())
}
}
// proxyClientToBackend proxies messages from client to backend. pongTimeout
// is the effective (defaulted+clamped) value computed once by
// handleConnection via effectiveKeepalive; a value of 0 means keepalive is
// disabled and no read deadline is refreshed here.
func (wsp *WebSocketProxy) proxyClientToBackend(ctx context.Context, client *websocket.Conn, backend *gorillaws.Conn, connectionID string, pongTimeout time.Duration) {
for {
select {
case <-ctx.Done():
return
default:
messageType, message, err := client.ReadMessage()
if err != nil {
if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
if wsp.logger != nil {
wsp.logger.Debug(&libpack_logger.LogMessage{
Message: "Client WebSocket closed normally",
Pairs: map[string]any{
"connection_id": connectionID,
},
})
}
} else {
wsp.errors.Add(1)
if wsp.logger != nil {
wsp.logger.Error(&libpack_logger.LogMessage{
Message: "Error reading from client WebSocket",
Pairs: map[string]any{
"connection_id": connectionID,
"error": err.Error(),
},
})
}
}
return
}
wsp.messagesSent.Add(1)
// Refresh the read deadline on real client traffic so a client that
// sends data (even without pongs) is never reaped by the keepalive.
if pongTimeout > 0 {
_ = client.SetReadDeadline(time.Now().Add(pongTimeout))
}
// Forward message to backend
if err := backend.WriteMessage(messageType, message); err != nil {
wsp.errors.Add(1)
if wsp.logger != nil {
wsp.logger.Error(&libpack_logger.LogMessage{
Message: "Error writing to backend WebSocket",
Pairs: map[string]any{
"connection_id": connectionID,
"error": err.Error(),
},
})
}
return
}
if wsp.logger != nil {
wsp.logger.Debug(&libpack_logger.LogMessage{
Message: "Message proxied to backend",
Pairs: map[string]any{
"connection_id": connectionID,
"message_type": messageType,
"message_size": len(message),
},
})
}
}
}
}
// proxyBackendToClient proxies messages from backend to client. pongTimeout
// is the effective (defaulted+clamped) value computed once by
// handleConnection via effectiveKeepalive; a value of 0 means keepalive is
// disabled and no read deadline is refreshed here.
func (wsp *WebSocketProxy) proxyBackendToClient(ctx context.Context, backend *gorillaws.Conn, client *websocket.Conn, connectionID string, pongTimeout time.Duration) {
for {
select {
case <-ctx.Done():
return
default:
messageType, message, err := backend.ReadMessage()
if err != nil {
if gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
if wsp.logger != nil {
wsp.logger.Debug(&libpack_logger.LogMessage{
Message: "Backend WebSocket closed normally",
Pairs: map[string]any{
"connection_id": connectionID,
},
})
}
} else {
wsp.errors.Add(1)
if wsp.logger != nil {
wsp.logger.Error(&libpack_logger.LogMessage{
Message: "Error reading from backend WebSocket",
Pairs: map[string]any{
"connection_id": connectionID,
"error": err.Error(),
},
})
}
}
return
}
wsp.messagesReceived.Add(1)
// Refresh the backend read deadline on real traffic too: a backend
// that streams data but does not answer control pings must not be
// reaped by the keepalive.
if pongTimeout > 0 {
_ = backend.SetReadDeadline(time.Now().Add(pongTimeout))
}
// Forward message to client
if err := client.WriteMessage(messageType, message); err != nil {
wsp.errors.Add(1)
if wsp.logger != nil {
wsp.logger.Error(&libpack_logger.LogMessage{
Message: "Error writing to client WebSocket",
Pairs: map[string]any{
"connection_id": connectionID,
"error": err.Error(),
},
})
}
return
}
if wsp.logger != nil {
wsp.logger.Debug(&libpack_logger.LogMessage{
Message: "Message proxied to client",
Pairs: map[string]any{
"connection_id": connectionID,
"message_type": messageType,
"message_size": len(message),
},
})
}
}
}
}
// extractAuthFromPayload extracts authentication headers from GraphQL WebSocket connection_init payload
// This bridges the gap between clients sending auth in payload and Hasura expecting it in HTTP headers
func (wsp *WebSocketProxy) extractAuthFromPayload(message []byte, originalHeaders http.Header) http.Header {
// Create a copy of original headers
enrichedHeaders := make(http.Header)
for k, v := range originalHeaders {
enrichedHeaders[k] = v
}
// Try to parse as JSON to extract headers from payload
var msg map[string]any
if err := json.Unmarshal(message, &msg); err != nil {
// Not JSON or parse error, return original headers
return enrichedHeaders
}
// Check if this is a connection_init message
msgType, ok := msg["type"].(string)
if !ok || (msgType != "connection_init" && msgType != "start") {
// Not a connection_init, return original headers
return enrichedHeaders
}
// Extract payload
payload, ok := msg["payload"].(map[string]any)
if !ok {
return enrichedHeaders
}
// Try to extract headers from payload.headers (graphql-ws format)
if payloadHeaders, ok := payload["headers"].(map[string]any); ok {
for key, value := range payloadHeaders {
if strValue, ok := value.(string); ok {
enrichedHeaders.Set(key, strValue)
}
}
}
// Also check top-level payload keys that look like headers (Apollo format)
for key, value := range payload {
if strValue, ok := value.(string); ok {
// Common auth headers
if key == "Authorization" || key == "authorization" ||
key == "x-hasura-role" || key == "x-hasura-admin-secret" {
enrichedHeaders.Set(key, strValue)
}
}
}
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://
wsURL := wsp.backendURL
if len(wsURL) > 7 && wsURL[:7] == "http://" {
wsURL = "ws://" + wsURL[7:]
} else if len(wsURL) > 8 && wsURL[:8] == "https://" {
wsURL = "wss://" + wsURL[8:]
}
// Append GraphQL WebSocket path
wsURL = wsURL + "/v1/graphql"
// Extract subprotocols from headers (e.g., graphql-ws, graphql-transport-ws)
var subprotocols []string
if proto := headers.Get("Sec-WebSocket-Protocol"); proto != "" {
subprotocols = []string{proto}
// Remove from headers since it will be set via Subprotocols field
headers.Del("Sec-WebSocket-Protocol")
}
// Use gorilla websocket dialer
dialer := gorillaws.Dialer{
HandshakeTimeout: 10 * time.Second,
Subprotocols: subprotocols,
}
// Dial the backend with forwarded headers
conn, _, err := dialer.DialContext(ctx, wsURL, headers)
if err != nil {
return nil, fmt.Errorf("failed to dial backend WebSocket: %w", err)
}
return conn, nil
}
// GetStats returns WebSocket statistics. ping_interval/pong_timeout report
// the effective (defaulted-and-clamped) keepalive values actually in use -
// see computeEffectiveKeepalive - not the raw configured fields, so callers
// see what handleConnection really does with idle connections.
func (wsp *WebSocketProxy) GetStats() map[string]any {
return map[string]any{
"enabled": wsp.enabled,
"active_connections": wsp.activeConnections.Load(),
"total_connections": wsp.totalConnections.Load(),
"messages_sent": wsp.messagesSent.Load(),
"messages_received": wsp.messagesReceived.Load(),
"errors": wsp.errors.Load(),
"ping_interval": wsp.effectivePingInterval.String(),
"pong_timeout": wsp.effectivePongTimeout.String(),
"max_message_size": wsp.maxMessageSize,
}
}
// IsWebSocketRequest checks if the request is a WebSocket upgrade request
func IsWebSocketRequest(c fiber.Ctx) bool {
return websocket.IsWebSocketUpgrade(c) ||
c.Get("Upgrade") == "websocket" ||
c.Get("Connection") == "Upgrade"
}
// Global WebSocket proxy
var (
webSocketProxy *WebSocketProxy
webSocketProxyOnce sync.Once
)
// InitializeWebSocketProxy initializes the global WebSocket proxy
func InitializeWebSocketProxy(backendURL string, config WebSocketConfig, logger *libpack_logger.Logger, monitoring *libpack_monitoring.MetricsSetup) *WebSocketProxy {
webSocketProxyOnce.Do(func() {
webSocketProxy = NewWebSocketProxy(backendURL, config, logger, monitoring)
})
return webSocketProxy
}
// GetWebSocketProxy returns the global WebSocket proxy
func GetWebSocketProxy() *WebSocketProxy {
return webSocketProxy
}