From 4a752ffa2af87a75e9e9ab0a11407b3210445377 Mon Sep 17 00:00:00 2001 From: Steve Dudenhoeffer Date: Mon, 13 Jul 2026 00:50:21 -0400 Subject: [PATCH] fix: harden upstream path + binary-body validation (gadfly round 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - upstreamPath rejects '..' in model ids AND in the rest path — the rest can embed SERVER-SUPPLIED components (ACE-Step result file URLs), so dot-dot/scheme smuggling toward other proxy endpoints is refused - singleImageResult requires positive image evidence (sniffed magic OR declared image/*): an empty-Content-Type error page can no longer pass as 'the image' via sniffImageMIME's PNG-default labelling - upscale/background responses get a dedicated 256MB cap (the 64MB cap is JSON-sized; a 4x PNG legitimately exceeds it) - mesh JSON-detection widened (512-byte whitespace-tolerant peek + reject declared application/json) - Transcribe now reuses buildMultipart; transcriptionFilename takes (filename, mime) so diarize shares it without a fake request struct; truncateForError stops shadowing builtin cap; OnlyMask doc de-ambiguated Co-Authored-By: Claude Fable 5 --- imagegen/background.go | 6 ++--- provider/llamaswap/audio.go | 34 ++++++++++------------------ provider/llamaswap/diarize.go | 11 +-------- provider/llamaswap/mediautil.go | 30 +++++++++++++++++------- provider/llamaswap/mediautil_test.go | 29 ++++++++++++++++++++++++ provider/llamaswap/mesh.go | 20 +++++++++++----- provider/llamaswap/upstream.go | 9 +++++++- 7 files changed, 89 insertions(+), 50 deletions(-) diff --git a/imagegen/background.go b/imagegen/background.go index 93c52a4..2b7cb68 100644 --- a/imagegen/background.go +++ b/imagegen/background.go @@ -15,9 +15,9 @@ type BackgroundRemovalRequest struct { Net string // OnlyMask returns the black/white segmentation mask instead of the - // cutout — white where the subject is. Useful as an inpainting mask - // (EditRequest.Mask semantics after inversion: the mask marks the - // FOREGROUND, an inpaint mask marks the region to REPAINT). + // cutout — WHITE marks the SUBJECT. To inpaint (repaint) the subject, + // use it directly as EditRequest.Mask; to repaint the BACKGROUND, + // invert it first (EditRequest.Mask is white-means-repaint). OnlyMask bool } diff --git a/provider/llamaswap/audio.go b/provider/llamaswap/audio.go index d4e168d..024ad17 100644 --- a/provider/llamaswap/audio.go +++ b/provider/llamaswap/audio.go @@ -7,7 +7,6 @@ import ( "fmt" "io" "mime" - "mime/multipart" "net/http" "net/url" "strings" @@ -119,28 +118,19 @@ func (m *transcriptionModel) Transcribe(ctx context.Context, req audio.Transcrip return nil, fmt.Errorf("%w: transcription requires audio bytes", llm.ErrUnsupported) } - var buf bytes.Buffer - w := multipart.NewWriter(&buf) - fw, err := w.CreateFormFile("file", transcriptionFilename(req)) + buf, formType, err := buildMultipart("build transcription form", + filePart{field: "file", filename: transcriptionFilename(req.Filename, req.MIME), data: req.Audio}, + []formField{ + {"model", m.id, true}, + {"response_format", "json", true}, + {"language", req.Language, false}, + {"prompt", req.Prompt, false}, + }) if err != nil { - return nil, fmt.Errorf("llama-swap: build transcription form: %w", err) - } - if _, err := fw.Write(req.Audio); err != nil { - return nil, fmt.Errorf("llama-swap: build transcription form: %w", err) - } - if err := writeFormFields(w, "build transcription form", []formField{ - {"model", m.id, true}, - {"response_format", "json", true}, - {"language", req.Language, false}, - {"prompt", req.Prompt, false}, - }); err != nil { return nil, err } - if err := w.Close(); err != nil { - return nil, fmt.Errorf("llama-swap: build transcription form: %w", err) - } - raw, _, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/transcriptions", m.id, w.FormDataContentType(), &buf, maxResponseBytes) + raw, _, err := m.p.doRaw(ctx, http.MethodPost, "/v1/audio/transcriptions", m.id, formType, buf, maxResponseBytes) if err != nil { return nil, err } @@ -158,11 +148,11 @@ func (m *transcriptionModel) Transcribe(ctx context.Context, req audio.Transcrip // headers), else one derived from the MIME subtype ("audio.mp3"), else // "audio". MIME parameters ("audio/ogg; codecs=opus") are stripped before // matching. -func transcriptionFilename(req audio.TranscriptionRequest) string { - if name := sanitizeFilename(req.Filename); name != "" { +func transcriptionFilename(filename, mimeType string) string { + if name := sanitizeFilename(filename); name != "" { return name } - mt := strings.ToLower(strings.TrimSpace(req.MIME)) + mt := strings.ToLower(strings.TrimSpace(mimeType)) if parsed, _, err := mime.ParseMediaType(mt); err == nil { mt = parsed } diff --git a/provider/llamaswap/diarize.go b/provider/llamaswap/diarize.go index aae29fc..fc77ae6 100644 --- a/provider/llamaswap/diarize.go +++ b/provider/llamaswap/diarize.go @@ -79,7 +79,7 @@ func (m *diarizationModel) Diarize(ctx context.Context, req audio.DiarizationReq } body, contentType, err := buildMultipart("build diarization form", - filePart{field: "audio_file", filename: diarizationFilename(req), data: req.Audio}, + filePart{field: "audio_file", filename: transcriptionFilename(req.Filename, req.MIME), data: req.Audio}, nil) if err != nil { return nil, err @@ -111,12 +111,3 @@ func (m *diarizationModel) Diarize(ctx context.Context, req audio.DiarizationReq } return res, nil } - -// diarizationFilename mirrors transcriptionFilename for the diarization -// request shape (same untrusted-metadata sanitization rules). -func diarizationFilename(req audio.DiarizationRequest) string { - return transcriptionFilename(audio.TranscriptionRequest{ - Filename: req.Filename, - MIME: req.MIME, - }) -} diff --git a/provider/llamaswap/mediautil.go b/provider/llamaswap/mediautil.go index 3387b4c..334ed49 100644 --- a/provider/llamaswap/mediautil.go +++ b/provider/llamaswap/mediautil.go @@ -22,6 +22,12 @@ import ( "gitea.stevedudenhoeffer.com/steve/majordomo/videogen" ) +// maxImageResponseBytes caps the raw-image bodies from the upscale and +// background-removal endpoints. The 64MB maxResponseBytes is a JSON cap; +// a 4x-upscaled PNG legitimately exceeds it. Bounded (not video-sized) +// because a single still image past this is an upstream bug, not data. +const maxImageResponseBytes = 256 << 20 + // --- upscale --- // UpscaleModel implements imagegen.UpscaleProvider against the mediautils @@ -62,7 +68,7 @@ func (m *upscaleModel) Upscale(ctx context.Context, req imagegen.UpscaleRequest, if err != nil { return nil, err } - raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxResponseBytes) + raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes) if err != nil { return nil, err } @@ -112,23 +118,31 @@ func (m *backgroundModel) RemoveBackground(ctx context.Context, req imagegen.Bac if err != nil { return nil, err } - raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxResponseBytes) + raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, contentType, body, maxImageResponseBytes) if err != nil { return nil, err } return singleImageResult(m.p.name, m.id, "background removal", raw, respType) } -// singleImageResult wraps one raw image body into an imagegen.Result, -// rejecting empty or non-image payloads (a proxy error page must not become -// "the image"). +// singleImageResult wraps one raw image body into an imagegen.Result. +// Acceptance requires positive evidence of image-ness — sniffed magic +// bytes or a declared image/* Content-Type — so a proxy error page with +// an empty Content-Type can never become "the image" (sniffImageMIME's +// PNG default is a labelling fallback, not a validator). func singleImageResult(provider, model, verb string, raw []byte, contentType string) (*imagegen.Result, error) { if len(raw) == 0 { return nil, &llm.APIError{Provider: provider, Model: model, Message: verb + " response contained no image"} } - mimeType := sniffImageMIME(raw) - if ct := strings.TrimSpace(contentType); ct != "" && !strings.HasPrefix(ct, "image/") && !strings.HasPrefix(http.DetectContentType(raw), "image/") { - return nil, &llm.APIError{Provider: provider, Model: model, Message: fmt.Sprintf("%s response is not an image (Content-Type %q)", verb, ct)} + sniffed := http.DetectContentType(raw) + declared := strings.TrimSpace(contentType) + if !strings.HasPrefix(sniffed, "image/") && !strings.HasPrefix(declared, "image/") { + return nil, &llm.APIError{Provider: provider, Model: model, + Message: fmt.Sprintf("%s response is not an image (Content-Type %q, sniffed %q)", verb, declared, sniffed)} + } + mimeType := sniffed + if !strings.HasPrefix(mimeType, "image/") { + mimeType = declared } return &imagegen.Result{Images: []llm.ImagePart{{MIME: mimeType, Data: raw}}}, nil } diff --git a/provider/llamaswap/mediautil_test.go b/provider/llamaswap/mediautil_test.go index 6737c12..466ec9e 100644 --- a/provider/llamaswap/mediautil_test.go +++ b/provider/llamaswap/mediautil_test.go @@ -251,3 +251,32 @@ func TestUpstreamPathRejectsSeparators(t *testing.T) { t.Errorf("path prefix wrong: %q", got) } } + +func TestUpscaleRejectsEmptyContentTypeNonImage(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header()["Content-Type"] = nil // suppress auto-detection: NO Content-Type at all + _, _ = w.Write([]byte("502 bad gateway")) + })) + defer srv.Close() + + p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client())) + up, _ := p.UpscaleModel("mediautils") + _, err := up.Upscale(context.Background(), + imagegen.UpscaleRequest{Image: imagegen.Image{Data: pngFixture(t)}}) + var apiErr *llm.APIError + if !errors.As(err, &apiErr) { + t.Fatalf("err = %v, want APIError for headerless non-image body", err) + } +} + +func TestUpstreamPathRejectsDotDot(t *testing.T) { + if _, err := upstreamPath("..", "/x"); err == nil { + t.Error("model '..' accepted") + } + if _, err := upstreamPath("m", "/v1/audio?path=../../api/models/unload"); err == nil { + t.Error("dot-dot rest accepted") + } + if _, err := upstreamPath("m", "/v1/audio?path=https://evil.example/x"); err == nil { + t.Error("absolute-URL rest accepted") + } +} diff --git a/provider/llamaswap/mesh.go b/provider/llamaswap/mesh.go index 796a008..ff8b44f 100644 --- a/provider/llamaswap/mesh.go +++ b/provider/llamaswap/mesh.go @@ -102,7 +102,7 @@ func (m *meshModel) Generate(ctx context.Context, req meshgen.Request, opts ...m if err != nil { return nil, fmt.Errorf("llama-swap: encode mesh request: %w", err) } - raw, _, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, "application/json", bytes.NewReader(encoded), maxMeshResponseBytes) + raw, respType, err := m.p.doRaw(ctx, http.MethodPost, path, m.id, "application/json", bytes.NewReader(encoded), maxMeshResponseBytes) if err != nil { return nil, err } @@ -110,8 +110,16 @@ func (m *meshModel) Generate(ctx context.Context, req meshgen.Request, opts ...m return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "mesh response contained no data"} } // A JSON body on the success path is the server reporting a soft error - // (or an API drift) — never hand it back as "the mesh". - if first := strings.TrimLeft(string(raw[:min(len(raw), 64)]), " \t\r\n"); strings.HasPrefix(first, "{") || strings.HasPrefix(first, "[") { + // (or an API drift) — never hand it back as "the mesh". Two signals: + // the declared Content-Type, and a whitespace-tolerant peek at the + // leading bytes (512 covers any indented error envelope; a binary STL + // header theoretically CAN start with '{', but a real one also won't + // be all-whitespace-then-brace). + if strings.Contains(strings.ToLower(respType), "application/json") { + return nil, &llm.APIError{Provider: m.p.name, Model: m.id, + Message: "mesh response is JSON, not mesh bytes: " + truncateForError(raw)} + } + if first := strings.TrimLeft(string(raw[:min(len(raw), 512)]), " \t\r\n"); strings.HasPrefix(first, "{") || strings.HasPrefix(first, "[") { return nil, &llm.APIError{Provider: m.p.name, Model: m.id, Message: "mesh response is JSON, not mesh bytes: " + truncateForError(raw)} } @@ -120,9 +128,9 @@ func (m *meshModel) Generate(ctx context.Context, req meshgen.Request, opts ...m // truncateForError bounds a payload quoted into an error message. func truncateForError(b []byte) string { - const cap = 500 - if len(b) > cap { - return string(b[:cap]) + "..." + const maxErrLen = 500 + if len(b) > maxErrLen { + return string(b[:maxErrLen]) + "..." } return string(b) } diff --git a/provider/llamaswap/upstream.go b/provider/llamaswap/upstream.go index 85cd2f6..1a33bc2 100644 --- a/provider/llamaswap/upstream.go +++ b/provider/llamaswap/upstream.go @@ -21,12 +21,19 @@ func upstreamPath(model, rest string) (string, error) { if strings.TrimSpace(model) == "" { return "", fmt.Errorf("llama-swap: upstream call requires a model id") } - if strings.ContainsAny(model, "/?#") { + if strings.ContainsAny(model, "/?#") || strings.Contains(model, "..") { return "", fmt.Errorf("llama-swap: invalid model id %q for upstream call (contains a path separator)", model) } if !strings.HasPrefix(rest, "/") { rest = "/" + rest } + // rest may embed SERVER-SUPPLIED components (e.g. ACE-Step's result + // file URL) — refuse dot-dot segments and absolute-URL smuggling so a + // hostile/buggy upstream cannot redirect the follow-up request at + // another proxy endpoint (/api/models/unload, ...). + if strings.Contains(rest, "..") || strings.Contains(rest, "://") { + return "", fmt.Errorf("llama-swap: invalid upstream path %q (dot-dot or scheme)", rest) + } return "/upstream/" + model + rest, nil }