From 0e8a8fe33f2600d6d11ff6e45835a61fba5eb150 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 13 Aug 2026 03:10:04 +0100 Subject: [PATCH] fix(websocket): tear down peer connection and bound backend messages Two defects in the bidirectional proxy: 1. Goroutine/socket leak on client disconnect. proxyClientToBackend and proxyBackendToClient each block on their own Read; when the client drops, the client-direction goroutine returns but the backend-direction goroutine stays blocked on backend.ReadMessage, so wg.Wait() never returns and the goroutine plus backend socket leak until the backend independently closes or sends. Each direction now closes the opposite connection on exit, unblocking the peer. 2. Backend->client messages were unbounded. client.SetReadLimit only bound what the client could send; a large or misbehaving backend could push arbitrarily large frames through. backend.SetReadLimit now bounds it. Tests: TestWebSocketProxy_TeardownOnClientDisconnect (backend must be closed after client disconnect), TestWebSocketProxy_BackendReadLimit (connection torn down when backend exceeds the limit). --- concerns_test.go | 111 +++++++++++++++++++++++++++++++++++++++++++++++ websocket.go | 12 +++++ 2 files changed, 123 insertions(+) diff --git a/concerns_test.go b/concerns_test.go index e36bb0e..a84cf2e 100644 --- a/concerns_test.go +++ b/concerns_test.go @@ -16,6 +16,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "strings" "testing" "time" @@ -434,3 +435,113 @@ func TestStartCacheMemoryMonitoring_NoPanicWithMinimalSetup(t *testing.T) { t.Fatal("startCacheMemoryMonitoring did not exit within 1s") } } + +// TestWebSocketProxy_TeardownOnClientDisconnect verifies that when the client +// side of the WebSocket proxy closes (logs off / drops), the backend +// connection is torn down too. Previously the backend-direction goroutine +// stayed blocked on backend.ReadMessage, so the goroutine and the backend +// socket leaked until the backend independently closed or sent a message. +func TestWebSocketProxy_TeardownOnClientDisconnect(t *testing.T) { + backendClosed := make(chan struct{}) + + // Backend WS server that stays open until the proxy disconnects it. + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&gorillaws.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + defer func() { _ = conn.Close() }() + // Block reading; if the proxy closes the connection (the fix), this + // returns and we signal teardown. + for { + if _, _, err := conn.ReadMessage(); err != nil { + close(backendClosed) + return + } + } + })) + defer backend.Close() + + wsp := NewWebSocketProxy(backend.URL, WebSocketConfig{ + Enabled: true, MaxMessageSize: 64 * 1024, + }, libpack_logger.New(), nil) + + app := fiber.New() + app.Get("/ws", websocket.New(func(c *websocket.Conn) { + wsp.handleConnection(context.Background(), c, http.Header{}) + })) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + go func() { _ = app.Listener(ln) }() + t.Cleanup(func() { _ = app.Shutdown() }) + + client, _, err := gorillaws.DefaultDialer.Dial("ws://"+ln.Addr().String()+"/ws", nil) + require.NoError(t, err) + + // Trigger backend dial + first-message forward. + err = client.WriteMessage(gorillaws.TextMessage, []byte(`{"type":"connection_init","payload":{}}`)) + require.NoError(t, err) + // Give the proxy a moment to connect to the backend before we drop the client. + time.Sleep(200 * time.Millisecond) + + // Client logs off. + _ = client.Close() + + select { + case <-backendClosed: + // Backend was torn down — no leak. + case <-time.After(3 * time.Second): + t.Fatal("backend connection was not closed after client disconnect (goroutine / socket leak)") + } +} + +// TestWebSocketProxy_BackendReadLimit verifies that messages received from the +// backend are bounded by MaxMessageSize (previously only the client-side read +// was bounded), so an oversized backend frame closes the connection instead of +// being forwarded to the client unbounded. +func TestWebSocketProxy_BackendReadLimit(t *testing.T) { + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&gorillaws.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + defer func() { _ = conn.Close() }() + // Send a 16KB frame; the proxy's read limit is set far below this. + payload := strings.Repeat("x", 16*1024) + for { + if err := conn.WriteMessage(gorillaws.TextMessage, []byte(payload)); err != nil { + return + } + time.Sleep(50 * time.Millisecond) + } + })) + defer backend.Close() + + wsp := NewWebSocketProxy(backend.URL, WebSocketConfig{ + Enabled: true, MaxMessageSize: 1024, // tiny limit + }, libpack_logger.New(), nil) + + app := fiber.New() + app.Get("/ws", websocket.New(func(c *websocket.Conn) { + wsp.handleConnection(context.Background(), c, http.Header{}) + })) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + go func() { _ = app.Listener(ln) }() + t.Cleanup(func() { _ = app.Shutdown() }) + + client, _, err := gorillaws.DefaultDialer.Dial("ws://"+ln.Addr().String()+"/ws", nil) + require.NoError(t, err) + defer func() { _ = client.Close() }() + + err = client.WriteMessage(gorillaws.TextMessage, []byte(`{"type":"connection_init","payload":{}}`)) + require.NoError(t, err) + + // Backend pushes frames that exceed the read limit; the proxy must tear + // the connection down rather than forwarding unbounded frames. + client.SetReadDeadline(time.Now().Add(3 * time.Second)) //nolint:errcheck + _, _, readErr := client.ReadMessage() + require.Error(t, readErr, "expected connection to be closed after backend exceeded read limit") +} diff --git a/websocket.go b/websocket.go index 2134ec0..34a6480 100644 --- a/websocket.go +++ b/websocket.go @@ -207,6 +207,11 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web }) } + // 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) + // Set up bidirectional proxying var wg sync.WaitGroup wg.Add(2) @@ -214,12 +219,19 @@ func (wsp *WebSocketProxy) handleConnection(ctx context.Context, clientConn *web // 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) }() // 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) }()