Fix transparent websocket recovery for terminals
Close retired websocket connections in stopWSClient so clients reconnect promptly instead of remaining in a stdin-only state with no returning output. Add regression coverage to verify stopWSClient actively disconnects the websocket. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
@@ -343,6 +343,9 @@ func (s *LocalServer) stopWSClient(routeKey string, expected *wsClient) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
close(client.send)
|
close(client.send)
|
||||||
|
if client.conn != nil {
|
||||||
|
_ = client.conn.Close()
|
||||||
|
}
|
||||||
<-client.done
|
<-client.done
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"io"
|
"io"
|
||||||
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"strings"
|
"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) {
|
func TestStaleSessionConnectorCloseDoesNotDropReassignedRouteClient(t *testing.T) {
|
||||||
server, httpServer, _ := newServerForTests(t, false)
|
server, httpServer, _ := newServerForTests(t, false)
|
||||||
wsURL := "ws" + strings.TrimPrefix(httpServer.URL, "http") + "/ws/shell"
|
wsURL := "ws" + strings.TrimPrefix(httpServer.URL, "http") + "/ws/shell"
|
||||||
|
|||||||
Reference in New Issue
Block a user