Add HTTP size limits
This commit is contained in:
@@ -4,10 +4,12 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"gitea.maximumdirect.net/eric/scriptorium/internal/artifact"
|
||||
"gitea.maximumdirect.net/eric/scriptorium/internal/defaults"
|
||||
"gitea.maximumdirect.net/eric/scriptorium/internal/domain"
|
||||
"gitea.maximumdirect.net/eric/scriptorium/internal/profile"
|
||||
"gitea.maximumdirect.net/eric/scriptorium/internal/promptdef"
|
||||
@@ -19,11 +21,24 @@ type Runner interface {
|
||||
}
|
||||
|
||||
type Handler struct {
|
||||
runner Runner
|
||||
runner Runner
|
||||
options HandlerOptions
|
||||
}
|
||||
|
||||
type HandlerOptions struct {
|
||||
MaxRequestBytes int64
|
||||
MaxResponseBytes int64
|
||||
}
|
||||
|
||||
func NewHandler(runner Runner) *Handler {
|
||||
return &Handler{runner: runner}
|
||||
return NewHandlerWithOptions(runner, HandlerOptions{
|
||||
MaxRequestBytes: defaults.HTTPMaxRequestBytesDefault,
|
||||
MaxResponseBytes: defaults.HTTPMaxResponseBytesDefault,
|
||||
})
|
||||
}
|
||||
|
||||
func NewHandlerWithOptions(runner Runner, options HandlerOptions) *Handler {
|
||||
return &Handler{runner: runner, options: options}
|
||||
}
|
||||
|
||||
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -37,9 +52,26 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
|
||||
var req runRequestDTO
|
||||
dec := json.NewDecoder(r.Body)
|
||||
body := r.Body
|
||||
if h.options.MaxRequestBytes > 0 {
|
||||
body = http.MaxBytesReader(w, r.Body, h.options.MaxRequestBytes)
|
||||
}
|
||||
dec := json.NewDecoder(body)
|
||||
dec.DisallowUnknownFields()
|
||||
if err := dec.Decode(&req); err != nil {
|
||||
if isRequestTooLarge(err) {
|
||||
writeError(w, http.StatusRequestEntityTooLarge, "request_too_large", "request body is too large")
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusBadRequest, "invalid_json", "invalid JSON request body")
|
||||
return
|
||||
}
|
||||
var trailing any
|
||||
if err := dec.Decode(&trailing); err != io.EOF {
|
||||
if isRequestTooLarge(err) {
|
||||
writeError(w, http.StatusRequestEntityTooLarge, "request_too_large", "request body is too large")
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusBadRequest, "invalid_json", "invalid JSON request body")
|
||||
return
|
||||
}
|
||||
@@ -121,7 +153,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
raw := res.RawOutput
|
||||
resp.RawModelOutput = &raw
|
||||
}
|
||||
writeJSON(w, http.StatusOK, resp)
|
||||
writeLimitedJSON(w, http.StatusOK, resp, h.options.MaxResponseBytes)
|
||||
}
|
||||
|
||||
func executionTargetOverrideFromModelOverrideDTO(dto *modelOverrideRequestDTO) *domain.ExecutionTargetOverride {
|
||||
@@ -190,6 +222,8 @@ func mapRunError(err error) (int, string, string) {
|
||||
return http.StatusBadRequest, "profile_load_failed", "failed to load execution profile"
|
||||
case errors.Is(err, artifact.ErrFileNotAllowed), errors.Is(err, artifact.ErrFileOutsideRoot):
|
||||
return http.StatusBadRequest, "artifact_not_allowed", "file input artifact is not allowed"
|
||||
case errors.Is(err, artifact.ErrFileTooLarge):
|
||||
return http.StatusRequestEntityTooLarge, "artifact_too_large", "file input artifact is too large"
|
||||
case errors.Is(err, usecase.ErrArtifactLoad):
|
||||
return http.StatusBadRequest, "artifact_read_failed", "failed to read input artifact"
|
||||
case errors.Is(err, usecase.ErrPromptRender):
|
||||
@@ -204,9 +238,23 @@ func mapRunError(err error) (int, string, string) {
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, v any) {
|
||||
writeLimitedJSON(w, status, v, 0)
|
||||
}
|
||||
|
||||
func writeLimitedJSON(w http.ResponseWriter, status int, v any, maxBytes int64) {
|
||||
data, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "internal_error", "internal server error")
|
||||
return
|
||||
}
|
||||
data = append(data, '\n')
|
||||
if maxBytes > 0 && int64(len(data)) > maxBytes {
|
||||
writeError(w, http.StatusRequestEntityTooLarge, "response_too_large", "response body is too large")
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(v)
|
||||
_, _ = w.Write(data)
|
||||
}
|
||||
|
||||
func writeError(w http.ResponseWriter, status int, code, message string) {
|
||||
@@ -217,3 +265,8 @@ func writeError(w http.ResponseWriter, status int, code, message string) {
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func isRequestTooLarge(err error) bool {
|
||||
var maxBytesErr *http.MaxBytesError
|
||||
return errors.As(err, &maxBytesErr)
|
||||
}
|
||||
|
||||
@@ -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