diff --git a/provider/llamaswap/mediautil_test.go b/provider/llamaswap/mediautil_test.go index 466ec9e..9bb42c3 100644 --- a/provider/llamaswap/mediautil_test.go +++ b/provider/llamaswap/mediautil_test.go @@ -238,7 +238,7 @@ func TestInterpolateRejectsBadArgs(t *testing.T) { } func TestUpstreamPathRejectsSeparators(t *testing.T) { - for _, bad := range []string{"", "a/b", "a?b", "a#b"} { + for _, bad := range []string{"", "a/b", "a?b", "a#b", "a%2Fb", "%2e%2e", "a%b"} { if _, err := upstreamPath(bad, "/x"); err == nil { t.Errorf("upstreamPath(%q) succeeded; want error", bad) } diff --git a/provider/llamaswap/segment.go b/provider/llamaswap/segment.go index 404d592..c92123f 100644 --- a/provider/llamaswap/segment.go +++ b/provider/llamaswap/segment.go @@ -14,6 +14,7 @@ package llamaswap import ( "context" "fmt" + "math" "net/http" "strconv" @@ -45,7 +46,9 @@ func (m *segmentationModel) Segment(ctx context.Context, req imagegen.Segmentati if req.Prompt == "" { return nil, fmt.Errorf("%w: segmentation requires a prompt", llm.ErrUnsupported) } - if req.Threshold < 0 || req.Threshold > 1 { + // NaN fails every comparison, so it would sail through a bare range + // check and reach the upstream as the literal string "NaN". + if math.IsNaN(req.Threshold) || req.Threshold < 0 || req.Threshold > 1 { return nil, fmt.Errorf("%w: segmentation threshold must be in [0,1], got %g", llm.ErrUnsupported, req.Threshold) } path, err := upstreamPath(m.id, "/v1/segment") diff --git a/provider/llamaswap/segment_restore_test.go b/provider/llamaswap/segment_restore_test.go index fa39d91..dcb4f21 100644 --- a/provider/llamaswap/segment_restore_test.go +++ b/provider/llamaswap/segment_restore_test.go @@ -3,6 +3,7 @@ package llamaswap import ( "context" "errors" + "math" "net/http" "net/http/httptest" "testing" @@ -95,6 +96,16 @@ func TestSegmentRejectsBadArgs(t *testing.T) { imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: 1.5}); !errors.Is(err, llm.ErrUnsupported) { t.Errorf("threshold 1.5: err = %v, want ErrUnsupported", err) } + for name, bad := range map[string]float64{ + "NaN": math.NaN(), + "+Inf": math.Inf(1), + "-Inf": math.Inf(-1), + } { + if _, err := sg.Segment(context.Background(), + imagegen.SegmentationRequest{Image: imagegen.Image{Data: []byte{1}}, Prompt: "dog", Threshold: bad}); !errors.Is(err, llm.ErrUnsupported) { + t.Errorf("threshold %s: err = %v, want ErrUnsupported", name, err) + } + } } func TestSegmentRejectsNonImageResponse(t *testing.T) { diff --git a/provider/llamaswap/upstream.go b/provider/llamaswap/upstream.go index 1a33bc2..48998d2 100644 --- a/provider/llamaswap/upstream.go +++ b/provider/llamaswap/upstream.go @@ -16,12 +16,14 @@ import ( // // Why reject rather than escape: same rationale as Unload — model ids // legitimately contain ":" but never path-structure characters, and escaping -// would mask a config error instead of surfacing it. +// would mask a config error instead of surfacing it. '%' is rejected too: +// ids never legitimately carry percent-escapes, and %2F/%2E%2E would decode +// back into path structure on the server side. 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, "/?#") || strings.Contains(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, "/") {