diff --git a/websocket.go b/websocket.go index 8d5495c..b409228 100644 --- a/websocket.go +++ b/websocket.go @@ -26,6 +26,15 @@ type WebSocketProxy struct { 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 @@ -64,6 +73,9 @@ func NewWebSocketProxy(backendURL string, config WebSocketConfig, logger *libpac 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", @@ -120,6 +132,62 @@ func (wsp *WebSocketProxy) HandleWebSocket(c fiber.Ctx) error { }, 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) @@ -218,44 +286,55 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web // 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. - pingInterval := wsp.pingInterval - pongTimeout := wsp.pongTimeout - if pingInterval <= 0 { - pingInterval = 30 * time.Second - } - if pongTimeout <= 0 { - pongTimeout = 60 * time.Second - } + // 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() - _ = 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)) - }) + 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) - defer cancelPing() - go func() { - 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) + 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 @@ -268,7 +347,7 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web // 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) + wsp.proxyClientToBackend(ctx, clientConn, backendConn, connectionID, pongTimeout) }() // Backend -> Client @@ -277,12 +356,18 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web // 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) + 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 { @@ -302,8 +387,11 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web } } -// proxyClientToBackend proxies messages from client to backend -func (wsp *WebSocketProxy) proxyClientToBackend(ctx context.Context, client *websocket.Conn, backend *gorillaws.Conn, connectionID string) { +// 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(): @@ -339,8 +427,8 @@ func (wsp *WebSocketProxy) proxyClientToBackend(ctx context.Context, client *web // Refresh the read deadline on real client traffic so a client that // sends data (even without pongs) is never reaped by the keepalive. - if wsp.pongTimeout > 0 { - _ = client.SetReadDeadline(time.Now().Add(wsp.pongTimeout)) + if pongTimeout > 0 { + _ = client.SetReadDeadline(time.Now().Add(pongTimeout)) } // Forward message to backend @@ -372,8 +460,11 @@ func (wsp *WebSocketProxy) proxyClientToBackend(ctx context.Context, client *web } } -// proxyBackendToClient proxies messages from backend to client -func (wsp *WebSocketProxy) proxyBackendToClient(ctx context.Context, backend *gorillaws.Conn, client *websocket.Conn, connectionID string) { +// 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(): @@ -410,8 +501,8 @@ func (wsp *WebSocketProxy) proxyBackendToClient(ctx context.Context, backend *go // 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 wsp.pongTimeout > 0 { - _ = backend.SetReadDeadline(time.Now().Add(wsp.pongTimeout)) + if pongTimeout > 0 { + _ = backend.SetReadDeadline(time.Now().Add(pongTimeout)) } // Forward message to client @@ -531,7 +622,10 @@ func (wsp *WebSocketProxy) dialBackend(ctx context.Context, headers http.Header) return conn, nil } -// GetStats returns WebSocket statistics +// 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, @@ -540,8 +634,8 @@ func (wsp *WebSocketProxy) GetStats() map[string]any { "messages_sent": wsp.messagesSent.Load(), "messages_received": wsp.messagesReceived.Load(), "errors": wsp.errors.Load(), - "ping_interval": wsp.pingInterval.String(), - "pong_timeout": wsp.pongTimeout.String(), + "ping_interval": wsp.effectivePingInterval.String(), + "pong_timeout": wsp.effectivePongTimeout.String(), "max_message_size": wsp.maxMessageSize, } } diff --git a/websocket_test.go b/websocket_test.go index 55c7ab1..d0a49bc 100644 --- a/websocket_test.go +++ b/websocket_test.go @@ -2,12 +2,14 @@ package main import ( "context" + "net" "net/http" "net/http/httptest" "strings" "testing" "time" + fiber "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" @@ -310,6 +312,13 @@ func TestWebSocketProxy_ConfigValidation(t *testing.T) { }) } +// TestWebSocketProxy_ReapsSilentBackend_Keepalive is a lower-level unit test +// of proxyBackendToClient's read loop in isolation: given an initial read +// deadline (set here exactly as handleConnection sets it), a silent backend +// must be reaped instead of blocking forever. It does NOT exercise +// handleConnection's own keepalive wiring (effectiveKeepalive, the ping +// goroutine, or the deadline being set from configuration) - see +// TestWebSocketProxy_HandleConnection_KeepaliveReapsIdleClient for that. func TestWebSocketProxy_ReapsSilentBackend_Keepalive(t *testing.T) { config := WebSocketConfig{ Enabled: true, @@ -350,7 +359,7 @@ func TestWebSocketProxy_ReapsSilentBackend_Keepalive(t *testing.T) { // client conn is nil and safe here: the silent backend sends no // message, so proxyBackendToClient returns on the read deadline // before any client.WriteMessage call. - wsp.proxyBackendToClient(context.Background(), backend, nil, "test-conn") + wsp.proxyBackendToClient(context.Background(), backend, nil, "test-conn", config.PongTimeout) close(done) }() @@ -395,3 +404,177 @@ func TestWebSocketProxy_StatsStructure(t *testing.T) { assert.True(t, hasPong) assert.True(t, hasSize) } + +// startTestWebSocketProxyApp starts a real fiber app on an ephemeral +// loopback port serving wsp at /ws, so tests can dial it with a real +// WebSocket client and exercise handleConnection through its actual entry +// point (HandleWebSocket) instead of calling unexported methods directly. +// It registers app shutdown via t.Cleanup and returns the "host:port" to +// dial. +func startTestWebSocketProxyApp(t *testing.T, wsp *WebSocketProxy) string { + t.Helper() + + app := fiber.New() + app.Get("/ws", func(c fiber.Ctx) error { + return wsp.HandleWebSocket(c) + }) + + listenReady := make(chan struct{}) + var addr string + go func() { + _ = app.Listen("127.0.0.1:0", fiber.ListenConfig{ + DisableStartupMessage: true, + ListenerAddrFunc: func(a net.Addr) { + addr = a.String() + close(listenReady) + }, + }) + }() + + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = app.ShutdownWithContext(ctx) + }) + + select { + case <-listenReady: + case <-time.After(2 * time.Second): + t.Fatal("test WebSocket proxy app did not start listening") + } + + return addr +} + +// TestWebSocketProxy_HandleConnection_KeepaliveReapsIdleClient exercises the +// keepalive wiring INSIDE handleConnection end to end, through the real +// HandleWebSocket entry point. This closes the gap left by +// TestWebSocketProxy_ReapsSilentBackend_Keepalive above: that test sets the +// read deadline itself and calls proxyBackendToClient directly, so it passes +// even if handleConnection stopped calling effectiveKeepalive / setting any +// deadline at all. This test sets no deadline itself - if handleConnection's +// own keepalive wiring were removed, the idle client below would never be +// reaped and this test would time out. +func TestWebSocketProxy_HandleConnection_KeepaliveReapsIdleClient(t *testing.T) { + // Backend that completes the upgrade, consumes the forwarded + // connection_init, then goes silent - mirrors a slow/half-open backend. + serverRelease := make(chan struct{}) + upgrader := gorillaws.Upgrader{} + backendSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + _, _, _ = conn.ReadMessage() + <-serverRelease + _ = conn.Close() + })) + defer backendSrv.Close() + defer close(serverRelease) + + // PongTimeout >= 2x PingInterval so effectiveKeepalive's clamp does not + // stretch it past the tight value this test relies on to stay fast. + config := WebSocketConfig{ + Enabled: true, + PingInterval: 100 * time.Millisecond, + PongTimeout: 250 * time.Millisecond, + } + wsp := NewWebSocketProxy(backendSrv.URL, config, libpack_logger.New(), nil) + + proxyAddr := startTestWebSocketProxyApp(t, wsp) + + clientConn, _, err := gorillaws.DefaultDialer.Dial("ws://"+proxyAddr+"/ws", nil) + if err != nil { + t.Fatalf("dial proxy: %v", err) + } + defer clientConn.Close() + + // Send the connection_init handleConnection reads before wiring up the + // backend and the keepalive, then go silent: no more messages, so this + // client never hits the traffic-refresh path in proxyClientToBackend + // either - only handleConnection's own read deadline can reap it. + if err := clientConn.WriteMessage(gorillaws.TextMessage, []byte(`{"type":"connection_init"}`)); err != nil { + t.Fatalf("write connection_init: %v", err) + } + + readErrCh := make(chan error, 1) + go func() { + _, _, err := clientConn.ReadMessage() + readErrCh <- err + }() + + select { + case err := <-readErrCh: + assert.Error(t, err, "expected the proxy to close the idle connection") + case <-time.After(2 * time.Second): + t.Fatal("handleConnection did not reap the idle client via its own keepalive wiring") + } + + // wsp's own bookkeeping (decremented by handleConnection's defer right + // before it returns) must reflect the teardown too. + assert.Eventually(t, func() bool { + return wsp.activeConnections.Load() == 0 + }, 2*time.Second, 10*time.Millisecond, "handleConnection did not return after reaping the idle client") +} + +// TestWebSocketProxy_HandleConnection_PingGoroutineJoinsBeforeReturn is a +// regression test for the ping goroutine leak: handleConnection must cancel +// AND join (wait for) its ping goroutine before returning, because gofiber's +// contrib/v3 websocket recycles the *websocket.Conn wrapper into a +// sync.Pool - nilling its embedded conn - the instant the handler returns +// (see releaseConn in contrib/v3/websocket). A ping goroutine still running +// after that either nil-derefs inside WriteControl (an unrecovered panic in +// a goroutine with no recover, crashing the whole test binary) or writes a +// stray ping onto a different, newly-accepted connection's socket. +// +// This cannot be proven with a single deterministic assertion: the failure +// is a timing race between the goroutine observing cancellation and its +// ticker firing. The test instead hammers many rapid connect/disconnect +// cycles with an aggressive ping interval to maximize the window in which a +// missing join would be exercised, and relies on this project's `-race` test +// gate to flag the resulting data race, or on the unrecovered panic crashing +// the test binary outright. +func TestWebSocketProxy_HandleConnection_PingGoroutineJoinsBeforeReturn(t *testing.T) { + upgrader := gorillaws.Upgrader{} + backendSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + defer backendSrv.Close() + + config := WebSocketConfig{ + Enabled: true, + PingInterval: 5 * time.Millisecond, + PongTimeout: 20 * time.Millisecond, + } + wsp := NewWebSocketProxy(backendSrv.URL, config, libpack_logger.New(), nil) + + proxyAddr := startTestWebSocketProxyApp(t, wsp) + + const iterations = 25 + for i := 0; i < iterations; i++ { + clientConn, _, err := gorillaws.DefaultDialer.Dial("ws://"+proxyAddr+"/ws", nil) + if err != nil { + t.Fatalf("iteration %d: dial proxy: %v", i, err) + } + if err := clientConn.WriteMessage(gorillaws.TextMessage, []byte(`{"type":"connection_init"}`)); err != nil { + t.Fatalf("iteration %d: write connection_init: %v", i, err) + } + // Disconnect immediately, no close handshake, while the ping + // goroutine's ticker is still armed - handleConnection's teardown + // and join then race a live ticker on every iteration. + _ = clientConn.Close() + + assert.Eventually(t, func() bool { + return wsp.activeConnections.Load() == 0 + }, 2*time.Second, time.Millisecond, "iteration %d: handleConnection did not tear down", i) + } +}