Add HTTP size limits
This commit is contained in:
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user