diff --git a/packages/sandbox/daemon-go/internal/dispatch/dispatch.go b/packages/sandbox/daemon-go/internal/dispatch/dispatch.go index 46eae015fc..50ae593c53 100644 --- a/packages/sandbox/daemon-go/internal/dispatch/dispatch.go +++ b/packages/sandbox/daemon-go/internal/dispatch/dispatch.go @@ -3,6 +3,7 @@ package dispatch import ( "context" "encoding/json" + "errors" "io" "log/slog" "net/http" @@ -18,6 +19,12 @@ import ( const tombstoneTTL = 60 * time.Second +// maxDispatchBodyBytes bounds the inline `/dispatch` request body. Matches +// maxOffloadBytes: a run whose messages are actually this large is expected to +// come in via messagesRef, not inline, so this cap just stops an unbounded +// read from parking the pod's memory on one request. +const maxDispatchBodyBytes = maxOffloadBytes + // Terminal code for a run this pod could not finish (shutdown / dropped // connection) as opposed to one that was cancelled on purpose. Studio maps it // to `SandboxUnreachableError` and continues the turn elsewhere — the literal @@ -282,8 +289,13 @@ func (reg *Registry) HandleDispatch(w http.ResponseWriter, r *http.Request, deps jsonError(w, 401, map[string]string{"error": "unauthorized"}) return } - body, err := io.ReadAll(r.Body) + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxDispatchBodyBytes)) if err != nil { + var tooLarge *http.MaxBytesError + if errors.As(err, &tooLarge) { + jsonError(w, 413, map[string]string{"error": "body_too_large"}) + return + } jsonError(w, 400, map[string]string{"error": "bad_json"}) return } diff --git a/packages/sandbox/daemon-go/internal/dispatch/dispatch_test.go b/packages/sandbox/daemon-go/internal/dispatch/dispatch_test.go index a8b162be6a..93dc621533 100644 --- a/packages/sandbox/daemon-go/internal/dispatch/dispatch_test.go +++ b/packages/sandbox/daemon-go/internal/dispatch/dispatch_test.go @@ -354,3 +354,18 @@ func TestUnauthorizedCancelLeavesNoTombstone(t *testing.T) { t.Fatal("a rejected cancel must not tombstone the run") } } + +// An oversized inline dispatch body must be rejected before it is buffered in +// full — the offload path already caps at this size, and the inline path had +// no cap at all. +func TestDispatchRejectsOversizedBody(t *testing.T) { + const token = "tkn" + body := strings.NewReader(strings.Repeat("a", maxDispatchBodyBytes+1)) + req := httptest.NewRequest("POST", "/dispatch", body) + req.Header.Set("Authorization", "Bearer "+token) + rec := httptest.NewRecorder() + NewRegistry().HandleDispatch(rec, req, Deps{DaemonToken: func() string { return token }}) + if rec.Code != 413 { + t.Fatalf("dispatch with oversized body returned %d, want 413", rec.Code) + } +}