From ccc25426dd8a17598216d7b3ea3bb8337d00b1c8 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Fri, 14 Aug 2026 00:01:03 +0100 Subject: [PATCH] fix(websocket): join ping goroutine and unify keepalive deadlines The handler returned while the ping goroutine could still write to a connection that gofiber had recycled into its pool. This caused a nil dereference panic or a ping on an unrelated socket. Join the ping goroutine through a done channel before the handler returns. Compute effective keepalive values once in the constructor, clamp pongTimeout to at least pingInterval with one warning, and pass the effective values to both proxy directions. A negative interval or timeout disables keepalive. New test drives the real HandleWebSocket wiring. --- websocket.go | 180 +++++++++++++++++++++++++++++++++----------- websocket_test.go | 185 +++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 321 insertions(+), 44 deletions(-) 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) + } +}