test(gateway): cover websocket terminal pumps, proxies and tunnels (#2587)
Co-authored-by: Kandev Agent <agent@kandev.dev> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
João Salavisa
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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(`<link href="/static/app.css"><script src="/static/app.js"></script>`))
|
||||
})
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user