steve
4a752ffa2a
- 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 <noreply@anthropic.com>
283 lines
9.2 KiB
Go
283 lines
9.2 KiB
Go
package llamaswap
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
|
|
"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
|
|
"gitea.stevedudenhoeffer.com/steve/majordomo/videogen"
|
|
)
|
|
|
|
func pngFixture(t *testing.T) []byte {
|
|
t.Helper()
|
|
raw, err := base64.StdEncoding.DecodeString(onePixelPNG)
|
|
if err != nil {
|
|
t.Fatalf("decode fixture: %v", err)
|
|
}
|
|
return raw
|
|
}
|
|
|
|
func TestUpscale(t *testing.T) {
|
|
png := pngFixture(t)
|
|
var gotPath, gotScale, gotFilename string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
|
t.Errorf("parse multipart: %v", err)
|
|
}
|
|
gotScale = r.FormValue("scale")
|
|
if f, hdr, err := r.FormFile("file"); err == nil {
|
|
gotFilename = hdr.Filename
|
|
f.Close()
|
|
} else {
|
|
t.Errorf("file part: %v", err)
|
|
}
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = w.Write(png)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
|
up, err := p.UpscaleModel("mediautils")
|
|
if err != nil {
|
|
t.Fatalf("UpscaleModel: %v", err)
|
|
}
|
|
res, err := up.Upscale(context.Background(),
|
|
imagegen.UpscaleRequest{Image: imagegen.Image{MIME: "image/png", Data: png}},
|
|
imagegen.WithUpscaleScale(2))
|
|
if err != nil {
|
|
t.Fatalf("Upscale: %v", err)
|
|
}
|
|
if gotPath != "/upstream/mediautils/v1/upscale" {
|
|
t.Errorf("path = %q", gotPath)
|
|
}
|
|
if gotScale != "2" || gotFilename != "image.png" {
|
|
t.Errorf("scale/filename = %q/%q", gotScale, gotFilename)
|
|
}
|
|
if len(res.Images) != 1 || res.Images[0].MIME != "image/png" {
|
|
t.Fatalf("images = %+v", res.Images)
|
|
}
|
|
}
|
|
|
|
func TestUpscaleRejectsBadArgs(t *testing.T) {
|
|
p := New(WithBaseURL("http://unused"))
|
|
up, _ := p.UpscaleModel("mediautils")
|
|
if _, err := up.Upscale(context.Background(), imagegen.UpscaleRequest{}); !errors.Is(err, llm.ErrUnsupported) {
|
|
t.Errorf("no image: err = %v, want ErrUnsupported", err)
|
|
}
|
|
if _, err := up.Upscale(context.Background(),
|
|
imagegen.UpscaleRequest{Image: imagegen.Image{Data: []byte{1}}, Scale: 3}); !errors.Is(err, llm.ErrUnsupported) {
|
|
t.Errorf("scale 3: err = %v, want ErrUnsupported", err)
|
|
}
|
|
}
|
|
|
|
func TestUpscaleRejectsNonImageResponse(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/html")
|
|
_, _ = w.Write([]byte("<html>proxy error page</html>"))
|
|
}))
|
|
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 non-image body", err)
|
|
}
|
|
}
|
|
|
|
func TestRemoveBackground(t *testing.T) {
|
|
png := pngFixture(t)
|
|
var gotPath, gotNet, gotOM string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
|
t.Errorf("parse multipart: %v", err)
|
|
}
|
|
gotNet = r.FormValue("model")
|
|
gotOM = r.FormValue("om")
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = w.Write(png)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
|
br, err := p.BackgroundRemovalModel("rembg")
|
|
if err != nil {
|
|
t.Fatalf("BackgroundRemovalModel: %v", err)
|
|
}
|
|
res, err := br.RemoveBackground(context.Background(),
|
|
imagegen.BackgroundRemovalRequest{Image: imagegen.Image{Data: png}},
|
|
imagegen.WithBackgroundNet("birefnet-general"), imagegen.WithOnlyMask())
|
|
if err != nil {
|
|
t.Fatalf("RemoveBackground: %v", err)
|
|
}
|
|
if gotPath != "/upstream/rembg/api/remove" {
|
|
t.Errorf("path = %q", gotPath)
|
|
}
|
|
if gotNet != "birefnet-general" || gotOM != "true" {
|
|
t.Errorf("net/om = %q/%q", gotNet, gotOM)
|
|
}
|
|
if len(res.Images) != 1 {
|
|
t.Fatalf("images = %+v", res.Images)
|
|
}
|
|
}
|
|
|
|
func TestRemoveBackgroundOmitsUnsetFields(t *testing.T) {
|
|
png := pngFixture(t)
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
|
t.Errorf("parse multipart: %v", err)
|
|
}
|
|
if _, ok := r.MultipartForm.Value["model"]; ok {
|
|
t.Error("model field sent for default net; want omitted")
|
|
}
|
|
if _, ok := r.MultipartForm.Value["om"]; ok {
|
|
t.Error("om field sent for cutout mode; want omitted")
|
|
}
|
|
w.Header().Set("Content-Type", "image/png")
|
|
_, _ = w.Write(png)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
|
br, _ := p.BackgroundRemovalModel("rembg")
|
|
if _, err := br.RemoveBackground(context.Background(),
|
|
imagegen.BackgroundRemovalRequest{Image: imagegen.Image{Data: png}}); err != nil {
|
|
t.Fatalf("RemoveBackground: %v", err)
|
|
}
|
|
}
|
|
|
|
// mp4Fixture is a minimal ftyp box so http.DetectContentType sniffs video/mp4.
|
|
func mp4Fixture() []byte {
|
|
return []byte{0, 0, 0, 24, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm',
|
|
0, 0, 2, 0, 'i', 's', 'o', 'm', 'i', 's', 'o', '2'}
|
|
}
|
|
|
|
func TestInterpolate(t *testing.T) {
|
|
var gotPath, gotMulti, gotMode, gotTargetFPS string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotPath = r.URL.Path
|
|
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
|
t.Errorf("parse multipart: %v", err)
|
|
}
|
|
gotMulti = r.FormValue("multi")
|
|
gotMode = r.FormValue("mode")
|
|
gotTargetFPS = r.FormValue("target_fps")
|
|
w.Header().Set("Content-Type", "video/mp4")
|
|
_, _ = w.Write(mp4Fixture())
|
|
}))
|
|
defer srv.Close()
|
|
|
|
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
|
ip, err := p.InterpolatorModel("mediautils")
|
|
if err != nil {
|
|
t.Fatalf("InterpolatorModel: %v", err)
|
|
}
|
|
res, err := ip.Interpolate(context.Background(),
|
|
videogen.InterpolateRequest{Video: videogen.Video{Data: mp4Fixture(), MIME: "video/mp4"}},
|
|
videogen.WithInterpolateMulti(2), videogen.WithTargetFPS(60))
|
|
if err != nil {
|
|
t.Fatalf("Interpolate: %v", err)
|
|
}
|
|
if gotPath != "/upstream/mediautils/v1/interpolate" {
|
|
t.Errorf("path = %q", gotPath)
|
|
}
|
|
if gotMulti != "2" || gotMode != "" || gotTargetFPS != "60" {
|
|
t.Errorf("multi/mode/target_fps = %q/%q/%q", gotMulti, gotMode, gotTargetFPS)
|
|
}
|
|
if res.Video.MIME != "video/mp4" || len(res.Video.Data) == 0 {
|
|
t.Fatalf("video = %q/%d bytes", res.Video.MIME, len(res.Video.Data))
|
|
}
|
|
}
|
|
|
|
func TestInterpolateSlowMoSetsModeAndDropsTargetFPS(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
|
t.Errorf("parse multipart: %v", err)
|
|
}
|
|
if got := r.FormValue("mode"); got != "slowmo" {
|
|
t.Errorf("mode = %q, want slowmo", got)
|
|
}
|
|
if _, ok := r.MultipartForm.Value["target_fps"]; ok {
|
|
t.Error("target_fps sent in slowmo mode; want omitted")
|
|
}
|
|
w.Header().Set("Content-Type", "video/mp4")
|
|
_, _ = w.Write(mp4Fixture())
|
|
}))
|
|
defer srv.Close()
|
|
|
|
p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
|
|
ip, _ := p.InterpolatorModel("mediautils")
|
|
_, err := ip.Interpolate(context.Background(),
|
|
videogen.InterpolateRequest{Video: videogen.Video{Data: mp4Fixture()}, TargetFPS: 60},
|
|
videogen.WithSlowMo())
|
|
if err != nil {
|
|
t.Fatalf("Interpolate: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestInterpolateRejectsBadArgs(t *testing.T) {
|
|
p := New(WithBaseURL("http://unused"))
|
|
ip, _ := p.InterpolatorModel("mediautils")
|
|
if _, err := ip.Interpolate(context.Background(), videogen.InterpolateRequest{}); !errors.Is(err, llm.ErrUnsupported) {
|
|
t.Errorf("no video: err = %v, want ErrUnsupported", err)
|
|
}
|
|
if _, err := ip.Interpolate(context.Background(),
|
|
videogen.InterpolateRequest{Video: videogen.Video{Data: []byte{1}}, Multi: 3}); !errors.Is(err, llm.ErrUnsupported) {
|
|
t.Errorf("multi 3: err = %v, want ErrUnsupported", err)
|
|
}
|
|
}
|
|
|
|
func TestUpstreamPathRejectsSeparators(t *testing.T) {
|
|
for _, bad := range []string{"", "a/b", "a?b", "a#b"} {
|
|
if _, err := upstreamPath(bad, "/x"); err == nil {
|
|
t.Errorf("upstreamPath(%q) succeeded; want error", bad)
|
|
}
|
|
}
|
|
got, err := upstreamPath("rembg", "api/remove")
|
|
if err != nil || got != "/upstream/rembg/api/remove" {
|
|
t.Errorf("upstreamPath = %q, %v", got, err)
|
|
}
|
|
if !strings.HasPrefix(got, "/upstream/") {
|
|
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")
|
|
}
|
|
}
|