diff --git a/webterm/server.go b/webterm/server.go index 7eaed74..9396ca5 100644 --- a/webterm/server.go +++ b/webterm/server.go @@ -343,6 +343,9 @@ func (s *LocalServer) stopWSClient(routeKey string, expected *wsClient) { return } close(client.send) + if client.conn != nil { + _ = client.conn.Close() + } <-client.done } diff --git a/webterm/server_test.go b/webterm/server_test.go index fbe0d41..f2347eb 100644 --- a/webterm/server_test.go +++ b/webterm/server_test.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "io" + "net" "net/http" "net/http/httptest" "strings" @@ -259,6 +260,47 @@ func TestWebSocketOldConnectionCloseDoesNotDropNewClient(t *testing.T) { } } +func TestStopWSClientClosesWebSocketConnection(t *testing.T) { + server, httpServer, _ := newServerForTests(t, false) + wsURL := "ws" + strings.TrimPrefix(httpServer.URL, "http") + "/ws/shell" + + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("ws dial error = %v", err) + } + defer conn.Close() + if err := conn.WriteJSON([]any{"resize", map[string]any{"width": 80, "height": 24}}); err != nil { + t.Fatalf("resize write: %v", err) + } + + var client *wsClient + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + server.mu.RLock() + client = server.wsClients["shell"] + server.mu.RUnlock() + if client != nil { + break + } + time.Sleep(10 * time.Millisecond) + } + if client == nil { + t.Fatalf("expected websocket client to be registered") + } + + server.stopWSClient("shell", client) + + _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + _, _, err = conn.ReadMessage() + if err == nil { + t.Fatalf("expected websocket close after stopWSClient") + } + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + t.Fatalf("expected immediate disconnect, got timeout: %v", err) + } +} + func TestStaleSessionConnectorCloseDoesNotDropReassignedRouteClient(t *testing.T) { server, httpServer, _ := newServerForTests(t, false) wsURL := "ws" + strings.TrimPrefix(httpServer.URL, "http") + "/ws/shell"