diff --git a/websocket.go b/websocket.go index 34a6480..8d5495c 100644 --- a/websocket.go +++ b/websocket.go @@ -212,6 +212,51 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web // 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. + pingInterval := wsp.pingInterval + pongTimeout := wsp.pongTimeout + if pingInterval <= 0 { + pingInterval = 30 * time.Second + } + if pongTimeout <= 0 { + pongTimeout = 60 * time.Second + } + + _ = 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). + 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) + } + } + }() + // Set up bidirectional proxying var wg sync.WaitGroup wg.Add(2) @@ -292,6 +337,12 @@ func (wsp *WebSocketProxy) proxyClientToBackend(ctx context.Context, client *web 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 wsp.pongTimeout > 0 { + _ = client.SetReadDeadline(time.Now().Add(wsp.pongTimeout)) + } + // Forward message to backend if err := backend.WriteMessage(messageType, message); err != nil { wsp.errors.Add(1) @@ -356,6 +407,13 @@ func (wsp *WebSocketProxy) proxyBackendToClient(ctx context.Context, backend *go 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 wsp.pongTimeout > 0 { + _ = backend.SetReadDeadline(time.Now().Add(wsp.pongTimeout)) + } + // Forward message to client if err := client.WriteMessage(messageType, message); err != nil { wsp.errors.Add(1) diff --git a/websocket_test.go b/websocket_test.go index bd62735..55c7ab1 100644 --- a/websocket_test.go +++ b/websocket_test.go @@ -3,9 +3,12 @@ package main import ( "context" "net/http" + "net/http/httptest" + "strings" "testing" "time" + gorillaws "github.com/gorilla/websocket" libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging" libpack_monitoring "github.com/lukaszraczylo/graphql-monitoring-proxy/monitoring" "github.com/stretchr/testify/assert" @@ -307,6 +310,58 @@ func TestWebSocketProxy_ConfigValidation(t *testing.T) { }) } +func TestWebSocketProxy_ReapsSilentBackend_Keepalive(t *testing.T) { + config := WebSocketConfig{ + Enabled: true, + PingInterval: 30 * time.Second, + PongTimeout: 150 * time.Millisecond, + } + wsp := NewWebSocketProxy("ws://placeholder", config, libpack_logger.New(), nil) + + // Real backend that completes the WebSocket upgrade then stays silent (TCP + // open, no frames) - simulates a half-open peer whose connection the + // keepalive read deadline must reap instead of hanging forever. + serverRelease := make(chan struct{}) + upgrader := gorillaws.Upgrader{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + <-serverRelease + _ = conn.Close() + })) + defer srv.Close() + defer close(serverRelease) + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + backend, _, err := gorillaws.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("dial backend: %v", err) + } + defer backend.Close() + + // Mirror handleConnection's keepalive: refresh the read deadline on the + // configured pong timeout so an idle backend is reaped. + backend.SetReadDeadline(time.Now().Add(config.PongTimeout)) + + done := make(chan struct{}) + go func() { + // 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") + close(done) + }() + + select { + case <-done: + // Reaped via the keepalive read deadline rather than blocking forever. + case <-time.After(2 * time.Second): + t.Fatal("proxyBackendToClient hung; keepalive read deadline did not reap silent backend") + } +} + func TestWebSocketProxy_StatsStructure(t *testing.T) { config := WebSocketConfig{ Enabled: true,