From 590ffb31f09ae745269698dded561f5a061c4154 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Salavisa?= Date: Wed, 12 Aug 2026 22:36:31 +0100 Subject: [PATCH] test(gateway): cover websocket terminal pumps, proxies and tunnels (#2587) Co-authored-by: Kandev Agent Co-authored-by: Claude Opus 5 (1M context) --- .../websocket/client_pump_integration_test.go | 254 +++++++++ .../websocket/client_unsubscribe_test.go | 260 +++++++++ .../gateway/websocket/gateway_wiring_test.go | 218 ++++++++ .../websocket/lsp_handler_proxy_test.go | 253 +++++++++ .../websocket/port_proxy_handler_test.go | 460 ++++++++++++++++ .../websocket/port_tunnel_serve_test.go | 378 +++++++++++++ .../websocket/terminal_handler_remote_test.go | 316 +++++++++++ .../websocket/terminal_handler_routes_test.go | 316 +++++++++++ .../websocket/terminal_pumps_bridge_test.go | 293 ++++++++++ .../websocket/terminal_pumps_pty_test.go | 514 ++++++++++++++++++ .../gateway/websocket/terminal_wsutil_test.go | 149 +++++ .../websocket/vscode_proxy_handler_test.go | 303 +++++++++++ 12 files changed, 3714 insertions(+) create mode 100644 apps/backend/internal/gateway/websocket/client_pump_integration_test.go create mode 100644 apps/backend/internal/gateway/websocket/client_unsubscribe_test.go create mode 100644 apps/backend/internal/gateway/websocket/gateway_wiring_test.go create mode 100644 apps/backend/internal/gateway/websocket/lsp_handler_proxy_test.go create mode 100644 apps/backend/internal/gateway/websocket/port_proxy_handler_test.go create mode 100644 apps/backend/internal/gateway/websocket/port_tunnel_serve_test.go create mode 100644 apps/backend/internal/gateway/websocket/terminal_handler_remote_test.go create mode 100644 apps/backend/internal/gateway/websocket/terminal_handler_routes_test.go create mode 100644 apps/backend/internal/gateway/websocket/terminal_pumps_bridge_test.go create mode 100644 apps/backend/internal/gateway/websocket/terminal_pumps_pty_test.go create mode 100644 apps/backend/internal/gateway/websocket/terminal_wsutil_test.go create mode 100644 apps/backend/internal/gateway/websocket/vscode_proxy_handler_test.go diff --git a/apps/backend/internal/gateway/websocket/client_pump_integration_test.go b/apps/backend/internal/gateway/websocket/client_pump_integration_test.go new file mode 100644 index 0000000000..6a4a950938 --- /dev/null +++ b/apps/backend/internal/gateway/websocket/client_pump_integration_test.go @@ -0,0 +1,254 @@ +package websocket + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + gorillaws "github.com/gorilla/websocket" + + "github.com/kandev/kandev/internal/auth/authn" + "github.com/kandev/kandev/internal/common/logger" + ws "github.com/kandev/kandev/pkg/websocket" +) + +// connectedClient is a live browser connection driving a real Client with both +// pumps running, plus the handles needed to join on their shutdown. +type connectedClient struct { + browser *gorillaws.Conn + client *Client + hub *Hub + stopped <-chan struct{} +} + +// newConnectedClient wires a hub, a real WebSocket connection, and both client +// pumps together so tests can exercise the full receive/dispatch/send path. +func newConnectedClient(t *testing.T) *connectedClient { + t.Helper() + log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "json"}) + if err != nil { + t.Fatalf("logger: %v", err) + } + hub := NewHub(nil, log) + + hubCtx, cancelHub := context.WithCancel(context.Background()) + hubDone := make(chan struct{}) + go func() { + defer close(hubDone) + hub.Run(hubCtx) + }() + t.Cleanup(func() { + cancelHub() + <-hubDone + }) + + ready := make(chan *Client, 1) + pumpsDone := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, upgradeErr := testWSUpgrader.Upgrade(w, r, nil) + if upgradeErr != nil { + return + } + client := NewClient("client-1", authn.Identity{}, conn, hub, log) + hub.Register(client) + ready <- client + + writeDone := make(chan struct{}) + go func() { + defer close(writeDone) + client.WritePump() + }() + client.ReadPump(r.Context()) + <-writeDone + close(pumpsDone) + })) + t.Cleanup(srv.Close) + + browser, resp, err := gorillaws.DefaultDialer.Dial("ws"+strings.TrimPrefix(srv.URL, "http"), nil) + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + if err != nil { + t.Fatalf("dial gateway: %v", err) + } + t.Cleanup(func() { _ = browser.Close() }) + + // Join the pumps from a cleanup registered before the assertions below, so + // an early t.Fatal cannot leave them running into a later test. + t.Cleanup(func() { + _ = browser.Close() + pumpTimer := time.NewTimer(wsTestTimeout) + defer pumpTimer.Stop() + select { + case <-pumpsDone: + case <-pumpTimer.C: + t.Errorf("client pumps did not return within %s", wsTestTimeout) + } + }) + + readyTimer := time.NewTimer(wsTestTimeout) + defer readyTimer.Stop() + var client *Client + select { + case client = <-ready: + case <-readyTimer.C: + t.Fatal("gateway never registered the client") + } + + return &connectedClient{browser: browser, client: client, hub: hub, stopped: pumpsDone} +} + +// roundTrip sends one frame and reads the single reply it produces. +func (cc *connectedClient) roundTrip(t *testing.T, raw []byte) ws.Message { + t.Helper() + if err := cc.browser.WriteMessage(gorillaws.TextMessage, raw); err != nil { + t.Fatalf("browser write: %v", err) + } + if err := cc.browser.SetReadDeadline(time.Now().Add(wsTestTimeout)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + _, data, err := cc.browser.ReadMessage() + if err != nil { + t.Fatalf("browser read: %v", err) + } + var msg ws.Message + if err := json.Unmarshal(data, &msg); err != nil { + t.Fatalf("decode gateway frame %s: %v", data, err) + } + return msg +} + +func (cc *connectedClient) request(t *testing.T, id, action string, payload any) ws.Message { + t.Helper() + raw, err := json.Marshal(ws.Message{ + ID: id, + Type: ws.MessageTypeRequest, + Action: action, + Payload: mustJSON(t, payload), + }) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + return cc.roundTrip(t, raw) +} + +func mustJSON(t *testing.T, v any) json.RawMessage { + t.Helper() + raw, err := json.Marshal(v) + if err != nil { + t.Fatalf("marshal payload: %v", err) + } + return raw +} + +func TestReadPumpRejectsMalformedFramesWithoutDroppingTheConnection(t *testing.T) { + cc := newConnectedClient(t) + + bad := cc.roundTrip(t, []byte("{not json")) + if bad.Type != ws.MessageTypeError { + t.Fatalf("response = %+v, want an error frame", bad) + } + if code := errorCode(t, bad); code != ws.ErrorCodeBadRequest { + t.Fatalf("error code = %q, want %q", code, ws.ErrorCodeBadRequest) + } + + // The connection must survive a bad frame and keep serving requests. + ack := cc.request(t, "req-after-bad", ws.ActionTaskSubscribe, SubscribeRequest{TaskID: "task-1"}) + if ack.Type != ws.MessageTypeResponse { + t.Fatalf("follow-up response = %+v, want a success response", ack) + } +} + +// MCP actions are PAT-authenticated over /mcp; the raw socket carries no +// credential of its own and must refuse them before dispatch. +func TestReadPumpRefusesMCPActionsOverTheRawSocket(t *testing.T) { + cc := newConnectedClient(t) + + resp := cc.request(t, "req-mcp", "mcp.tools.list", map[string]string{}) + if resp.Type != ws.MessageTypeError { + t.Fatalf("response = %+v, want an error frame", resp) + } + if code := errorCode(t, resp); code != ws.ErrorCodeForbidden { + t.Fatalf("error code = %q, want %q", code, ws.ErrorCodeForbidden) + } +} + +func TestConnectedClientSubscriptionRoundTrips(t *testing.T) { + cc := newConnectedClient(t) + + if resp := cc.request(t, "1", ws.ActionTaskSubscribe, SubscribeRequest{TaskID: "task-1"}); resp.Type != ws.MessageTypeResponse { + t.Fatalf("task.subscribe response = %+v", resp) + } + cc.assertTaskSubscribed(t, "task-1", true) + + if resp := cc.request(t, "2", ws.ActionTaskUnsubscribe, SubscribeRequest{TaskID: "task-1"}); resp.Type != ws.MessageTypeResponse { + t.Fatalf("task.unsubscribe response = %+v", resp) + } + cc.assertTaskSubscribed(t, "task-1", false) + + if resp := cc.request(t, "3", ws.ActionSystemMetricsSubscribe, map[string]string{}); resp.Type != ws.MessageTypeResponse { + t.Fatalf("system.metrics.subscribe response = %+v", resp) + } + cc.assertMetricsSubscribed(t, true) + + if resp := cc.request(t, "4", ws.ActionSystemMetricsUnsubscribe, map[string]string{}); resp.Type != ws.MessageTypeResponse { + t.Fatalf("system.metrics.unsubscribe response = %+v", resp) + } + cc.assertMetricsSubscribed(t, false) + + if resp := cc.request(t, "5", ws.ActionRunSubscribe, RunSubscribeRequest{RunID: "run-1"}); resp.Type != ws.MessageTypeResponse { + t.Fatalf("run.subscribe response = %+v", resp) + } + if resp := cc.request(t, "6", ws.ActionRunUnsubscribe, RunSubscribeRequest{RunID: "run-1"}); resp.Type != ws.MessageTypeResponse { + t.Fatalf("run.unsubscribe response = %+v", resp) + } +} + +// A browser disconnect must unwind both pumps and drop the client from the hub, +// otherwise every closed tab leaks a goroutine pair and a subscriber entry. +func TestClientPumpsUnwindOnBrowserDisconnect(t *testing.T) { + cc := newConnectedClient(t) + + if resp := cc.request(t, "1", ws.ActionTaskSubscribe, SubscribeRequest{TaskID: "task-1"}); resp.Type != ws.MessageTypeResponse { + t.Fatalf("task.subscribe response = %+v", resp) + } + if got := cc.hub.GetClientCount(); got != 1 { + t.Fatalf("hub client count = %d, want 1", got) + } + + if err := cc.browser.Close(); err != nil { + t.Fatalf("close browser: %v", err) + } + joinWithin(t, cc.stopped, "client pumps") + + if got := cc.hub.GetClientCount(); got != 0 { + t.Fatalf("hub client count = %d after disconnect, want 0", got) + } + cc.assertTaskSubscribed(t, "task-1", false) +} + +func (cc *connectedClient) assertTaskSubscribed(t *testing.T, taskID string, want bool) { + t.Helper() + cc.hub.mu.RLock() + defer cc.hub.mu.RUnlock() + _, got := cc.hub.taskSubscribers[taskID][cc.client] + if got != want { + t.Fatalf("hub task subscription for %q = %v, want %v", taskID, got, want) + } +} + +func (cc *connectedClient) assertMetricsSubscribed(t *testing.T, want bool) { + t.Helper() + cc.hub.mu.RLock() + defer cc.hub.mu.RUnlock() + if got := cc.hub.systemMetricsSubscribers[cc.client]; got != want { + t.Fatalf("hub metrics subscription = %v, want %v", got, want) + } + if got := cc.client.systemMetricsSubscribed; got != want { + t.Fatalf("client metrics flag = %v, want %v", got, want) + } +} diff --git a/apps/backend/internal/gateway/websocket/client_unsubscribe_test.go b/apps/backend/internal/gateway/websocket/client_unsubscribe_test.go new file mode 100644 index 0000000000..eaa0812092 --- /dev/null +++ b/apps/backend/internal/gateway/websocket/client_unsubscribe_test.go @@ -0,0 +1,260 @@ +package websocket + +import ( + "encoding/json" + "testing" + + ws "github.com/kandev/kandev/pkg/websocket" +) + +// nextClientFrame pops one queued frame from the client's send channel. +func nextClientFrame(t *testing.T, c *Client) ws.Message { + t.Helper() + select { + case data := <-c.send: + var msg ws.Message + if err := json.Unmarshal(data, &msg); err != nil { + t.Fatalf("decode client frame: %v", err) + } + return msg + default: + t.Fatal("expected a frame in the client send queue") + return ws.Message{} + } +} + +func assertNoMoreFrames(t *testing.T, c *Client) { + t.Helper() + select { + case data := <-c.send: + t.Fatalf("unexpected extra frame: %s", data) + default: + } +} + +// errorCode extracts the error code from an error-typed frame. +func errorCode(t *testing.T, msg ws.Message) string { + t.Helper() + var payload ws.ErrorPayload + if err := json.Unmarshal(msg.Payload, &payload); err != nil { + t.Fatalf("decode error payload %s: %v", msg.Payload, err) + } + return payload.Code +} + +func requestMessage(t *testing.T, action string, payload any) *ws.Message { + t.Helper() + raw, err := json.Marshal(payload) + if err != nil { + t.Fatalf("marshal payload: %v", err) + } + return &ws.Message{ID: "req-1", Type: ws.MessageTypeRequest, Action: action, Payload: raw} +} + +func TestHandleUnsubscribeRemovesTaskSubscription(t *testing.T) { + h := newTestHub(t) + c := newTestClient("c-unsub") + c.hub = h + + h.SubscribeToTask(c, "task-1") + if !c.subscriptions["task-1"] { + t.Fatal("precondition failed: client is not subscribed to task-1") + } + + c.handleUnsubscribe(requestMessage(t, ws.ActionTaskUnsubscribe, SubscribeRequest{TaskID: "task-1"})) + + if c.subscriptions["task-1"] { + t.Fatal("client is still subscribed to task-1") + } + h.mu.RLock() + _, stillTracked := h.taskSubscribers["task-1"] + h.mu.RUnlock() + if stillTracked { + t.Fatal("hub still tracks a subscriber set for task-1") + } + + resp := nextClientFrame(t, c) + if resp.Type != ws.MessageTypeResponse || resp.Action != ws.ActionTaskUnsubscribe { + t.Fatalf("response = %+v, want a task.unsubscribe response", resp) + } + assertNoMoreFrames(t, c) +} + +func TestHandleSessionUnsubscribeRemovesSessionSubscription(t *testing.T) { + h := newTestHub(t) + c := newTestClient("c-sess-unsub") + c.hub = h + + h.SubscribeToSession(c, "sess-1") + if !c.sessionSubscriptions["sess-1"] { + t.Fatal("precondition failed: client is not subscribed to sess-1") + } + + c.handleSessionUnsubscribe( + requestMessage(t, ws.ActionSessionUnsubscribe, SessionSubscribeRequest{SessionID: "sess-1"}), + ) + + if c.sessionSubscriptions["sess-1"] { + t.Fatal("client is still subscribed to sess-1") + } + if resp := nextClientFrame(t, c); resp.Type != ws.MessageTypeResponse { + t.Fatalf("response = %+v, want a success response", resp) + } +} + +func TestHandleSessionUnfocusReleasesFocusButKeepsSubscription(t *testing.T) { + h := newTestHub(t) + c := newTestClient("c-unfocus") + c.hub = h + + h.SubscribeToSession(c, "sess-1") + h.FocusSession(c, "sess-1") + if !c.sessionFocus["sess-1"] { + t.Fatal("precondition failed: session is not focused") + } + + c.handleSessionUnfocus( + requestMessage(t, ws.ActionSessionUnfocus, SessionSubscribeRequest{SessionID: "sess-1"}), + ) + + if c.sessionFocus["sess-1"] { + t.Fatal("focus was not released") + } + if !c.sessionSubscriptions["sess-1"] { + t.Fatal("unfocus must not drop the subscription") + } + if resp := nextClientFrame(t, c); resp.Action != ws.ActionSessionUnfocus { + t.Fatalf("response action = %q, want %q", resp.Action, ws.ActionSessionUnfocus) + } +} + +func TestHandleUserUnsubscribeIsScopedToOwnTopic(t *testing.T) { + h := newTestHub(t) + c := newTestClient("c-user-unsub") + c.hub = h + + own := c.ownUserTopic() + h.SubscribeToUser(c, own) + if !c.userSubscriptions[own] { + t.Fatal("precondition failed: client is not subscribed to its own user topic") + } + + c.handleUserUnsubscribe( + requestMessage(t, ws.ActionUserUnsubscribe, UserSubscribeRequest{UserID: "someone-else"}), + ) + forbidden := nextClientFrame(t, c) + if forbidden.Type != ws.MessageTypeError { + t.Fatalf("response = %+v, want an error for a foreign user topic", forbidden) + } + if !c.userSubscriptions[own] { + t.Fatal("a rejected unsubscribe must not drop the client's own subscription") + } + + c.handleUserUnsubscribe(requestMessage(t, ws.ActionUserUnsubscribe, UserSubscribeRequest{})) + if c.userSubscriptions[own] { + t.Fatal("an empty user_id must unsubscribe the client's own topic") + } + if resp := nextClientFrame(t, c); resp.Type != ws.MessageTypeResponse { + t.Fatalf("response = %+v, want a success response", resp) + } +} + +func TestHandleRunUnsubscribeRemovesRunSubscription(t *testing.T) { + h := newTestHub(t) + c := newTestClient("c-run-unsub") + c.hub = h + + h.SubscribeToRun(c, "run-1") + if !c.runSubscriptions["run-1"] { + t.Fatal("precondition failed: client is not subscribed to run-1") + } + + c.handleRunUnsubscribe(requestMessage(t, ws.ActionRunUnsubscribe, RunSubscribeRequest{RunID: "run-1"})) + + if c.runSubscriptions["run-1"] { + t.Fatal("client is still subscribed to run-1") + } + h.mu.RLock() + _, stillTracked := h.runSubscribers["run-1"] + h.mu.RUnlock() + if stillTracked { + t.Fatal("hub still tracks a subscriber set for run-1") + } + if resp := nextClientFrame(t, c); resp.Type != ws.MessageTypeResponse { + t.Fatalf("response = %+v, want a success response", resp) + } +} + +// Every subscription handler must reject a malformed envelope and a missing +// identifier before touching hub state. +func TestSubscriptionHandlersRejectBadPayloads(t *testing.T) { + handlers := map[string]struct { + action string + invoke func(*Client, *ws.Message) + empty any + }{ + "task.subscribe": { + action: ws.ActionTaskSubscribe, + invoke: (*Client).handleSubscribe, + empty: SubscribeRequest{}, + }, + "task.unsubscribe": { + action: ws.ActionTaskUnsubscribe, + invoke: (*Client).handleUnsubscribe, + empty: SubscribeRequest{}, + }, + "session.subscribe": { + action: ws.ActionSessionSubscribe, + invoke: (*Client).handleSessionSubscribe, + empty: SessionSubscribeRequest{}, + }, + "session.unsubscribe": { + action: ws.ActionSessionUnsubscribe, + invoke: (*Client).handleSessionUnsubscribe, + empty: SessionSubscribeRequest{}, + }, + "session.unfocus": { + action: ws.ActionSessionUnfocus, + invoke: (*Client).handleSessionUnfocus, + empty: SessionSubscribeRequest{}, + }, + "run.subscribe": { + action: ws.ActionRunSubscribe, + invoke: (*Client).handleRunSubscribe, + empty: RunSubscribeRequest{}, + }, + "run.unsubscribe": { + action: ws.ActionRunUnsubscribe, + invoke: (*Client).handleRunUnsubscribe, + empty: RunSubscribeRequest{}, + }, + } + + for name, tc := range handlers { + t.Run(name, func(t *testing.T) { + h := newTestHub(t) + c := newTestClient("c-" + name) + c.hub = h + + malformed := &ws.Message{ + ID: "req-1", + Type: ws.MessageTypeRequest, + Action: tc.action, + Payload: json.RawMessage(`"not-an-object"`), + } + tc.invoke(c, malformed) + if resp := nextClientFrame(t, c); resp.Type != ws.MessageTypeError { + t.Fatalf("malformed payload response = %+v, want an error", resp) + } + + tc.invoke(c, requestMessage(t, tc.action, tc.empty)) + resp := nextClientFrame(t, c) + if resp.Type != ws.MessageTypeError { + t.Fatalf("missing-identifier response = %+v, want an error", resp) + } + if code := errorCode(t, resp); code != ws.ErrorCodeValidation { + t.Fatalf("missing-identifier error code = %q, want %q", code, ws.ErrorCodeValidation) + } + }) + } +} diff --git a/apps/backend/internal/gateway/websocket/gateway_wiring_test.go b/apps/backend/internal/gateway/websocket/gateway_wiring_test.go new file mode 100644 index 0000000000..2e8e56557c --- /dev/null +++ b/apps/backend/internal/gateway/websocket/gateway_wiring_test.go @@ -0,0 +1,218 @@ +package websocket + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sort" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/kandev/kandev/internal/common/logger" + ws "github.com/kandev/kandev/pkg/websocket" +) + +func newGatewayForTest(t *testing.T) *Gateway { + t.Helper() + log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "json"}) + if err != nil { + t.Fatalf("logger: %v", err) + } + gateway, err := Provide(log) + if err != nil { + t.Fatalf("Provide: %v", err) + } + return gateway +} + +// registeredRoutes returns "METHOD PATH" strings for everything on the engine. +func registeredRoutes(engine *gin.Engine) []string { + routes := make([]string, 0, len(engine.Routes())) + for _, route := range engine.Routes() { + routes = append(routes, route.Method+" "+route.Path) + } + sort.Strings(routes) + return routes +} + +func hasRoute(routes []string, want string) bool { + for _, route := range routes { + if route == want { + return true + } + } + return false +} + +// A gateway with no optional handlers wired must expose only /ws — the terminal, +// LSP and proxy surfaces are opt-in and must not appear by accident. +func TestSetupRoutesRegistersOnlyWiredSurfaces(t *testing.T) { + gateway := newGatewayForTest(t) + gin.SetMode(gin.TestMode) + + bare := gin.New() + gateway.SetupRoutes(bare) + bareRoutes := registeredRoutes(bare) + if !hasRoute(bareRoutes, "GET /ws") { + t.Fatalf("routes = %v, want GET /ws", bareRoutes) + } + for _, unwanted := range []string{ + "GET /terminal/*target", + "GET /lsp/:sessionId", + "GET /vscode/:sessionId/*path", + "GET /port-proxy/:sessionId/:port/*path", + } { + if hasRoute(bareRoutes, unwanted) { + t.Fatalf("routes = %v, want no %q before the handler is wired", bareRoutes, unwanted) + } + } + + gateway.SetLifecycleManager(nil, nil, nil) + gateway.SetLSPHandler(nil, nil) + gateway.SetVscodeProxy(nil) + gateway.SetPortProxy(nil) + gateway.SetPortTunnel(nil) + + if gateway.TerminalHandler == nil || gateway.LSPHandler == nil || + gateway.VscodeProxyHandler == nil || gateway.PortProxyHandler == nil || + gateway.TunnelManager == nil { + t.Fatalf("gateway handlers were not all constructed: %+v", gateway) + } + + wired := gin.New() + gateway.SetupRoutes(wired) + wiredRoutes := registeredRoutes(wired) + for _, want := range []string{ + "GET /ws", + "GET /terminal/*target", + "GET /lsp/:sessionId", + "GET /vscode/:sessionId/*path", + "GET /port-proxy/:sessionId/:port", + "GET /port-proxy/:sessionId/:port/*path", + } { + if !hasRoute(wiredRoutes, want) { + t.Fatalf("routes = %v, want %q", wiredRoutes, want) + } + } +} + +func TestGatewayHealthActionIsDispatchable(t *testing.T) { + gateway := newGatewayForTest(t) + + dispatcher := gateway.Hub.GetDispatcher() + if dispatcher != gateway.Dispatcher { + t.Fatal("Hub.GetDispatcher() did not return the gateway's dispatcher") + } + + resp, err := dispatcher.Dispatch( + context.Background(), + &ws.Message{ID: "health-1", Type: ws.MessageTypeRequest, Action: ws.ActionHealthCheck}, + ) + if err != nil { + t.Fatalf("Dispatch(health) error = %v", err) + } + var payload map[string]string + if err := json.Unmarshal(resp.Payload, &payload); err != nil { + t.Fatalf("decode health payload: %v", err) + } + if payload["status"] != "ok" || payload["service"] != "kandev" { + t.Fatalf("health payload = %v, want status=ok service=kandev", payload) + } +} + +func TestBroadcastToTaskReachesOnlySubscribers(t *testing.T) { + h := newTestHub(t) + subscribed := newTestClient("c-subscribed") + subscribed.hub = h + other := newTestClient("c-other") + other.hub = h + + h.SubscribeToTask(subscribed, "task-1") + h.SubscribeToTask(other, "task-2") + + msg, err := ws.NewNotification("task.updated", map[string]string{"task_id": "task-1"}) + if err != nil { + t.Fatalf("build notification: %v", err) + } + h.BroadcastToTask("task-1", msg) + + got := nextClientFrame(t, subscribed) + if got.Action != "task.updated" { + t.Fatalf("subscriber received action %q, want %q", got.Action, "task.updated") + } + assertNoMoreFrames(t, other) +} + +func TestGetSessionModeReflectsSubscriptionAndFocus(t *testing.T) { + h := newTestHub(t) + c := newTestClient("c-mode") + c.hub = h + + if got := h.GetSessionMode("sess-1"); got != SessionModePaused { + t.Fatalf("mode with no interest = %q, want %q", got, SessionModePaused) + } + + h.SubscribeToSession(c, "sess-1") + if got := h.GetSessionMode("sess-1"); got != SessionModeSlow { + t.Fatalf("mode when subscribed = %q, want %q", got, SessionModeSlow) + } + + h.FocusSession(c, "sess-1") + if got := h.GetSessionMode("sess-1"); got != SessionModeFast { + t.Fatalf("mode when focused = %q, want %q", got, SessionModeFast) + } +} + +func TestRemoveTunnelDropsASingleCacheKey(t *testing.T) { + log, _ := observedTerminalLogger(t) + manager := newTunnelManagerForTest(t, log) + + manager.mu.Lock() + manager.tunnels["sess-a:3000"] = &tunnelEntry{port: 3000, tunnelPort: 40001} + manager.tunnels["sess-a:4000"] = &tunnelEntry{port: 4000, tunnelPort: 40002} + manager.mu.Unlock() + t.Cleanup(func() { + manager.mu.Lock() + manager.tunnels = make(map[string]*tunnelEntry) + manager.mu.Unlock() + }) + + manager.removeTunnel("sess-a:3000") + + manager.mu.Lock() + defer manager.mu.Unlock() + if _, ok := manager.tunnels["sess-a:3000"]; ok { + t.Fatal("removeTunnel did not drop its cache key") + } + if _, ok := manager.tunnels["sess-a:4000"]; !ok { + t.Fatal("removeTunnel dropped an unrelated cache key") + } +} + +func TestShouldUseWorkspaceShellIsFalseWithoutAContainerExecution(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + + if handler.shouldUseWorkspaceShell(context.Background(), "session-unknown") { + t.Fatal("shouldUseWorkspaceShell() = true for a session with no execution") + } +} + +// The /ws route is guarded by requireConnectionAuth, which must not be skipped +// just because the request is a WebSocket upgrade. +func TestSetupRoutesGuardsWebSocketRouteWithConnectionAuth(t *testing.T) { + gateway := newGatewayForTest(t) + gateway.SetAuthPolicy(AuthPolicy{Enforced: func() bool { return true }}) + gin.SetMode(gin.TestMode) + engine := gin.New() + gateway.SetupRoutes(engine) + + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/ws", nil)) + + if recorder.Code != http.StatusUnauthorized { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusUnauthorized) + } +} diff --git a/apps/backend/internal/gateway/websocket/lsp_handler_proxy_test.go b/apps/backend/internal/gateway/websocket/lsp_handler_proxy_test.go new file mode 100644 index 0000000000..23d572257a --- /dev/null +++ b/apps/backend/internal/gateway/websocket/lsp_handler_proxy_test.go @@ -0,0 +1,253 @@ +package websocket + +import ( + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + gorillaws "github.com/gorilla/websocket" + + "github.com/kandev/kandev/internal/agent/runtime/lifecycle" + "github.com/kandev/kandev/internal/agentruntime" +) + +// dialLSPHandler serves handler.HandleLSPConnection over a real HTTP server and +// dials it. The returned channel closes once the handler has returned, which is +// what lets a test assert on post-handler state (capacity release) without +// polling. +func dialLSPHandler(t *testing.T, handler *LSPHandler, target string) (*gorillaws.Conn, <-chan struct{}) { + t.Helper() + gin.SetMode(gin.TestMode) + + served := make(chan struct{}) + router := gin.New() + router.GET("/lsp/:sessionId", func(c *gin.Context) { + defer close(served) + handler.HandleLSPConnection(c) + }) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + conn, resp, err := gorillaws.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http")+target, nil) + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + if err != nil { + t.Fatalf("dial %s: %v", target, err) + } + t.Cleanup(func() { _ = conn.Close() }) + return conn, served +} + +// expectLSPCloseCode reads from conn and asserts the peer closed with wantCode. +func expectLSPCloseCode(t *testing.T, conn *gorillaws.Conn, wantCode int, wantText string) { + t.Helper() + if err := conn.SetReadDeadline(time.Now().Add(wsTestTimeout)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + _, _, err := conn.ReadMessage() + var closeErr *gorillaws.CloseError + if !errors.As(err, &closeErr) { + t.Fatalf("ReadMessage() error = %v, want a websocket close error", err) + } + if closeErr.Code != wantCode { + t.Fatalf("close code = %d, want %d", closeErr.Code, wantCode) + } + if wantText != "" && closeErr.Text != wantText { + t.Fatalf("close text = %q, want %q", closeErr.Text, wantText) + } +} + +func TestNewLSPHandlerInitializesCapacityAndLogger(t *testing.T) { + handler := NewLSPHandler(nil, nil, testLogger()) + if handler.capacity == nil { + t.Fatal("capacity limiter was not initialized") + } + if handler.logger == nil { + t.Fatal("logger was not initialized") + } + if !handler.capacity.TryAcquire() { + t.Fatal("a freshly built handler must have capacity available") + } + handler.capacity.Release() +} + +func TestHandleLSPConnectionRejectsInvalidRequests(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + sessionID string + query string + wantError string + }{ + {name: "missing session", sessionID: "", query: "?language=go", wantError: "sessionId is required"}, + {name: "missing language", sessionID: "s-1", query: "", wantError: "language query parameter is required"}, + { + name: "unsupported language", + sessionID: "s-1", + query: "?language=cobol", + wantError: "unsupported language: cobol", + }, + } + handler := &LSPHandler{logger: testLogger(), capacity: newLSPCapacityLimiter(1)} + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/lsp/"+tt.sessionID+tt.query, nil) + c.Params = gin.Params{{Key: "sessionId", Value: tt.sessionID}} + handler.HandleLSPConnection(c) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != tt.wantError { + t.Fatalf("error = %q, want %q", got, tt.wantError) + } + if !handler.capacity.TryAcquire() { + t.Fatal("a rejected request must not consume a capacity slot") + } + handler.capacity.Release() + }) + } +} + +func TestHandleLSPConnectionClosesWithSessionNotFoundCode(t *testing.T) { + manager := &recordingLSPLifecycleManager{resolveErr: errors.New("no such session")} + handler := &LSPHandler{ + lifecycleMgr: manager, + capacity: newLSPCapacityLimiter(1), + logger: testLogger(), + } + + conn, served := dialLSPHandler(t, handler, "/lsp/ghost?language=go") + expectLSPCloseCode(t, conn, lspCloseSessionNotFound, "session not found") + joinWithin(t, served, "HandleLSPConnection") + + if !handler.capacity.TryAcquire() { + t.Fatal("capacity was not released after a failed resolve") + } + handler.capacity.Release() +} + +func TestHandleLSPConnectionClosesWithUnsupportedExecutorCode(t *testing.T) { + manager := &recordingLSPLifecycleManager{runtimeName: agentruntime.RuntimeSSH} + handler := &LSPHandler{ + lifecycleMgr: manager, + capacity: newLSPCapacityLimiter(1), + logger: testLogger(), + } + + conn, served := dialLSPHandler(t, handler, "/lsp/remote?language=go") + expectLSPCloseCode(t, conn, lspCloseUnsupportedExecutor, lspCloseUnsupportedCloseText) + joinWithin(t, served, "HandleLSPConnection") + + if manager.ensureCalls != 0 { + t.Fatalf("GetOrEnsureExecution calls = %d, want 0 for an unsupported runtime", manager.ensureCalls) + } +} + +// The capacity slot is acquired inside resolveLSPExecution; a later failure +// (here: an execution with no agentctl client) must still give it back. +func TestHandleLSPConnectionReleasesCapacityWhenAgentctlIsMissing(t *testing.T) { + manager := &recordingLSPLifecycleManager{ + runtimeName: agentruntime.RuntimeStandalone, + execution: &lifecycle.AgentExecution{RuntimeName: agentruntime.RuntimeStandalone}, + } + handler := &LSPHandler{ + lifecycleMgr: manager, + capacity: newLSPCapacityLimiter(1), + logger: testLogger(), + } + + conn, served := dialLSPHandler(t, handler, "/lsp/no-agentctl?language=go") + expectLSPCloseCode(t, conn, lspCloseSessionNotFound, "agentctl unavailable") + joinWithin(t, served, "HandleLSPConnection") + + if !handler.capacity.TryAcquire() { + t.Fatal("capacity was not released after the agentctl client was found missing") + } + handler.capacity.Release() +} + +func TestProxyLSPConnectionsForwardsMessagesInBothDirections(t *testing.T) { + handler := &LSPHandler{logger: testLogger(), capacity: newLSPCapacityLimiter(1)} + + browser, handlerBrowserSide := newTerminalWSPair(t) + handlerUpstreamSide, upstream := newTerminalWSPair(t) + + done := spawnJoinable(t, "proxyLSPConnections", func() { _ = browser.Close() }, func() { + handler.proxyLSPConnections(handlerBrowserSide, handlerUpstreamSide, "sess-lsp", "go") + }) + + request := []byte(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`) + if err := browser.WriteMessage(gorillaws.TextMessage, request); err != nil { + t.Fatalf("browser write: %v", err) + } + if err := upstream.SetReadDeadline(time.Now().Add(wsTestTimeout)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + msgType, got, err := upstream.ReadMessage() + if err != nil { + t.Fatalf("upstream read: %v", err) + } + if msgType != gorillaws.TextMessage || string(got) != string(request) { + t.Fatalf("upstream received (%d, %q), want (%d, %q)", + msgType, got, gorillaws.TextMessage, request) + } + + response := []byte(`{"jsonrpc":"2.0","id":1,"result":{}}`) + if err := upstream.WriteMessage(gorillaws.TextMessage, response); err != nil { + t.Fatalf("upstream write: %v", err) + } + if err := browser.SetReadDeadline(time.Now().Add(wsTestTimeout)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + _, got, err = browser.ReadMessage() + if err != nil { + t.Fatalf("browser read: %v", err) + } + if string(got) != string(response) { + t.Fatalf("browser received %q, want %q", got, response) + } + + if err := browser.Close(); err != nil { + t.Fatalf("close browser: %v", err) + } + joinWithin(t, done, "proxyLSPConnections") + + if err := upstream.SetReadDeadline(time.Now().Add(wsTestTimeout)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + if _, _, err := upstream.ReadMessage(); err == nil { + t.Fatal("upstream connection still readable after the browser disconnected") + } +} + +// A task-host LSP crash carries an application close code; the browser needs +// that exact code to decide whether to retry or surface an install prompt. +func TestProxyLSPConnectionsPropagatesUpstreamCloseCode(t *testing.T) { + handler := &LSPHandler{logger: testLogger(), capacity: newLSPCapacityLimiter(1)} + + browser, handlerBrowserSide := newTerminalWSPair(t) + handlerUpstreamSide, upstream := newTerminalWSPair(t) + + done := spawnJoinable(t, "proxyLSPConnections", func() { _ = browser.Close() }, func() { + handler.proxyLSPConnections(handlerBrowserSide, handlerUpstreamSide, "sess-lsp", "go") + }) + + closeFrame := gorillaws.FormatCloseMessage(lspCloseInstallFailed, "install failed") + if err := upstream.WriteControl( + gorillaws.CloseMessage, closeFrame, time.Now().Add(wsTestTimeout), + ); err != nil { + t.Fatalf("upstream close: %v", err) + } + + expectLSPCloseCode(t, browser, lspCloseInstallFailed, "install failed") + joinWithin(t, done, "proxyLSPConnections") +} diff --git a/apps/backend/internal/gateway/websocket/port_proxy_handler_test.go b/apps/backend/internal/gateway/websocket/port_proxy_handler_test.go new file mode 100644 index 0000000000..16a47927a1 --- /dev/null +++ b/apps/backend/internal/gateway/websocket/port_proxy_handler_test.go @@ -0,0 +1,460 @@ +package websocket + +import ( + "context" + "encoding/json" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "sync" + "testing" + + "github.com/gin-gonic/gin" + + agentctlclient "github.com/kandev/kandev/internal/agent/runtime/agentctl" + "github.com/kandev/kandev/internal/agent/runtime/lifecycle" + "github.com/kandev/kandev/internal/common/logger" + "github.com/kandev/kandev/internal/events/bus" +) + +// recordedRequest captures what a fake agentctl saw, so proxy tests can assert +// on the rewritten path and injected headers rather than only on status codes. +type recordedRequest struct { + path string + auth string + host string + method string +} + +// fakeAgentctl is an httptest server standing in for agentctl inside an executor. +type fakeAgentctl struct { + server *httptest.Server + mu sync.Mutex + requests []recordedRequest +} + +func (f *fakeAgentctl) record(r *http.Request) { + f.mu.Lock() + defer f.mu.Unlock() + f.requests = append(f.requests, recordedRequest{ + path: r.URL.Path, + auth: r.Header.Get("Authorization"), + host: r.Host, + method: r.Method, + }) +} + +func (f *fakeAgentctl) seen() []recordedRequest { + f.mu.Lock() + defer f.mu.Unlock() + return append([]recordedRequest(nil), f.requests...) +} + +func newFakeAgentctl(t *testing.T, handler http.HandlerFunc) *fakeAgentctl { + t.Helper() + fake := &fakeAgentctl{} + fake.server = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fake.record(r) + handler(w, r) + })) + // Keep-alive-free so the reverse proxy's pooled connections do not outlive + // the test and trip this package's goroutine-leak check. + fake.server.Config.SetKeepAlivesEnabled(false) + fake.server.Start() + t.Cleanup(fake.server.Close) + return fake +} + +// gatewayResponse is the outcome of a request made against the gateway under test. +type gatewayResponse struct { + status int + body string + header http.Header +} + +// serveGateway runs handler behind a real HTTP server. httptest.ResponseRecorder +// is not a http.CloseNotifier, which httputil.ReverseProxy requires, so proxy +// paths have to be driven over a real connection. +func serveGateway(t *testing.T, handler http.Handler) func(method, path string) gatewayResponse { + t.Helper() + srv := httptest.NewUnstartedServer(handler) + srv.Config.SetKeepAlivesEnabled(false) + srv.Start() + t.Cleanup(srv.Close) + + transport := &http.Transport{DisableKeepAlives: true} + t.Cleanup(transport.CloseIdleConnections) + client := &http.Client{ + Transport: transport, + CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }, + } + + return func(method, path string) gatewayResponse { + t.Helper() + req, err := http.NewRequest(method, srv.URL+path, nil) + if err != nil { + t.Fatalf("build request %s: %v", path, err) + } + resp, err := client.Do(req) + if err != nil { + t.Fatalf("request %s: %v", path, err) + } + defer func() { _ = resp.Body.Close() }() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read body for %s: %v", path, err) + } + return gatewayResponse{status: resp.StatusCode, body: string(body), header: resp.Header} + } +} + +func newProxyTestManager(t *testing.T, log *logger.Logger) *lifecycle.Manager { + t.Helper() + return lifecycle.NewManager( + nil, + bus.NewMemoryEventBus(log), + nil, + nil, + nil, + nil, + lifecycle.ExecutorFallbackDeny, + t.TempDir(), + log, + ) +} + +// addExecutionForURL registers a synthetic execution whose agentctl client +// points at serverURL and authenticates with authToken. +func addExecutionForURL( + t *testing.T, + manager *lifecycle.Manager, + sessionID, serverURL, authToken string, + log *logger.Logger, +) *lifecycle.AgentExecution { + t.Helper() + parsed, err := url.Parse(serverURL) + if err != nil { + t.Fatalf("parse server URL: %v", err) + } + host, portString, err := net.SplitHostPort(parsed.Host) + if err != nil { + t.Fatalf("split host/port: %v", err) + } + port, err := strconv.Atoi(portString) + if err != nil { + t.Fatalf("parse port: %v", err) + } + + execution := &lifecycle.AgentExecution{ID: "exec-" + sessionID, SessionID: sessionID} + execution.SetAgentCtlClientForTesting( + agentctlclient.NewClient(host, port, log, agentctlclient.WithAuthToken(authToken)), + ) + if err := manager.ExecutionStoreForTesting().Add(execution); err != nil { + t.Fatalf("add execution: %v", err) + } + return execution +} + +func newPortProxyEngine(handler *PortProxyHandler) *gin.Engine { + gin.SetMode(gin.TestMode) + engine := gin.New() + engine.Any("/port-proxy/:sessionId/:port/*path", handler.HandlePortProxy) + return engine +} + +func decodeErrorBody(t *testing.T, body []byte) string { + t.Helper() + var parsed map[string]string + if err := json.Unmarshal(body, &parsed); err != nil { + t.Fatalf("decode error body %q: %v", body, err) + } + return parsed["error"] +} + +func TestHandlePortProxyRewritesPathAndInjectsAgentctlToken(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/javascript") + w.WriteHeader(http.StatusTeapot) + _, _ = w.Write([]byte("console.log(1)")) + }) + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-1", upstream.server.URL, "port-token", log) + gateway := serveGateway(t, newPortProxyEngine(NewPortProxyHandler(manager, log))) + + got := gateway(http.MethodGet, "/port-proxy/sess-1/3000/assets/app.js?v=2") + if got.status != http.StatusTeapot { + t.Fatalf("status = %d, want %d (upstream status must pass through)", got.status, http.StatusTeapot) + } + if got.body != "console.log(1)" { + t.Fatalf("body = %q, want the upstream body", got.body) + } + + seen := upstream.seen() + if len(seen) != 1 { + t.Fatalf("upstream requests = %d, want 1", len(seen)) + } + if seen[0].path != "/api/v1/port-proxy/3000/assets/app.js" { + t.Fatalf("upstream path = %q, want %q", seen[0].path, "/api/v1/port-proxy/3000/assets/app.js") + } + if seen[0].auth != "Bearer port-token" { + t.Fatalf("upstream Authorization = %q, want %q", seen[0].auth, "Bearer port-token") + } +} + +// Root-absolute URLs in proxied HTML must be re-anchored on the proxy prefix, +// or every asset request escapes the proxy chain and hits the gateway origin. +func TestHandlePortProxyRewritesRootAbsoluteURLsInHTML(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + _, _ = w.Write([]byte(``)) + }) + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-html", upstream.server.URL, "", log) + gateway := serveGateway(t, newPortProxyEngine(NewPortProxyHandler(manager, log))) + + got := gateway(http.MethodGet, "/port-proxy/sess-html/8080/index.html") + for _, want := range []string{ + `href="/port-proxy/sess-html/8080/static/app.css"`, + `src="/port-proxy/sess-html/8080/static/app.js"`, + } { + if !strings.Contains(got.body, want) { + t.Fatalf("rewritten body %q missing %q", got.body, want) + } + } + if seen := upstream.seen(); len(seen) != 1 || seen[0].path != "/api/v1/port-proxy/8080/index.html" { + t.Fatalf("upstream requests = %+v, want a single /api/v1/port-proxy/8080/index.html", seen) + } +} + +func TestHandlePortProxyRejectsPortsOutsideTheAllowedRange(t *testing.T) { + tests := []struct { + name string + port string + }{ + {name: "privileged", port: "80"}, + {name: "boundary below", port: "1023"}, + {name: "above range", port: "70000"}, + {name: "not a number", port: "http"}, + } + log, _ := observedTerminalLogger(t) + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-1", "http://127.0.0.1:1", "", log) + gateway := serveGateway(t, newPortProxyEngine(NewPortProxyHandler(manager, log))) + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := gateway(http.MethodGet, "/port-proxy/sess-1/"+tt.port+"/x") + if got.status != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", got.status, http.StatusBadRequest) + } + if msg := decodeErrorBody(t, []byte(got.body)); msg != "invalid port: must be 1024-65535" { + t.Fatalf("error = %q", msg) + } + }) + } +} + +func TestHandlePortProxyRequiresSessionID(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler := NewPortProxyHandler(newProxyTestManager(t, log), log) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/port-proxy//3000/", nil) + c.Params = gin.Params{{Key: "sessionId", Value: ""}, {Key: "port", Value: "3000"}} + handler.HandlePortProxy(c) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != "sessionId is required" { + t.Fatalf("error = %q", got) + } +} + +// The proxy prefix ends exactly at the port, so a request for the app root must +// still forward a "/" path to agentctl rather than an empty one. +func TestHandlePortProxyForwardsRootPathForPrefixOnlyRequests(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-root", upstream.server.URL, "", log) + handler := NewPortProxyHandler(manager, log) + + // The gin wildcard route would redirect the prefix-only form, so the handler + // is driven directly with the params it reads. + gateway := serveGateway(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c, _ := gin.CreateTestContext(w) + c.Request = r + c.Params = gin.Params{{Key: "sessionId", Value: "sess-root"}, {Key: "port", Value: "3000"}} + handler.HandlePortProxy(c) + })) + gateway(http.MethodGet, "/port-proxy/sess-root/3000") + + seen := upstream.seen() + if len(seen) != 1 || seen[0].path != "/api/v1/port-proxy/3000/" { + t.Fatalf("upstream requests = %+v, want a single /api/v1/port-proxy/3000/", seen) + } +} + +// A bare execution lookup skips the service-layer chokepoint, so the proxy must +// authorize session ownership itself before touching the per-session cache. +func TestHandlePortProxyDeniesForeignSessionBeforeResolvingProxy(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-other", upstream.server.URL, "", log) + manager.SetSessionAccessChecker(func(context.Context, string) error { + return errors.New("not your session") + }) + handler := NewPortProxyHandler(manager, log) + gateway := serveGateway(t, newPortProxyEngine(handler)) + + got := gateway(http.MethodGet, "/port-proxy/sess-other/3000/secret") + if got.status != http.StatusNotFound { + t.Fatalf("status = %d, want %d", got.status, http.StatusNotFound) + } + if msg := decodeErrorBody(t, []byte(got.body)); msg != "session not found" { + t.Fatalf("error = %q", msg) + } + if n := len(upstream.seen()); n != 0 { + t.Fatalf("upstream requests = %d, want 0 for a denied session", n) + } + handler.mu.Lock() + defer handler.mu.Unlock() + if len(handler.proxies) != 0 { + t.Fatalf("cached proxies = %d, want 0 for a denied session", len(handler.proxies)) + } +} + +func TestHandlePortProxyReportsMissingExecutionAndClient(t *testing.T) { + log, _ := observedTerminalLogger(t) + manager := newProxyTestManager(t, log) + clientless := &lifecycle.AgentExecution{ID: "exec-noclient", SessionID: "sess-noclient"} + if err := manager.ExecutionStoreForTesting().Add(clientless); err != nil { + t.Fatalf("add execution: %v", err) + } + gateway := serveGateway(t, newPortProxyEngine(NewPortProxyHandler(manager, log))) + + tests := []struct { + name string + sessionID string + wantStatus int + wantError string + }{ + { + name: "no execution", + sessionID: "sess-missing", + wantStatus: http.StatusNotFound, + wantError: "session not found or no active execution", + }, + { + name: "no agentctl client", + sessionID: "sess-noclient", + wantStatus: http.StatusServiceUnavailable, + wantError: "agentctl client not available", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := gateway(http.MethodGet, "/port-proxy/"+tt.sessionID+"/3000/") + if got.status != tt.wantStatus { + t.Fatalf("status = %d, want %d", got.status, tt.wantStatus) + } + if msg := decodeErrorBody(t, []byte(got.body)); msg != tt.wantError { + t.Fatalf("error = %q, want %q", msg, tt.wantError) + } + }) + } +} + +func TestHandlePortProxyReusesCachedProxyPerSessionAndPort(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-cache", upstream.server.URL, "", log) + handler := NewPortProxyHandler(manager, log) + gateway := serveGateway(t, newPortProxyEngine(handler)) + + for _, path := range []string{ + "/port-proxy/sess-cache/3000/a", + "/port-proxy/sess-cache/3000/b", + "/port-proxy/sess-cache/4000/a", + } { + if got := gateway(http.MethodGet, path); got.status != http.StatusOK { + t.Fatalf("status for %s = %d, want %d", path, got.status, http.StatusOK) + } + } + + handler.mu.Lock() + defer handler.mu.Unlock() + if len(handler.proxies) != 2 { + t.Fatalf("cached proxies = %d, want 2 (one per session:port)", len(handler.proxies)) + } + for _, key := range []string{"sess-cache:3000", "sess-cache:4000"} { + if _, ok := handler.proxies[key]; !ok { + t.Fatalf("cache key %q missing from %v", key, handler.proxies) + } + } +} + +// A dead upstream must surface as 502 and drop the cached proxy, so the next +// request re-resolves the (possibly moved) agentctl target. +func TestHandlePortProxyInvalidatesCacheOnUpstreamFailure(t *testing.T) { + log, _ := observedTerminalLogger(t) + dead := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + deadURL := dead.URL + dead.Close() + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-dead", deadURL, "", log) + handler := NewPortProxyHandler(manager, log) + gateway := serveGateway(t, newPortProxyEngine(handler)) + + got := gateway(http.MethodGet, "/port-proxy/sess-dead/3000/") + if got.status != http.StatusBadGateway { + t.Fatalf("status = %d, want %d", got.status, http.StatusBadGateway) + } + handler.mu.Lock() + defer handler.mu.Unlock() + if len(handler.proxies) != 0 { + t.Fatalf("cached proxies = %d, want 0 after an upstream failure", len(handler.proxies)) + } +} + +func TestPortProxyInvalidateSessionDropsOnlyThatSession(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler := NewPortProxyHandler(nil, log) + + handler.mu.Lock() + handler.proxies["sess-a:3000"] = &portProxyEntry{target: "a-3000"} + handler.proxies["sess-a:4000"] = &portProxyEntry{target: "a-4000"} + handler.proxies["sess-b:3000"] = &portProxyEntry{target: "b-3000"} + handler.mu.Unlock() + + handler.InvalidateSession("sess-a") + + handler.mu.Lock() + defer handler.mu.Unlock() + if len(handler.proxies) != 1 { + t.Fatalf("cached proxies = %d, want 1", len(handler.proxies)) + } + if _, ok := handler.proxies["sess-b:3000"]; !ok { + t.Fatalf("sess-b entry was dropped: %v", handler.proxies) + } +} diff --git a/apps/backend/internal/gateway/websocket/port_tunnel_serve_test.go b/apps/backend/internal/gateway/websocket/port_tunnel_serve_test.go new file mode 100644 index 0000000000..bec80c759a --- /dev/null +++ b/apps/backend/internal/gateway/websocket/port_tunnel_serve_test.go @@ -0,0 +1,378 @@ +package websocket + +import ( + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/kandev/kandev/internal/agent/runtime/lifecycle" + "github.com/kandev/kandev/internal/common/logger" +) + +// newTunnelClient returns a keep-alive-free client so tunnel connections do not +// outlive the test and trip this package's goroutine-leak check. +func newTunnelClient(t *testing.T) *http.Client { + t.Helper() + transport := &http.Transport{DisableKeepAlives: true} + t.Cleanup(transport.CloseIdleConnections) + return &http.Client{Transport: transport, Timeout: wsTestTimeout} +} + +// getThroughTunnel performs a GET against a tunnel port and returns status and body. +func getThroughTunnel(t *testing.T, client *http.Client, tunnelPort int, path string) (int, string) { + t.Helper() + resp, err := client.Get(fmt.Sprintf("http://127.0.0.1:%d%s", tunnelPort, path)) + if err != nil { + t.Fatalf("GET through tunnel port %d: %v", tunnelPort, err) + } + defer func() { _ = resp.Body.Close() }() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read tunnel response: %v", err) + } + return resp.StatusCode, string(body) +} + +func newTunnelManagerForTest(t *testing.T, log *logger.Logger) *TunnelManager { + t.Helper() + manager := NewTunnelManager(newProxyTestManager(t, log), log) + t.Cleanup(manager.Shutdown) + return manager +} + +func tunnelList(t *testing.T, manager *TunnelManager, sessionID string) []TunnelInfo { + t.Helper() + listed, ok := manager.ListTunnels(sessionID).([]TunnelInfo) + if !ok { + t.Fatalf("ListTunnels() returned %T, want []TunnelInfo", manager.ListTunnels(sessionID)) + } + return listed +} + +// A tunnel serves the app at "/", so the port must be re-attached on the way +// out and the original Host preserved for the app's CORS/Origin checks. +func TestStartTunnelForwardsRootPathsToAgentctlPortProxy(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte("app-response")) + }) + + lifecycleMgr := newProxyTestManager(t, log) + addExecutionForURL(t, lifecycleMgr, "sess-tunnel", upstream.server.URL, "tunnel-token", log) + manager := NewTunnelManager(lifecycleMgr, log) + t.Cleanup(manager.Shutdown) + + tunnelPort, err := manager.StartTunnel("sess-tunnel", 5173, 0) + if err != nil { + t.Fatalf("StartTunnel: %v", err) + } + if tunnelPort <= 0 { + t.Fatalf("StartTunnel() port = %d, want an OS-assigned port", tunnelPort) + } + + status, body := getThroughTunnel(t, newTunnelClient(t), tunnelPort, "/assets/index.css") + if status != http.StatusOK { + t.Fatalf("status = %d, want %d", status, http.StatusOK) + } + if body != "app-response" { + t.Fatalf("body = %q, want %q", body, "app-response") + } + + seen := upstream.seen() + if len(seen) != 1 { + t.Fatalf("upstream requests = %d, want 1", len(seen)) + } + if seen[0].path != "/api/v1/port-proxy/5173/assets/index.css" { + t.Fatalf("upstream path = %q, want %q", seen[0].path, "/api/v1/port-proxy/5173/assets/index.css") + } + if seen[0].auth != "Bearer tunnel-token" { + t.Fatalf("upstream Authorization = %q, want %q", seen[0].auth, "Bearer tunnel-token") + } + wantHost := fmt.Sprintf("127.0.0.1:%d", tunnelPort) + if seen[0].host != wantHost { + t.Fatalf("upstream Host = %q, want the tunnel host %q", seen[0].host, wantHost) + } +} + +func TestStartTunnelBindsRequestedPortAndIsIdempotent(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + lifecycleMgr := newProxyTestManager(t, log) + addExecutionForURL(t, lifecycleMgr, "sess-fixed", upstream.server.URL, "", log) + manager := NewTunnelManager(lifecycleMgr, log) + t.Cleanup(manager.Shutdown) + + // Reserving a port means binding it, reading it, and releasing it, so another + // process can always win the gap before StartTunnel re-binds. Retry rather + // than let a loaded CI machine turn that into a flake; losing every attempt + // is not plausible. Loopback-only, to match the tunnel listener and to keep + // the window as narrow as possible. + var wanted, got int + var err error + for attempt := 1; attempt <= tunnelPortBindAttempts; attempt++ { + wanted = reserveLoopbackPort(t) + got, err = manager.StartTunnel("sess-fixed", 3000, wanted) + if err == nil { + break + } + if !strings.Contains(err.Error(), "already in use") { + t.Fatalf("StartTunnel: %v", err) + } + t.Logf("attempt %d: port %d was taken between reserve and bind, retrying", attempt, wanted) + } + if err != nil { + t.Fatalf("StartTunnel never won the port race in %d attempts: %v", tunnelPortBindAttempts, err) + } + if got != wanted { + t.Fatalf("StartTunnel() port = %d, want the requested %d", got, wanted) + } + + again, err := manager.StartTunnel("sess-fixed", 3000, 0) + if err != nil { + t.Fatalf("second StartTunnel: %v", err) + } + if again != wanted { + t.Fatalf("second StartTunnel() port = %d, want the existing %d", again, wanted) + } + if listed := tunnelList(t, manager, "sess-fixed"); len(listed) != 1 { + t.Fatalf("tunnels = %d, want 1 (a duplicate start must not bind a second port)", len(listed)) + } +} + +// StartTunnel reserves its cache key before resolving; a failed resolve must +// release it, otherwise the session/port pair is permanently unusable. +func TestStartTunnelReleasesReservationWhenResolutionFails(t *testing.T) { + log, _ := observedTerminalLogger(t) + lifecycleMgr := newProxyTestManager(t, log) + clientless := &lifecycle.AgentExecution{ID: "exec-noclient", SessionID: "sess-noclient"} + if err := lifecycleMgr.ExecutionStoreForTesting().Add(clientless); err != nil { + t.Fatalf("add execution: %v", err) + } + manager := NewTunnelManager(lifecycleMgr, log) + t.Cleanup(manager.Shutdown) + + tests := []struct { + name string + sessionID string + wantError string + }{ + {name: "no execution", sessionID: "sess-missing", wantError: "session not found or no active execution"}, + {name: "no agentctl client", sessionID: "sess-noclient", wantError: "agentctl client not available"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + port, err := manager.StartTunnel(tt.sessionID, 3000, 0) + if err == nil { + t.Fatal("StartTunnel() error = nil, want a failure") + } + if !strings.Contains(err.Error(), tt.wantError) { + t.Fatalf("error = %v, want substring %q", err, tt.wantError) + } + if port != 0 { + t.Fatalf("StartTunnel() port = %d, want 0 on failure", port) + } + manager.mu.Lock() + _, stillReserved := manager.tunnels[tt.sessionID+":3000"] + manager.mu.Unlock() + if stillReserved { + t.Fatal("the reserved cache key was not released after the failure") + } + }) + } +} + +func TestStopTunnelClosesTheListenerAndForgetsTheEntry(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + lifecycleMgr := newProxyTestManager(t, log) + addExecutionForURL(t, lifecycleMgr, "sess-stop", upstream.server.URL, "", log) + manager := NewTunnelManager(lifecycleMgr, log) + t.Cleanup(manager.Shutdown) + + tunnelPort, err := manager.StartTunnel("sess-stop", 3000, 0) + if err != nil { + t.Fatalf("StartTunnel: %v", err) + } + if status, _ := getThroughTunnel(t, newTunnelClient(t), tunnelPort, "/"); status != http.StatusOK { + t.Fatalf("status before stop = %d, want %d", status, http.StatusOK) + } + + if err := manager.StopTunnel("sess-stop", 3000); err != nil { + t.Fatalf("StopTunnel: %v", err) + } + if listed := tunnelList(t, manager, "sess-stop"); len(listed) != 0 { + t.Fatalf("tunnels after stop = %v, want none", listed) + } + assertPortClosed(t, tunnelPort) + + err = manager.StopTunnel("sess-stop", 3000) + if err == nil { + t.Fatal("second StopTunnel() error = nil, want a not-found error") + } + if !strings.Contains(err.Error(), "no tunnel found for session sess-stop port 3000") { + t.Fatalf("error = %v", err) + } +} + +func TestListTunnelsIsScopedToTheSession(t *testing.T) { + log, _ := observedTerminalLogger(t) + manager := newTunnelManagerForTest(t, log) + + manager.mu.Lock() + manager.tunnels["sess-a:3000"] = &tunnelEntry{port: 3000, tunnelPort: 40001} + manager.tunnels["sess-a:4000"] = &tunnelEntry{port: 4000, tunnelPort: 40002} + manager.tunnels["sess-b:3000"] = &tunnelEntry{port: 3000, tunnelPort: 40003} + manager.mu.Unlock() + t.Cleanup(func() { + manager.mu.Lock() + manager.tunnels = make(map[string]*tunnelEntry) + manager.mu.Unlock() + }) + + listed := tunnelList(t, manager, "sess-a") + if len(listed) != 2 { + t.Fatalf("tunnels for sess-a = %d, want 2", len(listed)) + } + for _, info := range listed { + if info.TunnelPort == 40003 { + t.Fatalf("sess-b tunnel leaked into sess-a's list: %+v", listed) + } + } + if other := tunnelList(t, manager, "sess-unknown"); len(other) != 0 { + t.Fatalf("tunnels for an unknown session = %v, want none", other) + } +} + +func TestInvalidateSessionStopsOnlyThatSessionsTunnels(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + lifecycleMgr := newProxyTestManager(t, log) + addExecutionForURL(t, lifecycleMgr, "sess-gone", upstream.server.URL, "", log) + addExecutionForURL(t, lifecycleMgr, "sess-kept", upstream.server.URL, "", log) + manager := NewTunnelManager(lifecycleMgr, log) + t.Cleanup(manager.Shutdown) + + gonePort, err := manager.StartTunnel("sess-gone", 3000, 0) + if err != nil { + t.Fatalf("StartTunnel(sess-gone): %v", err) + } + keptPort, err := manager.StartTunnel("sess-kept", 3000, 0) + if err != nil { + t.Fatalf("StartTunnel(sess-kept): %v", err) + } + + manager.InvalidateSession("sess-gone") + + if listed := tunnelList(t, manager, "sess-gone"); len(listed) != 0 { + t.Fatalf("invalidated session still lists %v", listed) + } + assertPortClosed(t, gonePort) + + if listed := tunnelList(t, manager, "sess-kept"); len(listed) != 1 { + t.Fatalf("surviving session lists %v, want 1 tunnel", listed) + } + if status, _ := getThroughTunnel(t, newTunnelClient(t), keptPort, "/"); status != http.StatusOK { + t.Fatalf("surviving tunnel status = %d, want %d", status, http.StatusOK) + } +} + +func TestShutdownStopsEveryTunnel(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + lifecycleMgr := newProxyTestManager(t, log) + addExecutionForURL(t, lifecycleMgr, "sess-1", upstream.server.URL, "", log) + addExecutionForURL(t, lifecycleMgr, "sess-2", upstream.server.URL, "", log) + manager := NewTunnelManager(lifecycleMgr, log) + // Registered before the StartTunnel calls below: an early t.Fatalf there + // would otherwise strand the listeners this test already opened. + t.Cleanup(manager.Shutdown) + + first, err := manager.StartTunnel("sess-1", 3000, 0) + if err != nil { + t.Fatalf("StartTunnel(sess-1): %v", err) + } + second, err := manager.StartTunnel("sess-2", 4000, 0) + if err != nil { + t.Fatalf("StartTunnel(sess-2): %v", err) + } + + manager.Shutdown() + + manager.mu.Lock() + remaining := len(manager.tunnels) + manager.mu.Unlock() + if remaining != 0 { + t.Fatalf("tunnels after Shutdown = %d, want 0", remaining) + } + assertPortClosed(t, first) + assertPortClosed(t, second) + + // Shutdown must be idempotent: the registered cleanup calls it a second time. + manager.Shutdown() +} + +func TestTunnelProxyReportsBadGatewayWhenAgentctlIsDown(t *testing.T) { + log, _ := observedTerminalLogger(t) + dead := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + deadURL := dead.URL + dead.Close() + + lifecycleMgr := newProxyTestManager(t, log) + addExecutionForURL(t, lifecycleMgr, "sess-dead", deadURL, "", log) + manager := NewTunnelManager(lifecycleMgr, log) + t.Cleanup(manager.Shutdown) + + tunnelPort, err := manager.StartTunnel("sess-dead", 3000, 0) + if err != nil { + t.Fatalf("StartTunnel: %v", err) + } + + status, body := getThroughTunnel(t, newTunnelClient(t), tunnelPort, "/") + if status != http.StatusBadGateway { + t.Fatalf("status = %d, want %d", status, http.StatusBadGateway) + } + if !strings.Contains(body, "tunnel proxy error") { + t.Fatalf("body = %q, want the tunnel proxy error", body) + } +} + +// tunnelPortBindAttempts bounds the reserve/re-bind retry below. +const tunnelPortBindAttempts = 5 + +// reserveLoopbackPort returns a loopback port that was free a moment ago. +func reserveLoopbackPort(t *testing.T) int { + t.Helper() + probe, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("reserve a free port: %v", err) + } + port := probe.Addr().(*net.TCPAddr).Port + if err := probe.Close(); err != nil { + t.Fatalf("release probe listener: %v", err) + } + return port +} + +// assertPortClosed fails unless nothing is listening on the given local port. +func assertPortClosed(t *testing.T, port int) { + t.Helper() + conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 500*time.Millisecond) + if err == nil { + _ = conn.Close() + t.Fatalf("port %d still accepts connections, want it closed", port) + } +} diff --git a/apps/backend/internal/gateway/websocket/terminal_handler_remote_test.go b/apps/backend/internal/gateway/websocket/terminal_handler_remote_test.go new file mode 100644 index 0000000000..7eac4e3943 --- /dev/null +++ b/apps/backend/internal/gateway/websocket/terminal_handler_remote_test.go @@ -0,0 +1,316 @@ +package websocket + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + gorillaws "github.com/gorilla/websocket" + + "github.com/kandev/kandev/internal/agent/runtime/lifecycle" + "github.com/kandev/kandev/internal/common/logger" + "github.com/kandev/kandev/internal/common/scripts" +) + +// shellStartRequest mirrors the JSON body agentctl receives for shell starts. +type shellStartRequest struct { + TerminalID string `json:"terminal_id"` + Cols int `json:"cols"` + Rows int `json:"rows"` +} + +// fakeShellAgentctl serves the per-terminal shell endpoints an executor's +// agentctl exposes: a start call plus a binary WebSocket stream. +type fakeShellAgentctl struct { + server *httptest.Server + shells chan *gorillaws.Conn + + mu sync.Mutex + starts []shellStartRequest +} + +func newFakeShellAgentctl(t *testing.T, startStatus, streamStatus int) *fakeShellAgentctl { + t.Helper() + fake := &fakeShellAgentctl{shells: make(chan *gorillaws.Conn, 1)} + fake.server = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/v1/shell/terminal/start": + var req shellStartRequest + body, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(body, &req) + fake.mu.Lock() + fake.starts = append(fake.starts, req) + fake.mu.Unlock() + w.WriteHeader(startStatus) + case strings.HasSuffix(r.URL.Path, "/stream"): + if streamStatus != http.StatusOK { + w.WriteHeader(streamStatus) + return + } + conn, err := testWSUpgrader.Upgrade(w, r, nil) + if err != nil { + return + } + fake.shells <- conn + default: + w.WriteHeader(http.StatusNotFound) + } + })) + fake.server.Config.SetKeepAlivesEnabled(false) + fake.server.Start() + t.Cleanup(fake.server.Close) + return fake +} + +func (f *fakeShellAgentctl) startedShells() []shellStartRequest { + f.mu.Lock() + defer f.mu.Unlock() + return append([]shellStartRequest(nil), f.starts...) +} + +// serveOnce runs handle for a single request and returns the server plus a +// channel that closes when the handler has returned. +func serveOnce(t *testing.T, handle func(*gin.Context)) (*httptest.Server, <-chan struct{}) { + t.Helper() + gin.SetMode(gin.TestMode) + served := make(chan struct{}) + srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer close(served) + c, _ := gin.CreateTestContext(w) + c.Request = r + handle(c) + })) + srv.Config.SetKeepAlivesEnabled(false) + srv.Start() + t.Cleanup(srv.Close) + return srv, served +} + +// dialTerminal opens a terminal WebSocket against a one-shot handler server. +func dialTerminal(t *testing.T, srv *httptest.Server, path string) *gorillaws.Conn { + t.Helper() + conn, resp, err := gorillaws.DefaultDialer.Dial("ws"+strings.TrimPrefix(srv.URL, "http")+path, nil) + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + if err != nil { + t.Fatalf("dial terminal %s: %v", path, err) + } + t.Cleanup(func() { _ = conn.Close() }) + return conn +} + +// getTerminal performs a plain HTTP GET (no upgrade) against a one-shot handler. +func getTerminal(t *testing.T, srv *httptest.Server, path string) (int, string) { + t.Helper() + transport := &http.Transport{DisableKeepAlives: true} + t.Cleanup(transport.CloseIdleConnections) + client := &http.Client{Transport: transport, Timeout: wsTestTimeout} + resp, err := client.Get(srv.URL + path) + if err != nil { + t.Fatalf("GET %s: %v", path, err) + } + defer func() { _ = resp.Body.Close() }() + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + return resp.StatusCode, string(body) +} + +func remoteExecutionFor( + t *testing.T, + manager *lifecycle.Manager, + sessionID, serverURL string, + log *logger.Logger, +) *lifecycle.AgentExecution { + t.Helper() + execution := addExecutionForURL(t, manager, sessionID, serverURL, "shell-token", log) + execution.TaskEnvironmentID = "env-1" + return execution +} + +func TestHandleRemoteUserShellWSRequiresAgentctlClient(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + execution := &lifecycle.AgentExecution{ID: "exec-1", SessionID: "session-1", TaskEnvironmentID: "env-1"} + + srv, served := serveOnce(t, func(c *gin.Context) { + handler.handleRemoteUserShellWS(c, execution, "env-1", "term-1") + }) + status, body := getTerminal(t, srv, "/terminal/environment/env-1?terminalId=term-1") + joinWithin(t, served, "handleRemoteUserShellWS") + + if status != http.StatusNotFound { + t.Fatalf("status = %d, want %d", status, http.StatusNotFound) + } + if got := decodeErrorBody(t, []byte(body)); got != "remote execution not available" { + t.Fatalf("error = %q", got) + } +} + +func TestHandleRemoteUserShellWSRejectsInvalidScriptBeforeStartingAShell(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, manager := newBridgeTestHandler(t, log) + handler.scriptService = &stubScriptService{err: errors.New("no such script")} + agentctl := newFakeShellAgentctl(t, http.StatusOK, http.StatusOK) + execution := remoteExecutionFor(t, manager, "session-1", agentctl.server.URL, log) + + srv, served := serveOnce(t, func(c *gin.Context) { + handler.handleRemoteUserShellWS(c, execution, "env-1", "term-1") + }) + status, body := getTerminal(t, srv, "/terminal/environment/env-1?terminalId=term-1&scriptId=missing") + joinWithin(t, served, "handleRemoteUserShellWS") + + if status != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", status, http.StatusBadRequest) + } + if got := decodeErrorBody(t, []byte(body)); got != "invalid script ID" { + t.Fatalf("error = %q", got) + } + if starts := agentctl.startedShells(); len(starts) != 0 { + t.Fatalf("agentctl shell starts = %v, want none for an invalid script", starts) + } +} + +func TestHandleRemoteUserShellWSReportsShellStartFailure(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, manager := newBridgeTestHandler(t, log) + agentctl := newFakeShellAgentctl(t, http.StatusInternalServerError, http.StatusOK) + execution := remoteExecutionFor(t, manager, "session-1", agentctl.server.URL, log) + + srv, served := serveOnce(t, func(c *gin.Context) { + handler.handleRemoteUserShellWS(c, execution, "env-1", "term-1") + }) + status, body := getTerminal(t, srv, "/terminal/environment/env-1?terminalId=term-1") + joinWithin(t, served, "handleRemoteUserShellWS") + + if status != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", status, http.StatusServiceUnavailable) + } + if got := decodeErrorBody(t, []byte(body)); !strings.Contains(got, "start shell terminal failed") { + t.Fatalf("error = %q, want the agentctl start failure", got) + } +} + +func TestHandleRemoteUserShellWSReportsStreamFailure(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, manager := newBridgeTestHandler(t, log) + agentctl := newFakeShellAgentctl(t, http.StatusOK, http.StatusNotFound) + execution := remoteExecutionFor(t, manager, "session-1", agentctl.server.URL, log) + + srv, served := serveOnce(t, func(c *gin.Context) { + handler.handleRemoteUserShellWS(c, execution, "env-1", "term-1") + }) + status, body := getTerminal(t, srv, "/terminal/environment/env-1?terminalId=term-1") + joinWithin(t, served, "handleRemoteUserShellWS") + + if status != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", status, http.StatusServiceUnavailable) + } + if got := decodeErrorBody(t, []byte(body)); !strings.Contains(got, "shell terminal stream") { + t.Fatalf("error = %q, want the stream connect failure", got) + } + if starts := agentctl.startedShells(); len(starts) != 1 { + t.Fatalf("agentctl shell starts = %v, want exactly 1", starts) + } +} + +// The end-to-end remote shell path: start the terminal on agentctl, send the +// script's initial command, then bridge keystrokes and output in both directions. +func TestHandleRemoteUserShellWSBridgesAgentctlShell(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, manager := newBridgeTestHandler(t, log) + handler.scriptService = &stubScriptService{ + script: &scripts.RepositoryScript{Name: "Dev server", Command: "npm run dev"}, + } + agentctl := newFakeShellAgentctl(t, http.StatusOK, http.StatusOK) + execution := remoteExecutionFor(t, manager, "session-1", agentctl.server.URL, log) + + srv, served := serveOnce(t, func(c *gin.Context) { + handler.handleRemoteUserShellWS(c, execution, "env-1", "term-1") + }) + browser := dialTerminal(t, srv, "/terminal/environment/env-1?terminalId=term-1&scriptId=script-1") + + shellTimer := time.NewTimer(wsTestTimeout) + defer shellTimer.Stop() + var shell *gorillaws.Conn + select { + case shell = <-agentctl.shells: + case <-shellTimer.C: + t.Fatal("agentctl never received a shell stream connection") + } + t.Cleanup(func() { _ = shell.Close() }) + + starts := agentctl.startedShells() + if len(starts) != 1 { + t.Fatalf("agentctl shell starts = %v, want exactly 1", starts) + } + if starts[0].TerminalID != "term-1" || starts[0].Cols != 80 || starts[0].Rows != 24 { + t.Fatalf("shell start = %+v, want term-1 at 80x24", starts[0]) + } + + if got := readBinary(t, shell); string(got) != "npm run dev\n" { + t.Fatalf("initial command sent upstream = %q, want %q", got, "npm run dev\n") + } + + if err := browser.WriteMessage(gorillaws.BinaryMessage, []byte("ls\r")); err != nil { + t.Fatalf("browser write: %v", err) + } + if got := readBinary(t, shell); string(got) != "ls\r" { + t.Fatalf("keystrokes forwarded upstream = %q, want %q", got, "ls\r") + } + + if err := shell.WriteMessage(gorillaws.BinaryMessage, []byte("file.txt\r\n")); err != nil { + t.Fatalf("shell write: %v", err) + } + if got := readBinary(t, browser); string(got) != "file.txt\r\n" { + t.Fatalf("shell output forwarded to browser = %q, want %q", got, "file.txt\r\n") + } + + if err := browser.Close(); err != nil { + t.Fatalf("close browser: %v", err) + } + joinWithin(t, served, "handleRemoteUserShellWS") +} + +// The local (host-executor) shell path: start a shell process for the terminal, +// upgrade, and bridge until the browser disconnects. +func TestHandleUserShellWSStartsShellAndBridgesUntilDisconnect(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + t.Cleanup(func() { _ = runner.StopUserShell(context.Background(), "env-1", "term-1") }) + + execution := &lifecycle.AgentExecution{ + ID: "exec-1", + SessionID: "session-1", + TaskEnvironmentID: "env-1", + WorkspacePath: t.TempDir(), + } + srv, served := serveOnce(t, func(c *gin.Context) { + handler.handleUserShellWS(c, execution, "env-1", "term-1", runner) + }) + browser := dialTerminal(t, srv, "/terminal/environment/env-1?terminalId=term-1&label=Logs") + + shells := runner.ListUserShells("env-1") + if len(shells) != 1 { + t.Fatalf("user shells = %v, want exactly 1", shells) + } + if shells[0].Label != "Logs" { + t.Fatalf("shell label = %q, want %q", shells[0].Label, "Logs") + } + + if err := browser.Close(); err != nil { + t.Fatalf("close browser: %v", err) + } + joinWithin(t, served, "handleUserShellWS") +} diff --git a/apps/backend/internal/gateway/websocket/terminal_handler_routes_test.go b/apps/backend/internal/gateway/websocket/terminal_handler_routes_test.go new file mode 100644 index 0000000000..7b078f0c45 --- /dev/null +++ b/apps/backend/internal/gateway/websocket/terminal_handler_routes_test.go @@ -0,0 +1,316 @@ +package websocket + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/kandev/kandev/internal/agent/runtime/lifecycle" + "github.com/kandev/kandev/internal/common/scripts" +) + +// stubScriptService returns a fixed repository script (or error) for any ID. +type stubScriptService struct { + script *scripts.RepositoryScript + err error + calls int +} + +func (s *stubScriptService) GetRepositoryScript(context.Context, string) (*scripts.RepositoryScript, error) { + s.calls++ + return s.script, s.err +} + +type stubUserService struct { + shell string + err error +} + +func (s stubUserService) PreferredShell(context.Context) (string, error) { + return s.shell, s.err +} + +// terminalContext builds a gin context for a terminal route with the wildcard +// :target param the real route uses. +func terminalContext(t *testing.T, target, query string) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/terminal"+target+query, nil) + c.Params = gin.Params{{Key: "target", Value: target}} + return c, recorder +} + +func TestHandleTerminalWSRejectsRoutesWithoutATarget(t *testing.T) { + const wantError = "terminal route must be /environment/:environmentId for shell terminals " + + "or /session/:sessionId?mode=agent for agent terminals" + + for _, target := range []string{"/", "/legacy-session-id", "/session/", "/environment/"} { + t.Run(target, func(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + c, recorder := terminalContext(t, target, "") + + handler.HandleTerminalWS(c) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != wantError { + t.Fatalf("error = %q, want %q", got, wantError) + } + }) + } +} + +func TestHandleTerminalWSValidatesEnvironmentRouteQuery(t *testing.T) { + tests := []struct { + name string + query string + wantError string + }{ + { + name: "agent mode is not a shell", + query: "?mode=agent&terminalId=t-1", + wantError: "environment terminal route only supports shell mode", + }, + { + name: "missing terminal id", + query: "?mode=shell", + wantError: "terminalId required for shell mode", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + c, recorder := terminalContext(t, "/environment/env-1", tt.query) + + handler.HandleTerminalWS(c) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != tt.wantError { + t.Fatalf("error = %q, want %q", got, tt.wantError) + } + }) + } +} + +func TestHandleTerminalWSSessionRouteReportsMissingInteractiveRunner(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + c, recorder := terminalContext(t, "/session/session-1", "?mode=agent") + + handler.HandleTerminalWS(c) + + if recorder.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusServiceUnavailable) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != "interactive runner not available" { + t.Fatalf("error = %q", got) + } +} + +func TestHandleAgentTerminalRouteRequiresSessionID(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + c, recorder := terminalContext(t, "/session/", "?mode=agent") + + handler.handleAgentTerminalRoute(c, "") + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != "sessionId is required" { + t.Fatalf("error = %q", got) + } +} + +func TestHandleTerminalWSEnvironmentRouteReportsUnreadyExecution(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + c, recorder := terminalContext(t, "/environment/env-unknown", "?terminalId=t-1") + + handler.HandleTerminalWS(c) + + if recorder.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusServiceUnavailable) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got == "" { + t.Fatal("expected the lifecycle error to be reported to the caller") + } +} + +// The passthrough terminal reaches the lifecycle manager directly, so a denied +// session must be refused before any WebSocket upgrade happens. +func TestHandleAgentPassthroughWSRefusesDeniedSession(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, manager := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + manager.SetSessionAccessChecker(func(context.Context, string) error { + return errors.New("session belongs to another user") + }) + + c, recorder := terminalContext(t, "/session/session-denied", "?mode=agent") + handler.handleAgentPassthroughWS(c, "session-denied", runner) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != "session belongs to another user" { + t.Fatalf("error = %q", got) + } + if recorder.Header().Get("Upgrade") != "" { + t.Fatal("a denied session must not be upgraded to a WebSocket") + } +} + +func TestResolveShellLabelPrefersScriptOverLabel(t *testing.T) { + log, _ := observedTerminalLogger(t) + scriptService := &stubScriptService{ + script: &scripts.RepositoryScript{Name: "Dev server", Command: "npm run dev"}, + } + handler, _ := newBridgeTestHandler(t, log) + handler.scriptService = scriptService + + c, _ := terminalContext(t, "/environment/env-1", "?scriptId=script-1&label=ignored") + label, command, httpErr := handler.resolveShellLabel(c) + + if httpErr != "" { + t.Fatalf("resolveShellLabel() httpErr = %q, want none", httpErr) + } + if label != "Dev server" || command != "npm run dev" { + t.Fatalf("resolveShellLabel() = (%q, %q), want (%q, %q)", + label, command, "Dev server", "npm run dev") + } + if scriptService.calls != 1 { + t.Fatalf("script lookups = %d, want 1", scriptService.calls) + } +} + +func TestResolveShellLabelReportsInvalidScript(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + handler.scriptService = &stubScriptService{err: errors.New("no such script")} + + c, _ := terminalContext(t, "/environment/env-1", "?scriptId=missing") + label, command, httpErr := handler.resolveShellLabel(c) + + if httpErr != "invalid script ID" { + t.Fatalf("resolveShellLabel() httpErr = %q, want %q", httpErr, "invalid script ID") + } + if label != "" || command != "" { + t.Fatalf("resolveShellLabel() = (%q, %q), want empty values on failure", label, command) + } +} + +func TestResolveShellLabelFallsBackToLabelQuery(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + + c, _ := terminalContext(t, "/environment/env-1", "?label=Logs") + if label, command, httpErr := handler.resolveShellLabel(c); label != "Logs" || command != "" || httpErr != "" { + t.Fatalf("resolveShellLabel() = (%q, %q, %q), want (\"Logs\", \"\", \"\")", label, command, httpErr) + } + + bare, _ := terminalContext(t, "/environment/env-1", "") + if label, command, httpErr := handler.resolveShellLabel(bare); label != "" || command != "" || httpErr != "" { + t.Fatalf("resolveShellLabel() = (%q, %q, %q), want all empty", label, command, httpErr) + } +} + +func TestResolvePreferredShellFallsBackToDefault(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + ctx := context.Background() + + if got := handler.resolvePreferredShell(ctx); got != "" { + t.Fatalf("resolvePreferredShell() with no user service = %q, want empty", got) + } + + handler.userService = stubUserService{err: errors.New("settings unavailable")} + if got := handler.resolvePreferredShell(ctx); got != "" { + t.Fatalf("resolvePreferredShell() on error = %q, want empty", got) + } + + handler.userService = stubUserService{shell: "/bin/zsh"} + if got := handler.resolvePreferredShell(ctx); got != "/bin/zsh" { + t.Fatalf("resolvePreferredShell() = %q, want %q", got, "/bin/zsh") + } +} + +func TestResolveWorkingDirRequiresAMaterializedWorkspace(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + + if _, err := handler.resolveWorkingDir(nil); err == nil { + t.Fatal("resolveWorkingDir(nil) error = nil, want a failure") + } + + _, err := handler.resolveWorkingDir(&lifecycle.AgentExecution{TaskEnvironmentID: "env-1"}) + if !errors.Is(err, lifecycle.ErrSessionWorkspaceNotReady) { + t.Fatalf("resolveWorkingDir() error = %v, want ErrSessionWorkspaceNotReady", err) + } + + dir := t.TempDir() + got, err := handler.resolveWorkingDir(&lifecycle.AgentExecution{WorkspacePath: dir}) + if err != nil { + t.Fatalf("resolveWorkingDir() error = %v", err) + } + if got != dir { + t.Fatalf("resolveWorkingDir() = %q, want %q", got, dir) + } +} + +func TestHandleUserShellWSReportsWorkspaceNotReady(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + c, recorder := terminalContext(t, "/environment/env-1", "?terminalId=t-1") + execution := &lifecycle.AgentExecution{ID: "exec-1", SessionID: "session-1", TaskEnvironmentID: "env-1"} + handler.handleUserShellWS(c, execution, "env-1", "t-1", runner) + + if recorder.Code != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusServiceUnavailable) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got == "" { + t.Fatal("expected the workspace-not-ready reason to be reported") + } + if shells := runner.ListUserShells("env-1"); len(shells) != 0 { + t.Fatalf("user shells = %v, want none when the workspace is not ready", shells) + } +} + +func TestHandleUserShellWSRejectsInvalidScriptBeforeStartingAShell(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + handler.scriptService = &stubScriptService{err: errors.New("no such script")} + runner := newBridgeTestRunner(t, log) + + c, recorder := terminalContext(t, "/environment/env-1", "?terminalId=t-1&scriptId=missing") + execution := &lifecycle.AgentExecution{ + ID: "exec-1", + SessionID: "session-1", + TaskEnvironmentID: "env-1", + WorkspacePath: t.TempDir(), + } + handler.handleUserShellWS(c, execution, "env-1", "t-1", runner) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != "invalid script ID" { + t.Fatalf("error = %q", got) + } + if shells := runner.ListUserShells("env-1"); len(shells) != 0 { + t.Fatalf("user shells = %v, want none for an invalid script", shells) + } +} diff --git a/apps/backend/internal/gateway/websocket/terminal_pumps_bridge_test.go b/apps/backend/internal/gateway/websocket/terminal_pumps_bridge_test.go new file mode 100644 index 0000000000..5b5cc8528e --- /dev/null +++ b/apps/backend/internal/gateway/websocket/terminal_pumps_bridge_test.go @@ -0,0 +1,293 @@ +package websocket + +import ( + "errors" + "io" + "testing" + "time" + + gorillaws "github.com/gorilla/websocket" + + "github.com/kandev/kandev/internal/agent/runtime/lifecycle" + "github.com/kandev/kandev/internal/agentctl/server/process" + "github.com/kandev/kandev/internal/common/logger" + "github.com/kandev/kandev/internal/events/bus" +) + +// newBridgeTestHandler builds a TerminalHandler backed by a real (dependency-free) +// lifecycle manager, which is all the pump code needs. +func newBridgeTestHandler(t *testing.T, log *logger.Logger) (*TerminalHandler, *lifecycle.Manager) { + t.Helper() + manager := lifecycle.NewManager( + nil, + bus.NewMemoryEventBus(log), + nil, + nil, + nil, + nil, + lifecycle.ExecutorFallbackDeny, + t.TempDir(), + log, + ) + return NewTerminalHandler(manager, nil, nil, log), manager +} + +func newBridgeTestRunner(t *testing.T, log *logger.Logger) *process.InteractiveRunner { + t.Helper() + return process.NewInteractiveRunner(nil, log, 2*1024*1024) +} + +func TestWsWriterSendsBinaryFramesAndRefusesWritesAfterClose(t *testing.T) { + browser, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + payload := []byte("\x1b[32mready\x1b[0m$ ") + n, err := wsw.Write(payload) + if err != nil { + t.Fatalf("Write() error = %v", err) + } + if n != len(payload) { + t.Fatalf("Write() n = %d, want %d", n, len(payload)) + } + if got := readBinary(t, browser); string(got) != string(payload) { + t.Fatalf("browser received %q, want %q", got, payload) + } + + if err := wsw.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + n, err = wsw.Write([]byte("after close")) + if !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("Write() after Close error = %v, want io.ErrClosedPipe", err) + } + if n != 0 { + t.Fatalf("Write() after Close n = %d, want 0", n) + } + + // Nothing may reach the peer once the writer is closed. + if err := browser.SetReadDeadline(time.Now().Add(150 * time.Millisecond)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + if _, data, err := browser.ReadMessage(); err == nil { + t.Fatalf("browser received %q after wsWriter.Close(), want no frame", data) + } +} + +func TestWsWriterReportsUnderlyingConnectionFailure(t *testing.T) { + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + if err := server.Close(); err != nil { + t.Fatalf("close server conn: %v", err) + } + n, err := wsw.Write([]byte("orphaned")) + if err == nil { + t.Fatal("Write() on a closed connection returned nil error") + } + if n != 0 { + t.Fatalf("Write() n = %d, want 0 on failure", n) + } +} + +func TestBridgeBinaryWebSocketsForwardsBothDirections(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + + browser, bridgeClientSide := newTerminalWSPair(t) + bridgeUpstreamSide, agentctl := newTerminalWSPair(t) + + done := spawnJoinable(t, "bridgeBinaryWebSockets", func() { _ = browser.Close() }, func() { + handler.bridgeBinaryWebSockets(bridgeClientSide, bridgeUpstreamSide, "term-1") + }) + + // client → agentctl (keystrokes) + if err := browser.WriteMessage(gorillaws.BinaryMessage, []byte("echo hi\r")); err != nil { + t.Fatalf("browser write: %v", err) + } + if got := readBinary(t, agentctl); string(got) != "echo hi\r" { + t.Fatalf("agentctl received %q, want %q", got, "echo hi\r") + } + + // agentctl → client (shell output) + if err := agentctl.WriteMessage(gorillaws.BinaryMessage, []byte("hi\r\n$ ")); err != nil { + t.Fatalf("agentctl write: %v", err) + } + if got := readBinary(t, browser); string(got) != "hi\r\n$ " { + t.Fatalf("browser received %q, want %q", got, "hi\r\n$ ") + } + + // A client disconnect must tear the upstream connection down too, otherwise + // the agentctl→client goroutine blocks forever and the bridge never returns. + if err := browser.Close(); err != nil { + t.Fatalf("close browser conn: %v", err) + } + joinWithin(t, done, "bridgeBinaryWebSockets") + + if err := agentctl.SetReadDeadline(time.Now().Add(wsTestTimeout)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + if _, _, err := agentctl.ReadMessage(); err == nil { + t.Fatal("agentctl connection still readable after the client disconnected") + } +} + +func TestBridgeBinaryWebSocketsReturnsWhenUpstreamDisconnects(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + + browser, bridgeClientSide := newTerminalWSPair(t) + bridgeUpstreamSide, agentctl := newTerminalWSPair(t) + + done := spawnJoinable(t, "bridgeBinaryWebSockets", func() { _ = agentctl.Close() }, func() { + handler.bridgeBinaryWebSockets(bridgeClientSide, bridgeUpstreamSide, "term-2") + }) + + if err := agentctl.Close(); err != nil { + t.Fatalf("close agentctl conn: %v", err) + } + joinWithin(t, done, "bridgeBinaryWebSockets") + _ = browser.Close() +} + +func TestRunTerminalBridgeReleasesSessionOutputOnDisconnect(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + browser, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + + const sessionID = "session-bridge" + runner.TrackSessionWebSocket(sessionID, wsw) + if !runner.HasActiveWebSocketBySession(sessionID) { + t.Fatal("session WebSocket was not tracked before the bridge started") + } + + // "proc-gone" is deliberately unknown to the runner: PTY setup fails, so the + // bridge exercises its drop-input path without spawning a real PTY. + done := spawnJoinable(t, "runTerminalBridge", func() { _ = browser.Close() }, func() { + handler.runTerminalBridge(server, sessionID, "proc-gone", runner, wsw) + }) + + if err := browser.WriteMessage(gorillaws.BinaryMessage, nil); err != nil { + t.Fatalf("write empty frame: %v", err) + } + if err := browser.WriteMessage(gorillaws.BinaryMessage, []byte("ls\n")); err != nil { + t.Fatalf("write input frame: %v", err) + } + if err := browser.Close(); err != nil { + t.Fatalf("close browser conn: %v", err) + } + joinWithin(t, done, "runTerminalBridge") + + if runner.HasActiveWebSocketBySession(sessionID) { + t.Fatal("session direct output was not cleared when the terminal disconnected") + } + if _, err := wsw.Write([]byte("late")); !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("wsWriter usable after bridge exit: err = %v, want io.ErrClosedPipe", err) + } + if got := countLogged(recorded, "PTY not ready, dropping input"); got != 1 { + t.Fatalf("dropped-input records = %d, want 1 (the empty frame must be skipped)", got) + } +} + +func TestRunTerminalBridgeResizeFallsBackWhenProcessAndSessionAreGone(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + browser, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + + done := spawnJoinable(t, "runTerminalBridge", func() { _ = browser.Close() }, func() { + handler.runTerminalBridge(server, "session-resize", "proc-gone", runner, wsw) + }) + + resize := append([]byte{resizeCommandByte}, []byte(`{"cols":120,"rows":40}`)...) + if err := browser.WriteMessage(gorillaws.BinaryMessage, resize); err != nil { + t.Fatalf("write resize frame: %v", err) + } + if err := browser.Close(); err != nil { + t.Fatalf("close browser conn: %v", err) + } + joinWithin(t, done, "runTerminalBridge") + + if got := countLogged(recorded, "failed to resize PTY (both process and session)"); got != 1 { + t.Fatalf("resize-fallback records = %d, want 1", got) + } +} + +func TestRunUserShellBridgeClearsShellDirectOutputOnDisconnect(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + const ( + scopeID = "env-shell" + terminalID = "term-shell" + sessionID = "session-shell" + ) + info := startDeferredUserShell(t, runner, scopeID, sessionID, terminalID) + + browser, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + if err := runner.SetDirectOutput(info.ID, wsw); err != nil { + t.Fatalf("SetDirectOutput: %v", err) + } + if !runner.HasActiveWebSocket(info.ID) { + t.Fatal("direct output was not registered on the user shell process") + } + + done := spawnJoinable(t, "runUserShellBridge", func() { _ = browser.Close() }, func() { + handler.runUserShellBridge(server, sessionID, scopeID, terminalID, info.ID, runner, wsw) + }) + + if err := browser.WriteMessage(gorillaws.BinaryMessage, nil); err != nil { + t.Fatalf("write empty frame: %v", err) + } + if err := browser.WriteMessage(gorillaws.BinaryMessage, []byte("pwd\n")); err != nil { + t.Fatalf("write input frame: %v", err) + } + if err := browser.Close(); err != nil { + t.Fatalf("close browser conn: %v", err) + } + joinWithin(t, done, "runUserShellBridge") + + if runner.HasActiveWebSocket(info.ID) { + t.Fatal("user shell direct output was not cleared when the terminal disconnected") + } + if _, err := wsw.Write([]byte("late")); !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("wsWriter usable after bridge exit: err = %v, want io.ErrClosedPipe", err) + } + if got := countLogged(recorded, "PTY not ready, dropping input"); got != 1 { + t.Fatalf("dropped-input records = %d, want 1 (the empty frame must be skipped)", got) + } +} + +func TestRunUserShellBridgeReportsUnresizableShell(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + browser, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + + done := spawnJoinable(t, "runUserShellBridge", func() { _ = browser.Close() }, func() { + handler.runUserShellBridge(server, "session-x", "env-x", "term-unknown", "proc-gone", runner, wsw) + }) + + resize := append([]byte{resizeCommandByte}, []byte(`{"cols":100,"rows":30}`)...) + if err := browser.WriteMessage(gorillaws.BinaryMessage, resize); err != nil { + t.Fatalf("write resize frame: %v", err) + } + if err := browser.Close(); err != nil { + t.Fatalf("close browser conn: %v", err) + } + joinWithin(t, done, "runUserShellBridge") + + if got := countLogged(recorded, "failed to resize user shell PTY"); got != 1 { + t.Fatalf("user-shell resize failure records = %d, want 1", got) + } +} diff --git a/apps/backend/internal/gateway/websocket/terminal_pumps_pty_test.go b/apps/backend/internal/gateway/websocket/terminal_pumps_pty_test.go new file mode 100644 index 0000000000..360b50d869 --- /dev/null +++ b/apps/backend/internal/gateway/websocket/terminal_pumps_pty_test.go @@ -0,0 +1,514 @@ +package websocket + +import ( + "bytes" + "context" + "errors" + "io" + "testing" + + "github.com/kandev/kandev/internal/agent/runtime/lifecycle" + "github.com/kandev/kandev/internal/agentctl/server/process" + v1 "github.com/kandev/kandev/pkg/api/v1" +) + +// failingWriter stands in for a PTY whose process died underneath the bridge. +type failingWriter struct{ writes int } + +func (w *failingWriter) Write(p []byte) (int, error) { + w.writes++ + return 0, errors.New("pty gone") +} + +// startDeferredUserShell registers a user shell whose PTY start is deferred to +// the first resize, so no OS process is ever spawned. +func startDeferredUserShell( + t *testing.T, + runner *process.InteractiveRunner, + scopeID, sessionID, terminalID string, +) *process.InteractiveProcessInfo { + t.Helper() + info, err := runner.StartUserShell( + context.Background(), scopeID, sessionID, terminalID, t.TempDir(), "", &process.UserShellOptions{Label: "Terminal"}, + ) + if err != nil { + t.Fatalf("StartUserShell: %v", err) + } + t.Cleanup(func() { _ = runner.StopUserShell(context.Background(), scopeID, terminalID) }) + return info +} + +// startDeferredPassthrough registers an agent passthrough process whose PTY +// start is deferred, so no OS process is ever spawned. +func startDeferredPassthrough( + t *testing.T, + runner *process.InteractiveRunner, + sessionID string, +) *process.InteractiveProcessInfo { + t.Helper() + info, err := runner.Start(context.Background(), process.InteractiveStartRequest{ + SessionID: sessionID, + Command: []string{"/bin/sh"}, + WorkingDir: t.TempDir(), + DisableTurnDetection: true, + }) + if err != nil { + t.Fatalf("Start passthrough process: %v", err) + } + t.Cleanup(func() { _ = runner.Stop(context.Background(), info.ID) }) + return info +} + +func TestSetupPtyAccessCommonWiresWriterAndDirectOutput(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + info := startDeferredPassthrough(t, runner, "session-setup") + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + pty := &bytes.Buffer{} + var ptyWriter io.Writer + ok := handler.setupPtyAccessCommon( + "session-setup", info.ID, runner, wsw, &ptyWriter, + func() (io.Writer, error) { return pty, nil }, + ) + if !ok { + t.Fatal("setupPtyAccessCommon() = false, want true") + } + if ptyWriter != io.Writer(pty) { + t.Fatalf("ptyWriter = %#v, want the writer returned by getWriter", ptyWriter) + } + if !runner.HasActiveWebSocket(info.ID) { + t.Fatal("direct output was not registered on the process") + } +} + +func TestSetupPtyAccessCommonLeavesWriterUnsetWhenPtyUnavailable(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + info := startDeferredPassthrough(t, runner, "session-nopty") + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + var ptyWriter io.Writer + ok := handler.setupPtyAccessCommon( + "session-nopty", info.ID, runner, wsw, &ptyWriter, + func() (io.Writer, error) { return nil, errors.New("pty not ready") }, + ) + if ok { + t.Fatal("setupPtyAccessCommon() = true, want false when the PTY writer is unavailable") + } + if ptyWriter != nil { + t.Fatalf("ptyWriter = %#v, want nil", ptyWriter) + } + if runner.HasActiveWebSocket(info.ID) { + t.Fatal("direct output must not be registered when PTY setup fails") + } + if got := countLogged(recorded, "PTY writer not ready yet"); got != 1 { + t.Fatalf("not-ready records = %d, want 1", got) + } +} + +// A PTY writer must not be published to the bridge when direct output could not +// be attached: the bridge would then forward keystrokes into a process whose +// output nobody is reading. +func TestSetupPtyAccessCommonRejectsWhenDirectOutputCannotBeSet(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + pty := &bytes.Buffer{} + var ptyWriter io.Writer + ok := handler.setupPtyAccessCommon( + "session-gone", "proc-gone", runner, wsw, &ptyWriter, + func() (io.Writer, error) { return pty, nil }, + ) + if ok { + t.Fatal("setupPtyAccessCommon() = true, want false when SetDirectOutput fails") + } + if ptyWriter != nil { + t.Fatalf("ptyWriter = %#v, want nil when direct output could not be set", ptyWriter) + } + if got := countLogged(recorded, "failed to set direct output"); got != 1 { + t.Fatalf("direct-output failure records = %d, want 1", got) + } +} + +func TestSetupPtyAccessVariantsRejectDeferredProcesses(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + passthrough := startDeferredPassthrough(t, runner, "session-deferred") + var ptyWriter io.Writer + if handler.setupPtyAccess("session-deferred", passthrough.ID, runner, wsw, &ptyWriter) { + t.Fatal("setupPtyAccess() = true for a process that has not started") + } + + shell := startDeferredUserShell(t, runner, "env-deferred", "session-deferred", "term-deferred") + var shellWriter io.Writer + if handler.setupUserShellPtyAccess( + "session-deferred", "env-deferred", "term-deferred", shell.ID, runner, wsw, &shellWriter, + ) { + t.Fatal("setupUserShellPtyAccess() = true for a shell that has not started") + } +} + +func TestHandleResizeRejectsMalformedAndZeroDimensions(t *testing.T) { + tests := []struct { + name string + payload string + wantLog string + }{ + {name: "malformed json", payload: `{"cols":`, wantLog: "failed to parse resize command"}, + {name: "zero cols", payload: `{"cols":0,"rows":40}`, wantLog: "invalid resize dimensions"}, + {name: "zero rows", payload: `{"cols":80,"rows":0}`, wantLog: "invalid resize dimensions"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + got := handler.handleResize([]byte(tt.payload), "session-1", "proc-1", runner) + if got != "proc-1" { + t.Fatalf("handleResize() = %q, want the unchanged process id", got) + } + if n := countLogged(recorded, tt.wantLog); n != 1 { + t.Fatalf("%q records = %d, want 1", tt.wantLog, n) + } + if n := countLogged(recorded, "PTY resized"); n != 0 { + t.Fatalf("PTY resized records = %d, want 0", n) + } + }) + } +} + +func TestHandleResizeKeepsProcessIDWhenSessionFallbackFails(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + got := handler.handleResize([]byte(`{"cols":80,"rows":24}`), "session-missing", "proc-missing", runner) + if got != "proc-missing" { + t.Fatalf("handleResize() = %q, want %q", got, "proc-missing") + } + if n := countLogged(recorded, "failed to resize PTY (both process and session)"); n != 1 { + t.Fatalf("resize-failure records = %d, want 1", n) + } +} + +func TestDiscoverProcessBySessionAdoptsRestartedProcess(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + info := startDeferredPassthrough(t, runner, "session-restart") + + got := handler.discoverProcessBySession("session-restart", "stale-proc", runner) + if got != info.ID { + t.Fatalf("discoverProcessBySession() = %q, want the live process id %q", got, info.ID) + } + + // Same ID in, same ID out — the bridge must not churn state on a no-op. + if same := handler.discoverProcessBySession("session-restart", info.ID, runner); same != info.ID { + t.Fatalf("discoverProcessBySession() = %q, want %q", same, info.ID) + } +} + +func TestDiscoverProcessBySessionKeepsOldIDWhenSessionHasNoProcess(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + if got := handler.discoverProcessBySession("session-empty", "stale-proc", runner); got != "stale-proc" { + t.Fatalf("discoverProcessBySession() = %q, want %q", got, "stale-proc") + } +} + +func TestResizeBySessionFallbackKeepsOldProcessIDOnFailure(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + got := handler.resizeBySessionFallback("session-empty", "old-proc", 80, 24, runner) + if got != "old-proc" { + t.Fatalf("resizeBySessionFallback() = %q, want %q", got, "old-proc") + } + if n := countLogged(recorded, "failed to resize PTY (both process and session)"); n != 1 { + t.Fatalf("resize-failure records = %d, want 1", n) + } +} + +// Enter is the signal that the user submitted a command; only then may the +// passthrough execution be promoted to RUNNING. +func TestDetectInputSubmissionMarksPassthroughRunningOnlyOnEnter(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, manager := newBridgeTestHandler(t, log) + + execution := &lifecycle.AgentExecution{ + ID: "exec-1", + SessionID: "session-enter", + PassthroughProcessID: "proc-1", + Status: v1.AgentStatusStarting, + } + if err := manager.ExecutionStoreForTesting().Add(execution); err != nil { + t.Fatalf("add execution: %v", err) + } + + handler.detectInputSubmission("session-enter", []byte("echo hello")) + if execution.Status != v1.AgentStatusStarting { + t.Fatalf("status = %q after keystrokes without Enter, want %q", + execution.Status, v1.AgentStatusStarting) + } + + handler.detectInputSubmission("session-enter", []byte("\r")) + if execution.Status != v1.AgentStatusRunning { + t.Fatalf("status = %q after Enter, want %q", execution.Status, v1.AgentStatusRunning) + } +} + +func TestDetectInputSubmissionSurvivesUnknownSession(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + + handler.detectInputSubmission("session-unknown", []byte("ls\n")) + if n := countLogged(recorded, "failed to mark passthrough as running"); n != 1 { + t.Fatalf("mark-failure records = %d, want 1", n) + } +} + +func TestWriteToReadyPtyForwardsInputVerbatim(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + pty := &bytes.Buffer{} + var ptyWriter io.Writer = pty + gotID, gotDirect := handler.writeToReadyPty( + []byte("git status"), "session-1", "proc-1", runner, wsw, &ptyWriter, true, + ) + if gotID != "proc-1" || !gotDirect { + t.Fatalf("writeToReadyPty() = (%q, %v), want (\"proc-1\", true)", gotID, gotDirect) + } + if pty.String() != "git status" { + t.Fatalf("PTY received %q, want %q", pty.String(), "git status") + } +} + +// A dead PTY must reset bridge state immediately and re-discover the session's +// current process, instead of waiting for a resize that may never arrive. +func TestWriteToReadyPtyResetsAndRediscoversAfterWriteFailure(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + replacement := startDeferredPassthrough(t, runner, "session-dead") + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + pty := &failingWriter{} + var ptyWriter io.Writer = pty + gotID, gotDirect := handler.writeToReadyPty( + []byte("ls\n"), "session-dead", "stale-proc", runner, wsw, &ptyWriter, true, + ) + if gotID != replacement.ID { + t.Fatalf("writeToReadyPty() process id = %q, want the rediscovered %q", gotID, replacement.ID) + } + if gotDirect { + t.Fatal("writeToReadyPty() directOutputSet = true, want false after a PTY write failure") + } + if ptyWriter != nil { + t.Fatalf("ptyWriter = %#v, want nil after a PTY write failure", ptyWriter) + } + if pty.writes != 1 { + t.Fatalf("PTY writes = %d, want exactly 1", pty.writes) + } + if n := countLogged(recorded, "PTY write error"); n != 1 { + t.Fatalf("PTY write error records = %d, want 1", n) + } + // The failed keystroke must not be reported as a submitted command. + if n := countLogged(recorded, "failed to mark passthrough as running"); n != 0 { + t.Fatalf("passthrough-mark attempts = %d, want 0 after a failed write", n) + } +} + +func TestWriteToUnreadyPtyDropsInputAndAdoptsDiscoveredProcess(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + replacement := startDeferredPassthrough(t, runner, "session-unready") + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + var ptyWriter io.Writer + gotID, gotDirect := handler.writeToUnreadyPty( + []byte("ls\n"), "session-unready", "stale-proc", runner, wsw, &ptyWriter, false, + ) + if gotID != replacement.ID { + t.Fatalf("writeToUnreadyPty() process id = %q, want %q", gotID, replacement.ID) + } + if gotDirect { + t.Fatal("writeToUnreadyPty() directOutputSet = true, want false while the PTY is unavailable") + } + if n := countLogged(recorded, "PTY not ready, dropping input"); n != 1 { + t.Fatalf("dropped-input records = %d, want 1", n) + } +} + +func TestHandleTerminalInputRoutesByPtyReadiness(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + pty := &bytes.Buffer{} + var ready io.Writer = pty + if id, direct := handler.handleTerminalInput( + []byte("a"), "session-1", "proc-1", runner, wsw, &ready, true, + ); id != "proc-1" || !direct { + t.Fatalf("handleTerminalInput(ready) = (%q, %v), want (\"proc-1\", true)", id, direct) + } + if pty.String() != "a" { + t.Fatalf("PTY received %q, want %q", pty.String(), "a") + } + + var unready io.Writer + if id, direct := handler.handleTerminalInput( + []byte("b"), "session-1", "proc-1", runner, wsw, &unready, false, + ); id != "proc-1" || direct { + t.Fatalf("handleTerminalInput(unready) = (%q, %v), want (\"proc-1\", false)", id, direct) + } + if n := countLogged(recorded, "PTY not ready, dropping input"); n != 1 { + t.Fatalf("dropped-input records = %d, want 1", n) + } + if pty.String() != "a" { + t.Fatalf("PTY received %q, want the unready input to be dropped", pty.String()) + } +} + +func TestUserShellInputRoutesByPtyReadiness(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + pty := &bytes.Buffer{} + var ready io.Writer = pty + if direct := handler.handleUserShellInput( + []byte("pwd\n"), "session-1", "env-1", "term-1", "proc-1", runner, wsw, &ready, true, + ); !direct { + t.Fatal("handleUserShellInput(ready) = false, want true") + } + if pty.String() != "pwd\n" { + t.Fatalf("shell PTY received %q, want %q", pty.String(), "pwd\n") + } + + var unready io.Writer + if direct := handler.handleUserShellInput( + []byte("pwd\n"), "session-1", "env-1", "term-1", "proc-1", runner, wsw, &unready, false, + ); direct { + t.Fatal("handleUserShellInput(unready) = true, want false") + } + if n := countLogged(recorded, "PTY not ready, dropping input"); n != 1 { + t.Fatalf("dropped-input records = %d, want 1", n) + } +} + +func TestWriteToReadyUserShellPtyKeepsStateWhenReconnectFails(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + _, server := newTerminalWSPair(t) + wsw := newWsWriter(server) + t.Cleanup(func() { _ = wsw.Close() }) + + pty := &failingWriter{} + var ptyWriter io.Writer = pty + direct := handler.writeToReadyUserShellPty( + []byte("ls\n"), "session-1", "env-1", "term-gone", "proc-gone", runner, wsw, &ptyWriter, false, + ) + if direct { + t.Fatal("writeToReadyUserShellPty() = true, want false when the shell cannot be re-established") + } + if pty.writes != 1 { + t.Fatalf("PTY writes = %d, want exactly 1 (no retry without a new writer)", pty.writes) + } + if n := countLogged(recorded, "PTY write error"); n != 1 { + t.Fatalf("PTY write error records = %d, want 1", n) + } +} + +func TestHandleUserShellResizeRejectsBadPayloads(t *testing.T) { + tests := []struct { + name string + payload string + wantLog string + }{ + {name: "malformed json", payload: `not-json`, wantLog: "failed to parse resize command"}, + {name: "zero dimensions", payload: `{"cols":0,"rows":0}`, wantLog: "invalid resize dimensions"}, + {name: "unknown shell", payload: `{"cols":80,"rows":24}`, wantLog: "failed to resize user shell PTY"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + log, recorded := observedTerminalLogger(t) + handler, _ := newBridgeTestHandler(t, log) + runner := newBridgeTestRunner(t, log) + + handler.handleUserShellResize([]byte(tt.payload), "session-1", "env-1", "term-1", runner) + + if n := countLogged(recorded, tt.wantLog); n != 1 { + t.Fatalf("%q records = %d, want 1", tt.wantLog, n) + } + if n := countLogged(recorded, "user shell PTY resized"); n != 0 { + t.Fatalf("success records = %d, want 0", n) + } + }) + } +} + +func TestOSPIDForInteractiveProcessHandlesMissingInputs(t *testing.T) { + log, _ := observedTerminalLogger(t) + runner := newBridgeTestRunner(t, log) + deferredProc := startDeferredPassthrough(t, runner, "session-pid") + + if got := osPIDForInteractiveProcess(nil, "proc-1"); got != 0 { + t.Fatalf("osPIDForInteractiveProcess(nil runner) = %d, want 0", got) + } + if got := osPIDForInteractiveProcess(runner, ""); got != 0 { + t.Fatalf("osPIDForInteractiveProcess(empty id) = %d, want 0", got) + } + if got := osPIDForInteractiveProcess(runner, "proc-unknown"); got != 0 { + t.Fatalf("osPIDForInteractiveProcess(unknown id) = %d, want 0", got) + } + if got := osPIDForInteractiveProcess(runner, deferredProc.ID); got != 0 { + t.Fatalf("osPIDForInteractiveProcess(deferred process) = %d, want 0", got) + } +} diff --git a/apps/backend/internal/gateway/websocket/terminal_wsutil_test.go b/apps/backend/internal/gateway/websocket/terminal_wsutil_test.go new file mode 100644 index 0000000000..83d413be7f --- /dev/null +++ b/apps/backend/internal/gateway/websocket/terminal_wsutil_test.go @@ -0,0 +1,149 @@ +package websocket + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + gorillaws "github.com/gorilla/websocket" + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" + + "github.com/kandev/kandev/internal/common/logger" +) + +// wsTestTimeout bounds every blocking read/join in the terminal WebSocket +// tests. All of them join on an observable side effect, so the timeout only +// fires when the production code genuinely failed to act. +const wsTestTimeout = 3 * time.Second + +var testWSUpgrader = gorillaws.Upgrader{ + CheckOrigin: func(*http.Request) bool { return true }, +} + +// newTerminalWSPair returns a connected pair of gorilla WebSocket connections: +// the first is the dialing (browser) side, the second the accepting (server) +// side. Both are closed on cleanup, which is what lets the goleak TestMain in +// this package stay green even when a test fails early. +func newTerminalWSPair(t *testing.T) (dialed, accepted *gorillaws.Conn) { + t.Helper() + + accepts := make(chan *gorillaws.Conn, 1) + // Drain a connection the handler delivers after the timeout branch below: + // nothing else would ever close it, and its read goroutine would outlive + // the test. + t.Cleanup(func() { + select { + case late := <-accepts: + _ = late.Close() + default: + } + }) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := testWSUpgrader.Upgrade(w, r, nil) + if err != nil { + return + } + accepts <- conn + })) + t.Cleanup(srv.Close) + + conn, resp, err := gorillaws.DefaultDialer.Dial("ws"+strings.TrimPrefix(srv.URL, "http"), nil) + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + if err != nil { + t.Fatalf("dial websocket: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + + acceptTimer := time.NewTimer(wsTestTimeout) + defer acceptTimer.Stop() + select { + case server := <-accepts: + t.Cleanup(func() { _ = server.Close() }) + return conn, server + case <-acceptTimer.C: + t.Fatal("timed out waiting for the server side of the websocket pair") + return nil, nil + } +} + +// observedTerminalLogger returns a logger whose records are captured in memory. +// Several terminal-pump branches (dropped input, resize fallbacks) have no +// other in-process observable, so the emitted record is the assertion target. +func observedTerminalLogger(t *testing.T) (*logger.Logger, *observer.ObservedLogs) { + t.Helper() + core, recorded := observer.New(zap.DebugLevel) + log, err := logger.NewFromZap(zap.New(core)) + if err != nil { + t.Fatalf("create observed logger: %v", err) + } + return log, recorded +} + +// readBinary reads one message with a bounded deadline and asserts it is a +// binary frame. +func readBinary(t *testing.T, conn *gorillaws.Conn) []byte { + t.Helper() + if err := conn.SetReadDeadline(time.Now().Add(wsTestTimeout)); err != nil { + t.Fatalf("set read deadline: %v", err) + } + msgType, data, err := conn.ReadMessage() + if err != nil { + t.Fatalf("read message: %v", err) + } + if msgType != gorillaws.BinaryMessage { + t.Fatalf("message type = %d, want %d (binary)", msgType, gorillaws.BinaryMessage) + } + return data +} + +// spawnJoinable starts fn in a goroutine and registers, immediately, a cleanup +// that unblocks it and waits for it to return. +// +// Registering the join up front rather than only on the happy path is what +// keeps an early t.Fatal from leaving a pump running into later tests, where +// goleak would report it against the wrong test. The cleanup reports +// non-fatally so the remaining cleanups (which close the connections) still run. +func spawnJoinable(t *testing.T, what string, unblock func(), fn func()) <-chan struct{} { + t.Helper() + done := make(chan struct{}) + go func() { + defer close(done) + fn() + }() + t.Cleanup(func() { + if unblock != nil { + unblock() + } + timer := time.NewTimer(wsTestTimeout) + defer timer.Stop() + select { + case <-done: + case <-timer.C: + t.Errorf("%s did not return within %s", what, wsTestTimeout) + } + }) + return done +} + +// joinWithin waits for done to close, failing the test when the production +// goroutine under test did not terminate. +func joinWithin(t *testing.T, done <-chan struct{}, what string) { + t.Helper() + timer := time.NewTimer(wsTestTimeout) + defer timer.Stop() + select { + case <-done: + case <-timer.C: + t.Fatalf("%s did not return within %s", what, wsTestTimeout) + } +} + +// countLogged returns how many captured records carry the given message. +func countLogged(recorded *observer.ObservedLogs, message string) int { + return len(recorded.FilterMessage(message).All()) +} diff --git a/apps/backend/internal/gateway/websocket/vscode_proxy_handler_test.go b/apps/backend/internal/gateway/websocket/vscode_proxy_handler_test.go new file mode 100644 index 0000000000..b31113d80c --- /dev/null +++ b/apps/backend/internal/gateway/websocket/vscode_proxy_handler_test.go @@ -0,0 +1,303 @@ +package websocket + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/kandev/kandev/internal/agent/runtime/lifecycle" +) + +const vscodeStatusPath = "/api/v1/vscode/status" + +// newFakeCodeServer serves the agentctl VS Code status endpoint plus a stub +// code-server behind the proxy path. +func newFakeCodeServer(t *testing.T, status string) *fakeAgentctl { + t.Helper() + return newFakeAgentctl(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == vscodeStatusPath { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"status":"` + status + `","port":8443,"url":"http://127.0.0.1:8443"}`)) + return + } + w.Header().Set("Content-Type", "text/plain") + _, _ = w.Write([]byte("code-server:" + r.URL.Path)) + }) +} + +func newVscodeEngine(handler *VscodeProxyHandler) *gin.Engine { + gin.SetMode(gin.TestMode) + engine := gin.New() + engine.Any("/vscode/:sessionId/*path", handler.HandleVscodeProxy) + return engine +} + +// countStatusCalls reports how many times the upstream VS Code status endpoint +// was polled, which is how the per-session proxy cache is observed. +func countStatusCalls(seen []recordedRequest) int { + n := 0 + for _, req := range seen { + if req.path == vscodeStatusPath { + n++ + } + } + return n +} + +func TestHandleVscodeProxyRewritesPathAndInjectsAgentctlToken(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeCodeServer(t, "running") + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-vs", upstream.server.URL, "vscode-token", log) + gateway := serveGateway(t, newVscodeEngine(NewVscodeProxyHandler(manager, log))) + + got := gateway(http.MethodGet, "/vscode/sess-vs/static/out/main.js") + if got.status != http.StatusOK { + t.Fatalf("status = %d, want %d", got.status, http.StatusOK) + } + if got.body != "code-server:/api/v1/vscode/proxy/static/out/main.js" { + t.Fatalf("body = %q, want the proxied code-server body", got.body) + } + + // One snapshot: seen() copies, so indexing a second call would address a + // different slice than the one being ranged over. + seen := upstream.seen() + var proxied *recordedRequest + for i := range seen { + if strings.HasPrefix(seen[i].path, "/api/v1/vscode/proxy") { + proxied = &seen[i] + } + } + if proxied == nil { + t.Fatalf("no proxied request reached the upstream: %+v", seen) + } + if proxied.path != "/api/v1/vscode/proxy/static/out/main.js" { + t.Fatalf("upstream path = %q", proxied.path) + } + if proxied.auth != "Bearer vscode-token" { + t.Fatalf("upstream Authorization = %q, want %q", proxied.auth, "Bearer vscode-token") + } +} + +// Older agentctl instances only register /api/v1/vscode/proxy/*path, so both +// root forms — with and without a trailing slash — must reach +// /api/v1/vscode/proxy/ upstream. Only the no-slash form exercises the +// empty-remainder branch; "/vscode/:id/" would survive dropping it. +func TestHandleVscodeProxyKeepsTrailingSlashForRootRequests(t *testing.T) { + for _, requestPath := range []string{"/vscode/sess-root/", "/vscode/sess-root"} { + t.Run(requestPath, func(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeCodeServer(t, "running") + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-root", upstream.server.URL, "", log) + handler := NewVscodeProxyHandler(manager, log) + + // The gin wildcard route would redirect the no-slash form, so both + // variants are driven through the handler with the params it reads. + gateway := serveGateway(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c, _ := gin.CreateTestContext(w) + c.Request = r + c.Params = gin.Params{{Key: "sessionId", Value: "sess-root"}} + handler.HandleVscodeProxy(c) + })) + + got := gateway(http.MethodGet, requestPath) + if got.status != http.StatusOK { + t.Fatalf("status = %d, want %d", got.status, http.StatusOK) + } + if got.body != "code-server:/api/v1/vscode/proxy/" { + t.Fatalf("body = %q, want the root code-server path", got.body) + } + }) + } +} + +func TestHandleVscodeProxyRejectsStoppedCodeServer(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeCodeServer(t, "stopped") + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-stopped", upstream.server.URL, "", log) + handler := NewVscodeProxyHandler(manager, log) + gateway := serveGateway(t, newVscodeEngine(handler)) + + got := gateway(http.MethodGet, "/vscode/sess-stopped/") + if got.status != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", got.status, http.StatusServiceUnavailable) + } + if msg := decodeErrorBody(t, []byte(got.body)); msg != "code-server is not running" { + t.Fatalf("error = %q", msg) + } + handler.mu.Lock() + defer handler.mu.Unlock() + if len(handler.proxies) != 0 { + t.Fatalf("cached proxies = %d, want 0 while code-server is stopped", len(handler.proxies)) + } +} + +func TestHandleVscodeProxyReportsUnreachableStatusEndpoint(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeAgentctl(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + }) + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-broken", upstream.server.URL, "", log) + gateway := serveGateway(t, newVscodeEngine(NewVscodeProxyHandler(manager, log))) + + got := gateway(http.MethodGet, "/vscode/sess-broken/") + if got.status != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", got.status, http.StatusServiceUnavailable) + } + if msg := decodeErrorBody(t, []byte(got.body)); msg != "failed to get vscode status" { + t.Fatalf("error = %q", msg) + } +} + +func TestHandleVscodeProxyRequiresSessionID(t *testing.T) { + log, _ := observedTerminalLogger(t) + handler := NewVscodeProxyHandler(newProxyTestManager(t, log), log) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/vscode//", nil) + c.Params = gin.Params{{Key: "sessionId", Value: ""}} + handler.HandleVscodeProxy(c) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want %d", recorder.Code, http.StatusBadRequest) + } + if got := decodeErrorBody(t, recorder.Body.Bytes()); got != "sessionId is required" { + t.Fatalf("error = %q", got) + } +} + +func TestHandleVscodeProxyDeniesForeignSessionBeforeStatusLookup(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeCodeServer(t, "running") + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-foreign", upstream.server.URL, "", log) + manager.SetSessionAccessChecker(func(context.Context, string) error { + return errors.New("not your session") + }) + gateway := serveGateway(t, newVscodeEngine(NewVscodeProxyHandler(manager, log))) + + got := gateway(http.MethodGet, "/vscode/sess-foreign/") + if got.status != http.StatusNotFound { + t.Fatalf("status = %d, want %d", got.status, http.StatusNotFound) + } + if msg := decodeErrorBody(t, []byte(got.body)); msg != "session not found" { + t.Fatalf("error = %q", msg) + } + if n := len(upstream.seen()); n != 0 { + t.Fatalf("upstream requests = %d, want 0 for a denied session", n) + } +} + +func TestHandleVscodeProxyReportsMissingExecutionAndClient(t *testing.T) { + log, _ := observedTerminalLogger(t) + manager := newProxyTestManager(t, log) + clientless := &lifecycle.AgentExecution{ID: "exec-vs-noclient", SessionID: "sess-vs-noclient"} + if err := manager.ExecutionStoreForTesting().Add(clientless); err != nil { + t.Fatalf("add execution: %v", err) + } + gateway := serveGateway(t, newVscodeEngine(NewVscodeProxyHandler(manager, log))) + + tests := []struct { + name string + sessionID string + wantStatus int + wantError string + }{ + { + name: "no execution", + sessionID: "sess-vs-missing", + wantStatus: http.StatusNotFound, + wantError: "session not found or no active execution", + }, + { + name: "no agentctl client", + sessionID: "sess-vs-noclient", + wantStatus: http.StatusServiceUnavailable, + wantError: "agentctl client not available", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := gateway(http.MethodGet, "/vscode/"+tt.sessionID+"/") + if got.status != tt.wantStatus { + t.Fatalf("status = %d, want %d", got.status, tt.wantStatus) + } + if msg := decodeErrorBody(t, []byte(got.body)); msg != tt.wantError { + t.Fatalf("error = %q, want %q", msg, tt.wantError) + } + }) + } +} + +// The cached proxy is the whole point of the fast path: a second request must +// not re-poll the upstream VS Code status endpoint. +func TestHandleVscodeProxyChecksUpstreamStatusOnlyOnce(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeCodeServer(t, "running") + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-cached", upstream.server.URL, "", log) + handler := NewVscodeProxyHandler(manager, log) + gateway := serveGateway(t, newVscodeEngine(handler)) + + for _, path := range []string{"/vscode/sess-cached/", "/vscode/sess-cached/manifest.json"} { + if got := gateway(http.MethodGet, path); got.status != http.StatusOK { + t.Fatalf("status for %s = %d, want %d", path, got.status, http.StatusOK) + } + } + if n := countStatusCalls(upstream.seen()); n != 1 { + t.Fatalf("vscode status polls = %d, want 1 (second request must reuse the cache)", n) + } + + handler.mu.Lock() + defer handler.mu.Unlock() + if len(handler.proxies) != 1 { + t.Fatalf("cached proxies = %d, want 1", len(handler.proxies)) + } +} + +func TestHandleVscodeProxyInvalidatesCacheOnUpstreamFailure(t *testing.T) { + log, _ := observedTerminalLogger(t) + upstream := newFakeCodeServer(t, "running") + + manager := newProxyTestManager(t, log) + addExecutionForURL(t, manager, "sess-flaky", upstream.server.URL, "", log) + handler := NewVscodeProxyHandler(manager, log) + gateway := serveGateway(t, newVscodeEngine(handler)) + + if got := gateway(http.MethodGet, "/vscode/sess-flaky/"); got.status != http.StatusOK { + t.Fatalf("priming status = %d, want %d", got.status, http.StatusOK) + } + handler.mu.Lock() + cached := len(handler.proxies) + handler.mu.Unlock() + if cached != 1 { + t.Fatalf("cached proxies = %d after priming, want 1", cached) + } + + upstream.server.Close() + + got := gateway(http.MethodGet, "/vscode/sess-flaky/") + if got.status != http.StatusBadGateway { + t.Fatalf("status = %d, want %d after the upstream died", got.status, http.StatusBadGateway) + } + handler.mu.Lock() + defer handler.mu.Unlock() + if len(handler.proxies) != 0 { + t.Fatalf("cached proxies = %d, want 0 after an upstream failure", len(handler.proxies)) + } +}