Add HTTP size limits

This commit is contained in:
2026-07-05 00:17:07 +00:00
parent a16f66cbc7
commit f7d821067f
15 changed files with 609 additions and 46 deletions

View File

@@ -243,6 +243,24 @@ func TestHandlerFileRefsUnderArtifactRootWork(t *testing.T) {
}
}
func TestHandlerFileRefsAboveArtifactLimitAreRejected(t *testing.T) {
root := t.TempDir()
if err := os.WriteFile(filepath.Join(root, "large.txt"), []byte("123456"), 0o644); err != nil {
t.Fatal(err)
}
h := newArtifactRootHandlerWithLimit(t, root, 5)
req := httptest.NewRequest(http.MethodPost, "/v1/runs", bytes.NewBufferString(`{
"prompt_id":"p",
"inputs":{"x":{"type":"file","uri":"large.txt"}}
}`))
w := httptest.NewRecorder()
h.ServeHTTP(w, req)
assertHTTPErrorCode(t, w, http.StatusRequestEntityTooLarge, "artifact_too_large")
}
func TestHandlerFileRefsOutsideArtifactRootAreRejected(t *testing.T) {
root := t.TempDir()
outside := t.TempDir()
@@ -576,6 +594,69 @@ func TestHandlerInvalidJSON(t *testing.T) {
}
}
func TestHandlerRejectsTrailingJSON(t *testing.T) {
h := NewHandler(&fakeRunner{})
req := httptest.NewRequest(http.MethodPost, "/v1/runs", bytes.NewBufferString(`{"prompt_id":"p","inputs":{"x":{"type":"file","uri":"a"}}} {}`))
w := httptest.NewRecorder()
h.ServeHTTP(w, req)
assertHTTPErrorCode(t, w, http.StatusBadRequest, "invalid_json")
}
func TestHandlerRequestTooLarge(t *testing.T) {
h := NewHandlerWithOptions(&fakeRunner{}, HandlerOptions{MaxRequestBytes: 12})
req := httptest.NewRequest(http.MethodPost, "/v1/runs", bytes.NewBufferString(`{"prompt_id":"p","inputs":{"x":{"type":"file","uri":"a"}}}`))
w := httptest.NewRecorder()
h.ServeHTTP(w, req)
assertHTTPErrorCode(t, w, http.StatusRequestEntityTooLarge, "request_too_large")
}
func TestHandlerMalformedJSONBelowLimitStillBadRequest(t *testing.T) {
h := NewHandlerWithOptions(&fakeRunner{}, HandlerOptions{MaxRequestBytes: 1024})
req := httptest.NewRequest(http.MethodPost, "/v1/runs", bytes.NewBufferString("{"))
w := httptest.NewRecorder()
h.ServeHTTP(w, req)
assertHTTPErrorCode(t, w, http.StatusBadRequest, "invalid_json")
}
func TestHandlerResponseTooLarge(t *testing.T) {
h := NewHandlerWithOptions(&fakeRunner{result: &domain.RunResult{
Artifact: domain.Artifact{Body: []byte(strings.Repeat("x", 128))},
Validation: domain.ValidationResult{Status: domain.ValidationPassed, Mode: domain.ValidationBasic, IsValid: true},
EffectiveModelParams: domain.ExecutionTarget{Endpoint: "http://llm/v1", Model: "m1"},
}}, HandlerOptions{MaxRequestBytes: 1024, MaxResponseBytes: 64})
req := httptest.NewRequest(http.MethodPost, "/v1/runs", bytes.NewBufferString(`{"prompt_id":"p","inputs":{"x":{"type":"file","uri":"a"}}}`))
w := httptest.NewRecorder()
h.ServeHTTP(w, req)
assertHTTPErrorCode(t, w, http.StatusRequestEntityTooLarge, "response_too_large")
}
func TestHandlerRawOutputDoesNotBypassResponseLimit(t *testing.T) {
h := NewHandlerWithOptions(&fakeRunner{result: &domain.RunResult{
Artifact: domain.Artifact{Body: []byte("ok")},
RawOutput: strings.Repeat("raw", 80),
Validation: domain.ValidationResult{Status: domain.ValidationPassed, Mode: domain.ValidationBasic, IsValid: true},
EffectiveModelParams: domain.ExecutionTarget{Endpoint: "http://llm/v1", Model: "m1"},
}}, HandlerOptions{MaxRequestBytes: 1024, MaxResponseBytes: 128})
req := httptest.NewRequest(http.MethodPost, "/v1/runs", bytes.NewBufferString(`{
"prompt_id":"p",
"inputs":{"x":{"type":"file","uri":"a"}},
"include_raw_output":true
}`))
w := httptest.NewRecorder()
h.ServeHTTP(w, req)
assertHTTPErrorCode(t, w, http.StatusRequestEntityTooLarge, "response_too_large")
}
func TestHandlerMissingPromptID(t *testing.T) {
h := NewHandler(&fakeRunner{})
req := httptest.NewRequest(http.MethodPost, "/v1/runs", bytes.NewBufferString(`{"inputs":{"x":{"type":"file","uri":"a"}}}`))
@@ -770,7 +851,13 @@ func wrap(stage error, cause error) error {
func newArtifactRootHandler(t *testing.T, root string) *Handler {
t.Helper()
reader, err := artifact.NewRestrictedCompositeReader(root)
return newArtifactRootHandlerWithLimit(t, root, 0)
}
func newArtifactRootHandlerWithLimit(t *testing.T, root string, maxArtifactBytes int64) *Handler {
t.Helper()
reader, err := artifact.NewRestrictedCompositeReaderWithLimit(root, maxArtifactBytes)
if err != nil {
t.Fatalf("expected restricted artifact reader: %v", err)
}