diff --git a/internal/vision/vision.go b/internal/vision/vision.go index eda0ad8..a1a4667 100644 --- a/internal/vision/vision.go +++ b/internal/vision/vision.go @@ -37,6 +37,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net" "net/http" "net/url" @@ -93,6 +94,10 @@ func (Disabled) Describe(context.Context, media.Image, string) (string, error) { return "", ErrDisabled } +// MaxReplyBytes bounds what is read back from the vision server. A description +// is words; anything past a megabyte is a broken endpoint. +const MaxReplyBytes = 1 << 20 + // Config — how to reach the local vision server. Built from // config.VisionConfig by the daemon; kept separate so this package does not // import internal/config. @@ -154,13 +159,31 @@ func NewLocal(cfg Config) (*LocalProvider, error) { model: cfg.Model, prompt: prompt, maxTokens: maxTokens, - http: &http.Client{Timeout: timeout}, + http: &http.Client{ + Timeout: timeout, + // No redirects. checkPrivate validates the configured literal and + // nothing validated a hop, so a 302 from the local llama-server + // would send the photo, as a data URI in the POST body, wherever + // the redirect named. "No provider in this repo may upload a blob" + // has to be true of the second request as well as the first. + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + }, + }, }, nil } // Endpoint is the server this provider talks to. For logs and /dash. func (p *LocalProvider) Endpoint() string { return p.endpoint } +// ValidateEndpoint reports whether a configured endpoint is one this package +// would accept. Exported so config validation fails at startup on a typo, +// rather than logging once at wiring time and leaving the capability quietly +// off. +func ValidateEndpoint(raw string) error { + return checkPrivate(strings.TrimRight(strings.TrimSpace(raw), "/")) +} + // checkPrivate refuses any endpoint that is not on this box or its LAN. A // hostname that is not an IP literal is refused too: "vision.example.com" could // resolve anywhere, and resolving it here would be trusting DNS with his photos. @@ -260,7 +283,10 @@ func (p *LocalProvider) Describe(ctx context.Context, im media.Image, prompt str return "", fmt.Errorf("vision: status %d", resp.StatusCode) } var out chatResp - if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + // Capped: the decoder would otherwise read whatever the endpoint sends, and + // a local server answering with a stuck stream should not cost the daemon + // its memory. A description is a few hundred tokens. + if err := json.NewDecoder(io.LimitReader(resp.Body, MaxReplyBytes)).Decode(&out); err != nil { return "", fmt.Errorf("vision: decode: %w", err) } if len(out.Choices) == 0 { diff --git a/internal/vision/vision_test.go b/internal/vision/vision_test.go index 71787fc..a65cfc5 100644 --- a/internal/vision/vision_test.go +++ b/internal/vision/vision_test.go @@ -11,6 +11,7 @@ import ( "net/http" "net/http/httptest" "strings" + "sync/atomic" "testing" "time" @@ -196,3 +197,72 @@ func TestDescribeErrors(t *testing.T) { } }) } + +// checkPrivate validates the configured literal and used to validate nothing +// else. A 302 from the local llama-server would have sent the photo, as a data +// URI in the POST body, wherever the redirect named. +func TestLocalProviderDoesNotFollowARedirect(t *testing.T) { + var elsewhere int32 + away := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&elsewhere, 1) + _, _ = io.WriteString(w, `{"choices":[{"message":{"content":"leaked"}}]}`) + })) + defer away.Close() + local := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, away.URL+"/v1/chat/completions", http.StatusFound) + })) + defer local.Close() + + p, err := NewLocal(Config{Endpoint: local.URL}) + if err != nil { + t.Fatal(err) + } + im := media.Image{JPEG: []byte{0xFF, 0xD8, 0xFF}} + if _, err := p.Describe(context.Background(), im, "что это"); err == nil { + t.Fatal("a redirected describe must fail, not follow") + } + if n := atomic.LoadInt32(&elsewhere); n != 0 { + t.Fatalf("the image was sent to the redirect target %d time(s)", n) + } +} + +// The reply is read through a cap. A stuck endpoint should not cost the daemon +// its memory. +func TestLocalProviderCapsTheReply(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"choices":[{"message":{"content":"`) + for written := 0; written < MaxReplyBytes+(1<<20); written += 1 << 16 { + if _, err := io.WriteString(w, strings.Repeat("a", 1<<16)); err != nil { + return + } + } + })) + defer srv.Close() + p, err := NewLocal(Config{Endpoint: srv.URL}) + if err != nil { + t.Fatal(err) + } + im := media.Image{JPEG: []byte{0xFF, 0xD8, 0xFF}} + if _, err := p.Describe(context.Background(), im, ""); err == nil { + t.Fatal("an unbounded reply must fail rather than being read whole") + } +} + +// ValidateEndpoint is what config calls at startup, and it must agree with the +// constructor. +func TestValidateEndpointMatchesTheConstructor(t *testing.T) { + for _, raw := range []string{"http://127.0.0.1:8081", "http://localhost:8081/"} { + if err := ValidateEndpoint(raw); err != nil { + t.Errorf("ValidateEndpoint(%q) = %v", raw, err) + } + } + for _, raw := range []string{"http://8.8.8.8:8081", "http://vision.example.com", "ftp://127.0.0.1"} { + if err := ValidateEndpoint(raw); err == nil { + t.Errorf("ValidateEndpoint(%q) accepted a non-private endpoint", raw) + } + if _, err := NewLocal(Config{Endpoint: raw}); err == nil { + t.Errorf("NewLocal(%q) accepted what ValidateEndpoint should refuse", raw) + } + } +}