diff --git a/docs/ai-agent/agent-runner-service.md b/docs/ai-agent/agent-runner-service.md index eacccb202..158523f62 100644 --- a/docs/ai-agent/agent-runner-service.md +++ b/docs/ai-agent/agent-runner-service.md @@ -171,6 +171,8 @@ There is no REST equivalent to `services/ai-agent`'s `/conversations/:id/pause` The three internal endpoints require `X-Internal-Token` matching `INTERNAL_API_KEY` — `services/api`'s `AI_AGENT_INTERNAL_KEY` must equal this value. +`/agent-bridge/ws` authenticates after the upgrade, with the `hello` frame (agent id plus bridge token). Because that handshake is unauthenticated until the frame arrives, it is bounded: at most `maxPendingHellos` (128) handshakes in flight per agent-runner process — excess upgrades get `503` before the socket is accepted — the `hello` frame may not exceed 1 KiB, and a peer that sends no `hello` within 10 s is closed. The 10 MiB bridge frame limit applies only after the token checks out. Anyone exempting this path from an edge/browser authentication layer inherits that bound and must not remove it. + --- ## Environment Variables diff --git a/services/agent-runner/internal/acpbridge/server.go b/services/agent-runner/internal/acpbridge/server.go index c953e9771..a81350667 100644 --- a/services/agent-runner/internal/acpbridge/server.go +++ b/services/agent-runner/internal/acpbridge/server.go @@ -27,6 +27,19 @@ import ( // _HELLO_TIMEOUT_SECONDS. const helloTimeout = 10 * time.Second +// helloReadLimit bounds the pre-auth frame. A hello carries only an agent id +// and a token, so a peer that has not authenticated yet must not be able to +// allocate the full bridge budget. The limit is raised once the token checks +// out, below in handleWS. +const helloReadLimit = 1 << 10 + +// maxPendingHellos caps concurrent unauthenticated handshakes. Excess +// upgrades are refused up front instead of occupying the 10s hello window. +const maxPendingHellos = 128 + +// helloSlots is the process-wide budget of unauthenticated handshakes. +var helloSlots = make(chan struct{}, maxPendingHellos) + // bridgeMessageReadLimit overrides coder/websocket's 32KiB default (see // handleWS) — generous enough for a large tool-call payload (verbose // command output, a sizeable file diff) without leaving this per-connection @@ -145,19 +158,28 @@ type helloFrame struct { // self-describing handshake that degrades gracefully across daemon/server // version skew, which headers rejected pre-accept wouldn't). func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) { + select { + case helloSlots <- struct{}{}: + default: + // Refuse before upgrading: a saturated service must not accept more + // sockets it cannot authenticate or bound. + http.Error(w, "too many pending bridge handshakes", http.StatusServiceUnavailable) + return + } + held := true + release := func() { + if held { + <-helloSlots + held = false + } + } + defer release() + conn, err := websocket.Accept(w, r, nil) if err != nil { return } - // coder/websocket defaults to a 32KiB read limit — comfortably too - // small for a bridge daemon relaying a raw ACP event verbatim (a large - // command's output, or a sizeable file diff, both routinely exceed it), - // which surfaced as the daemon getting disconnected mid-conversation - // with "received close frame: status = StatusMessageTooBig". Bridge - // events are the only traffic this endpoint ever reads at any real - // size (hello/ping frames are tiny), so this is scoped to the whole - // connection rather than per-message. - conn.SetReadLimit(bridgeMessageReadLimit) + conn.SetReadLimit(helloReadLimit) helloCtx, cancel := context.WithTimeout(r.Context(), helloTimeout) var hello helloFrame @@ -186,6 +208,13 @@ func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) { return } + // Authenticated: free the handshake budget and allow bridge-sized frames. + release() + conn.SetReadLimit(bridgeMessageReadLimit) + + // TODO: cover this raise with a post-auth frame above helloReadLimit + // (the e2e fake bridge currently sends only a small event payload). + sessionID, err := s.Registry.Register(r.Context(), agentID, cfg.ProjectID, wsConn{conn}) if err != nil { s.Log.Error("acpbridge: failed to register bridge connection", "agent_id", agentID, "error", err) diff --git a/services/agent-runner/internal/acpbridge/server_hello_test.go b/services/agent-runner/internal/acpbridge/server_hello_test.go new file mode 100644 index 000000000..c50e90ccf --- /dev/null +++ b/services/agent-runner/internal/acpbridge/server_hello_test.go @@ -0,0 +1,102 @@ +package acpbridge + +import ( + "bytes" + "context" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/coder/websocket" +) + +func newHelloTestServer(t *testing.T) *httptest.Server { + t.Helper() + s := &Server{Log: slog.New(slog.NewTextHandler(io.Discard, nil))} + ts := httptest.NewServer(s.Routes()) + t.Cleanup(ts.Close) + return ts +} + +func dialBridge(t *testing.T, ts *httptest.Server) *websocket.Conn { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + conn, dialResp, err := websocket.Dial(ctx, "ws"+strings.TrimPrefix(ts.URL, "http")+"/agent-bridge/ws", nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + if dialResp != nil && dialResp.Body != nil { + _ = dialResp.Body.Close() + } + t.Cleanup(func() { _ = conn.CloseNow() }) + return conn +} + +// A peer that has not authenticated must not be able to write bridge-sized +// frames: the pre-auth limit is the only thing standing between an anonymous +// socket and the relay's 10 MiB buffer. +func TestHandleWSRejectsOversizedPreAuthFrame(t *testing.T) { + ts := newHelloTestServer(t) + conn := dialBridge(t, ts) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = conn.Write(ctx, websocket.MessageText, bytes.Repeat([]byte("a"), 2*helloReadLimit)) + _, _, err := conn.Read(ctx) + // 1009 proves the frame was refused by the pre-auth limit; a malformed + // hello under the bridge-sized limit closes with 1007 instead, because + // wsjson.Read rejects the JSON before the handler's 4400 close runs. + if status := websocket.CloseStatus(err); status != websocket.StatusMessageTooBig { + t.Fatalf("close status = %d, want %d (err=%v)", status, websocket.StatusMessageTooBig, err) + } +} + +func TestHandleWSRefusesWhenHandshakeBudgetIsExhausted(t *testing.T) { + ts := newHelloTestServer(t) + for i := 0; i < maxPendingHellos; i++ { + helloSlots <- struct{}{} + } + t.Cleanup(func() { + for i := 0; i < maxPendingHellos; i++ { + <-helloSlots + } + }) + + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, ts.URL+"/agent-bridge/ws", nil) + if err != nil { + t.Fatalf("request: %v", err) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable) + } +} + +func TestHandleWSReleasesHandshakeSlotAfterRejectedHello(t *testing.T) { + ts := newHelloTestServer(t) + conn := dialBridge(t, ts) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = conn.Write(ctx, websocket.MessageText, []byte(`{"type":"hello","agent_id":"not-a-uuid","token":"x"}`)) + if _, _, err := conn.Read(ctx); err == nil { + t.Fatal("expected the connection to be closed for an invalid agent_id") + } + + deadline := time.Now().Add(2 * time.Second) + for len(helloSlots) != 0 { + if time.Now().After(deadline) { + t.Fatalf("handshake slot not released, %d still held", len(helloSlots)) + } + time.Sleep(10 * time.Millisecond) + } +}