From 247c91986975effa3e3f4145334080d9e8165bda Mon Sep 17 00:00:00 2001 From: Simon Ferquel's Clanker Date: Thu, 7 May 2026 09:26:27 +0000 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20attach-time=20processing=20?= =?UTF-8?q?=E2=80=93=20ProcessAttachment=20(library-only,=20no=20TUI/CLI?= =?UTF-8?q?=20wiring)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds pkg/chat/attach.go and pkg/chat/attach_test.go. No changes to pkg/app or pkg/cli — wiring is intentionally deferred until the library API is stable. ## ProcessAttachment / ProcessAttachmentWithMetadata Converts a raw MessagePart into a fully-resolved chat.Document with InlineData or InlineText. Called once when a message is assembled. MessagePartTypeFile: - MIME-first routing (precedes IsTextFile heuristic to avoid misclassifying ASCII-content PDFs as text) - image/* → ResizeImage transcode (bounded by MaxInlineBinarySize) - application/pdf and other binary → verbatim InlineData (bounded by MaxInlineBinarySize = 20 MB) - text/* → ReadFileForInline InlineText (bounded by MaxInlineFileSize) - unknown binary → verbatim InlineData (bounded) MessagePartTypeImageURL: - data: URIs decoded inline - http(s):// fetched via safeHTTPClient with SSRF protection: - blocks loopback (127.0.0.0/8, ::1) - blocks link-local / metadata endpoints (169.254.0.0/16, fe80::/10) - blocks RFC-1918 (10/8, 172.16/12, 192.168/16) - blocks IPv6 unique-local (fc00::/7) - DialContext + CheckRedirect both enforce the block - DNS resolved with context via net.DefaultResolver.LookupHost - bytes capped at 20 MB; goes through ResizeImage MessagePartTypeDocument: - images with InlineData transcoded; others passed through - no inline content → error ProcessAttachmentWithMetadata also returns *ImageResizeResult so callers can emit a dimension note without a redundant second ResizeImage call. SetFetchHTTPClientForTest allows tests to inject a plain http.Client to reach loopback-bound httptest servers. ## Tests (25 cases) JPEG/PNG passthrough, PNG-with-alpha, oversize image resized, PDF passthrough, binary too large (sparse file), text/markdown inline, missing file, nil inputs, data: URI (JPEG/PNG), non-base64 data: URI, unsupported scheme, httptest server success path, httptest 404, SSRF loopback block, SSRF metadata endpoint (169.254.169.254) block, document passthrough (binary/text), document image transcode, document no-content error, unsupported part type. Part of #2595 Assisted-By: docker-agent --- pkg/chat/attach.go | 439 ++++++++++++++++++++++++++++++++++++++++ pkg/chat/attach_test.go | 424 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 863 insertions(+) create mode 100644 pkg/chat/attach.go create mode 100644 pkg/chat/attach_test.go diff --git a/pkg/chat/attach.go b/pkg/chat/attach.go new file mode 100644 index 000000000..d63339152 --- /dev/null +++ b/pkg/chat/attach.go @@ -0,0 +1,439 @@ +package chat + +// Package-level attach-time processing pipeline. +// +// The attach-time pipeline runs once when a user adds a file or image to a +// message. It produces a fully-resolved [Document] whose Source.InlineData or +// Source.InlineText is populated. Provider-level convertDocument functions then +// consume the Document at inference time — they never perform I/O or resizing. +// +// producer (app.go / runner.go) +// └─ ProcessAttachment(ctx, part) → Document +// └─ session.UserMessage(…, MessagePart{Type: MessagePartTypeDocument, Document: &doc}) +// └─ per-provider convertDocument(ctx, doc, modelID) → wire format + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "io" + "mime" + "net" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +const ( + // attachHTTPTimeout is the maximum time allowed to fetch a remote image URL. + attachHTTPTimeout = 10 * time.Second + + // attachMaxRemoteBytes is the maximum number of bytes read from a remote URL. + attachMaxRemoteBytes = 20 * 1024 * 1024 // 20 MB + + // MaxInlineBinarySize is the maximum byte size of a binary file (PDF, etc.) + // that will be read into memory for inline attachment. Matches the remote + // fetch cap so local and remote paths behave consistently. + MaxInlineBinarySize = 20 * 1024 * 1024 // 20 MB +) + +// privateIPNets lists address ranges that must not be dialled by the +// attach-time URL fetcher. Blocking these prevents SSRF attacks against +// cloud metadata services, internal APIs, and loopback services. +var privateIPNets = func() []*net.IPNet { + blocks := []string{ + // Loopback + "127.0.0.0/8", + "::1/128", + // Link-local — covers AWS/GCP/Azure metadata endpoints (169.254.169.254) + "169.254.0.0/16", + "fe80::/10", + // RFC-1918 private ranges + "10.0.0.0/8", + "172.16.0.0/12", + "192.168.0.0/16", + // IPv6 unique-local + "fc00::/7", + } + nets := make([]*net.IPNet, 0, len(blocks)) + for _, b := range blocks { + _, ipNet, err := net.ParseCIDR(b) + if err != nil { + panic("attachment: invalid built-in CIDR " + b + ": " + err.Error()) + } + nets = append(nets, ipNet) + } + return nets +}() + +// isPrivateIP reports whether ip falls in any of the blocked address ranges. +func isPrivateIP(ip net.IP) bool { + for _, block := range privateIPNets { + if block.Contains(ip) { + return true + } + } + return false +} + +// checkSafeHostCtx resolves host to IP addresses and returns an error if any +// resolved address is in a private/reserved range. This is called both at +// dial time (where a context is available) and on redirect destinations. +func checkSafeHostCtx(ctx context.Context, host string) error { + addrs, err := net.DefaultResolver.LookupHost(ctx, host) + if err != nil { + return fmt.Errorf("attachment: cannot resolve host %q: %w", host, err) + } + for _, addr := range addrs { + ip := net.ParseIP(addr) + if ip == nil { + continue + } + if isPrivateIP(ip) { + return fmt.Errorf("attachment: URL resolves to private/reserved address %s (SSRF protection)", addr) + } + } + return nil +} + +// safeHTTPClient is a shared HTTP client used by fetchRemoteImage. +// It enforces SSRF protection by refusing connections to private and +// reserved IP ranges at both dial time and on each redirect hop. +// +// Tests may override this variable to inject a custom client that bypasses +// SSRF filtering (e.g. to reach httptest.Server on loopback). Only do this +// in test code guarded by a t.Cleanup restore. +var safeHTTPClient = newSafeHTTPClient() + +func newSafeHTTPClient() *http.Client { + return &http.Client{ + Timeout: attachHTTPTimeout, + Transport: &http.Transport{ + DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + host, _, err := net.SplitHostPort(addr) + if err != nil { + return nil, fmt.Errorf("attachment: malformed dial address %q: %w", addr, err) + } + if err := checkSafeHostCtx(ctx, host); err != nil { + return nil, err + } + return (&net.Dialer{}).DialContext(ctx, network, addr) + }, + }, + CheckRedirect: func(req *http.Request, _ []*http.Request) error { + return checkSafeHostCtx(req.Context(), req.URL.Hostname()) + }, + } +} + +// SetFetchHTTPClientForTest replaces the HTTP client used by fetchRemoteImage +// and returns a restore function. Only for use in tests. +// +// defer chat.SetFetchHTTPClientForTest(t, myClient)() +func SetFetchHTTPClientForTest(client *http.Client) (restore func()) { + prev := safeHTTPClient + safeHTTPClient = client + return func() { safeHTTPClient = prev } +} + +// ProcessAttachment converts a raw [MessagePart] into a [Document] with fully +// resolved Source.InlineData or Source.InlineText. It is called once when a +// message is assembled — never at inference time. +// +// Supported input types: +// +// - [MessagePartTypeFile]: reads the file from the local filesystem, detects +// its MIME type, and either inlines text content (text/* files) or reads +// binary bytes (images are transcoded+resized via [ResizeImage]; PDFs and +// other supported types are read verbatim). +// +// - [MessagePartTypeImageURL]: handles data: URIs (decoded inline) and +// http(s):// URLs (fetched with a 10-second timeout). The image bytes are +// then passed through [ResizeImage] to normalise to JPEG or PNG. +// +// - [MessagePartTypeDocument]: if Source.InlineData or Source.InlineText is +// already set, the document is returned as-is after applying image +// transcoding to any image/* InlineData. A Document with no inline content +// is an error. +// +// The context is forwarded to any network operations; filesystem and image +// decoding operations are not yet context-aware. +func ProcessAttachment(ctx context.Context, part MessagePart) (Document, error) { + doc, _, err := ProcessAttachmentWithMetadata(ctx, part) + return doc, err +} + +// ProcessAttachmentWithMetadata is like [ProcessAttachment] but also returns +// the [ImageResizeResult] when the attachment was an image that went through +// [ResizeImage]. The metadata is nil for non-image attachments. +// +// Callers that need to emit a dimension note (for model coordinate-mapping) +// should use this variant and call [FormatDimensionNote] on the returned +// metadata. +func ProcessAttachmentWithMetadata(ctx context.Context, part MessagePart) (Document, *ImageResizeResult, error) { + switch part.Type { + case MessagePartTypeFile: + return processFilePart(part) + case MessagePartTypeImageURL: + return processImageURLPart(ctx, part) + case MessagePartTypeDocument: + return processDocumentPart(part) + default: + return Document{}, nil, fmt.Errorf("ProcessAttachment: unsupported part type %q", part.Type) + } +} + +// processFilePart handles MessagePartTypeFile: reads from disk, detects MIME, +// routes to text-inline or binary-inline as appropriate. +func processFilePart(part MessagePart) (Document, *ImageResizeResult, error) { + if part.File == nil { + return Document{}, nil, errors.New("ProcessAttachment: file part has nil File field") + } + absPath := part.File.Path + name := filepath.Base(absPath) + + fi, err := os.Stat(absPath) + if err != nil { + return Document{}, nil, fmt.Errorf("ProcessAttachment: cannot stat %q: %w", absPath, err) + } + if !fi.Mode().IsRegular() { + return Document{}, nil, fmt.Errorf("ProcessAttachment: %q is not a regular file", absPath) + } + + mimeType := DetectMimeType(absPath) + + // Route by MIME type. Note: MIME-type check must precede IsTextFile because + // some binary formats (e.g. PDF) may pass the text heuristic when the file + // content happens to be printable ASCII. + switch { + case IsImageMimeType(mimeType): + if fi.Size() > MaxInlineBinarySize { + return Document{}, nil, fmt.Errorf("ProcessAttachment: image file %q too large to inline (%d bytes, max %d)", absPath, fi.Size(), MaxInlineBinarySize) + } + data, err := os.ReadFile(absPath) + if err != nil { + return Document{}, nil, fmt.Errorf("ProcessAttachment: read image %q: %w", absPath, err) + } + return transcodeImageWithMeta(name, data, mimeType) + + case mimeType == "application/pdf" || (IsSupportedMimeType(mimeType) && !IsTextFile(absPath)): + // PDF and other supported binary types — read verbatim. + // The !IsTextFile guard ensures that binary formats whose extension + // is unknown but content is ASCII-printable are not incorrectly inlined. + if fi.Size() > MaxInlineBinarySize { + return Document{}, nil, fmt.Errorf("ProcessAttachment: binary file %q too large to inline (%d bytes, max %d)", absPath, fi.Size(), MaxInlineBinarySize) + } + data, err := os.ReadFile(absPath) + if err != nil { + return Document{}, nil, fmt.Errorf("ProcessAttachment: read binary file %q: %w", absPath, err) + } + return Document{ + Name: name, + MimeType: mimeType, + Size: int64(len(data)), + Source: DocumentSource{InlineData: data}, + }, nil, nil + + case IsTextFile(absPath): + if fi.Size() > MaxInlineFileSize { + return Document{}, nil, fmt.Errorf("ProcessAttachment: text file %q too large to inline (%d bytes, max %d)", absPath, fi.Size(), MaxInlineFileSize) + } + content, err := ReadFileForInline(absPath) + if err != nil { + return Document{}, nil, fmt.Errorf("ProcessAttachment: read text file %q: %w", absPath, err) + } + return Document{ + Name: name, + MimeType: mimeType, + Size: fi.Size(), + Source: DocumentSource{InlineText: content}, + }, nil, nil + + default: + // Unknown binary — read verbatim and let modelcaps gate it at inference time. + if fi.Size() > MaxInlineBinarySize { + return Document{}, nil, fmt.Errorf("ProcessAttachment: file %q too large to inline (%d bytes, max %d)", absPath, fi.Size(), MaxInlineBinarySize) + } + data, err := os.ReadFile(absPath) + if err != nil { + return Document{}, nil, fmt.Errorf("ProcessAttachment: read file %q: %w", absPath, err) + } + return Document{ + Name: name, + MimeType: mimeType, + Size: int64(len(data)), + Source: DocumentSource{InlineData: data}, + }, nil, nil + } +} + +// processImageURLPart handles MessagePartTypeImageURL. +// Supports: +// - data: URIs (data:;base64,) +// - http:// and https:// URLs (fetched with attachHTTPTimeout) +func processImageURLPart(ctx context.Context, part MessagePart) (Document, *ImageResizeResult, error) { + if part.ImageURL == nil { + return Document{}, nil, errors.New("ProcessAttachment: image-url part has nil ImageURL field") + } + rawURL := part.ImageURL.URL + + var ( + data []byte + mimeType string + name string + err error + ) + + switch { + case strings.HasPrefix(rawURL, "data:"): + mimeType, data, err = parseDataURI(rawURL) + if err != nil { + return Document{}, nil, fmt.Errorf("ProcessAttachment: parse data URI: %w", err) + } + name = "image" + + case strings.HasPrefix(rawURL, "http://") || strings.HasPrefix(rawURL, "https://"): + data, mimeType, name, err = fetchRemoteImage(ctx, rawURL) + if err != nil { + return Document{}, nil, fmt.Errorf("ProcessAttachment: fetch remote image %q: %w", rawURL, err) + } + + default: + return Document{}, nil, fmt.Errorf("ProcessAttachment: unsupported image URL scheme (must be data: or http(s)://): %q", rawURL) + } + + // When content detection returns an image type, prefer it over whatever + // the URI or Content-Type header said. (Only image types are trusted from + // the sniffer to avoid promoting an unknown binary to a valid MIME.) + if detected := DetectMimeTypeByContent(data); IsImageMimeType(detected) { + mimeType = detected + } + + return transcodeImageWithMeta(name, data, mimeType) +} + +// processDocumentPart handles MessagePartTypeDocument. +// Images with InlineData are transcoded; other already-resolved documents pass through. +func processDocumentPart(part MessagePart) (Document, *ImageResizeResult, error) { + if part.Document == nil { + return Document{}, nil, errors.New("ProcessAttachment: document part has nil Document field") + } + doc := *part.Document + + if len(doc.Source.InlineData) > 0 { + if IsImageMimeType(doc.MimeType) { + // Apply transcoding/resizing to normalise the image. + return transcodeImageWithMeta(doc.Name, doc.Source.InlineData, doc.MimeType) + } + // Non-image binary — pass through unchanged. + return doc, nil, nil + } + + if doc.Source.InlineText != "" { + // Already resolved text — pass through unchanged. + return doc, nil, nil + } + + return Document{}, nil, fmt.Errorf("ProcessAttachment: document %q has no inline content (InlineData and InlineText are both empty)", doc.Name) +} + +// transcodeImageWithMeta runs bytes through ResizeImage to normalise the image +// to JPEG or PNG within provider limits, then wraps the result in a Document. +// Returns the [ImageResizeResult] so callers can emit dimension notes. +func transcodeImageWithMeta(name string, data []byte, mimeType string) (Document, *ImageResizeResult, error) { + result, err := ResizeImage(data, mimeType) + if err != nil { + return Document{}, nil, fmt.Errorf("ProcessAttachment: transcode image %q: %w", name, err) + } + doc := Document{ + Name: name, + MimeType: result.MimeType, + Size: int64(len(result.Data)), + Source: DocumentSource{InlineData: result.Data}, + } + return doc, result, nil +} + +// parseDataURI parses a data URI of the form "data:;base64,". +// Returns the MIME type and decoded bytes. +func parseDataURI(uri string) (mimeType string, data []byte, err error) { + rest, ok := strings.CutPrefix(uri, "data:") + if !ok { + return "", nil, errors.New("not a data URI") + } + + header, payload, ok := strings.Cut(rest, ",") + if !ok { + return "", nil, errors.New("data URI missing comma separator") + } + + // Header is "[;charset=…];base64" or "" (plain text, unsupported here). + if !strings.HasSuffix(header, ";base64") { + return "", nil, errors.New("data URI is not base64-encoded (only base64 data URIs are supported)") + } + mimeType = strings.TrimSuffix(header, ";base64") + + // Strip any charset parameter (e.g. "image/png;charset=utf-8;base64" → "image/png"). + if idx := strings.Index(mimeType, ";"); idx >= 0 { + mimeType = mimeType[:idx] + } + mimeType = strings.TrimSpace(mimeType) + if mimeType == "" { + mimeType = "application/octet-stream" + } + + data, err = base64.StdEncoding.DecodeString(payload) + if err != nil { + return "", nil, fmt.Errorf("base64 decode: %w", err) + } + return mimeType, data, nil +} + +// fetchRemoteImage fetches an image from an http(s) URL. +// The safeHTTPClient enforces a timeout and blocks connections to +// private/reserved IP ranges (SSRF protection). +// Returns the image bytes, detected MIME type, and a filename derived from the URL. +func fetchRemoteImage(ctx context.Context, rawURL string) (data []byte, mimeType, name string, err error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, http.NoBody) + if err != nil { + return nil, "", "", fmt.Errorf("create request: %w", err) + } + + resp, err := safeHTTPClient.Do(req) + if err != nil { + return nil, "", "", fmt.Errorf("HTTP GET: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, "", "", fmt.Errorf("HTTP %d for %q", resp.StatusCode, rawURL) + } + + data, err = io.ReadAll(io.LimitReader(resp.Body, attachMaxRemoteBytes+1)) + if err != nil { + return nil, "", "", fmt.Errorf("read response body: %w", err) + } + if int64(len(data)) > attachMaxRemoteBytes { + return nil, "", "", fmt.Errorf("remote image too large (max %d bytes)", attachMaxRemoteBytes) + } + + // Determine MIME type: prefer Content-Type header, fall back to content sniff. + mimeType = resp.Header.Get("Content-Type") + if ct, _, parseErr := mime.ParseMediaType(mimeType); parseErr == nil { + mimeType = ct + } + if mimeType == "" || mimeType == "application/octet-stream" { + mimeType = DetectMimeTypeByContent(data) + } + + // Derive a filename from the URL path. + name = filepath.Base(req.URL.Path) + if name == "" || name == "." || name == "/" { + name = "image" + } + + return data, mimeType, name, nil +} diff --git a/pkg/chat/attach_test.go b/pkg/chat/attach_test.go new file mode 100644 index 000000000..c2f493003 --- /dev/null +++ b/pkg/chat/attach_test.go @@ -0,0 +1,424 @@ +package chat_test + +import ( + "bytes" + "encoding/base64" + "image" + "image/color" + "image/jpeg" + "image/png" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" +) + +// ────────────────────────────────────────────────────────────────────────────── +// Helpers +// ────────────────────────────────────────────────────────────────────────────── + +// encodeJPEGBytes returns a minimal JPEG image as raw bytes. +func encodeJPEGBytes(w, h int) []byte { + img := image.NewRGBA(image.Rect(0, 0, w, h)) + for y := range h { + for x := range w { + img.Set(x, y, color.RGBA{R: 200, G: 100, B: 50, A: 255}) + } + } + var buf bytes.Buffer + if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: 80}); err != nil { + panic(err) + } + return buf.Bytes() +} + +// encodePNGBytes returns a minimal PNG image as raw bytes. +func encodePNGBytes(w, h int, alpha bool) []byte { + if alpha { + img := image.NewNRGBA(image.Rect(0, 0, w, h)) + for y := range h { + for x := range w { + img.Set(x, y, color.NRGBA{R: 0, G: 128, B: 255, A: 128}) + } + } + var buf bytes.Buffer + if err := png.Encode(&buf, img); err != nil { + panic(err) + } + return buf.Bytes() + } + img := image.NewRGBA(image.Rect(0, 0, w, h)) + for y := range h { + for x := range w { + img.Set(x, y, color.RGBA{R: 0, G: 128, B: 255, A: 255}) + } + } + var buf bytes.Buffer + if err := png.Encode(&buf, img); err != nil { + panic(err) + } + return buf.Bytes() +} + +// writeTempFile writes data to a temp file with the given extension and returns its path. +func writeTempFile(t *testing.T, ext string, data []byte) string { + t.Helper() + f, err := os.CreateTemp(t.TempDir(), "attach-*"+ext) + require.NoError(t, err) + _, err = f.Write(data) + require.NoError(t, err) + require.NoError(t, f.Close()) + return f.Name() +} + +// ────────────────────────────────────────────────────────────────────────────── +// ProcessAttachment — MessagePartTypeFile +// ────────────────────────────────────────────────────────────────────────────── + +func TestProcessAttachment_JPEG_Passthrough(t *testing.T) { + data := encodeJPEGBytes(100, 100) + path := writeTempFile(t, ".jpg", data) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: path, MimeType: "image/jpeg"}, + }) + require.NoError(t, err) + assert.Equal(t, "image/jpeg", doc.MimeType) + assert.NotEmpty(t, doc.Source.InlineData) + assert.Empty(t, doc.Source.InlineText) + assert.Equal(t, filepath.Base(path), doc.Name) +} + +func TestProcessAttachment_PNG_Passthrough(t *testing.T) { + data := encodePNGBytes(100, 100, false) + path := writeTempFile(t, ".png", data) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: path, MimeType: "image/png"}, + }) + require.NoError(t, err) + assert.Equal(t, "image/png", doc.MimeType) + assert.NotEmpty(t, doc.Source.InlineData) +} + +func TestProcessAttachment_PNG_WithAlpha_StaysPNG(t *testing.T) { + // A PNG with alpha channel must not be converted to JPEG (lossy strips alpha). + data := encodePNGBytes(100, 100, true) + path := writeTempFile(t, ".png", data) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: path, MimeType: "image/png"}, + }) + require.NoError(t, err) + // ResizeImage picks the smaller of PNG/JPEG; for small images PNG is usually smaller + // but what matters is that the output is a valid image. + assert.True(t, doc.MimeType == "image/png" || doc.MimeType == "image/jpeg", + "expected png or jpeg, got %q", doc.MimeType) + assert.NotEmpty(t, doc.Source.InlineData) +} + +func TestProcessAttachment_ImageTooLarge_Resized(t *testing.T) { + // An image larger than MaxImageDimension should be resized. + bigData := encodeJPEGBytes(chat.MaxImageDimension+200, chat.MaxImageDimension+200) + path := writeTempFile(t, ".jpg", bigData) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: path, MimeType: "image/jpeg"}, + }) + require.NoError(t, err) + assert.NotEmpty(t, doc.Source.InlineData) + + // Verify the output image dimensions fit within limits. + img, _, decErr := image.Decode(bytes.NewReader(doc.Source.InlineData)) + require.NoError(t, decErr) + b := img.Bounds() + assert.LessOrEqual(t, b.Dx(), chat.MaxImageDimension) + assert.LessOrEqual(t, b.Dy(), chat.MaxImageDimension) +} + +func TestProcessAttachment_PDF_Passthrough(t *testing.T) { + pdfBytes := []byte("%PDF-1.4 fake pdf content for testing") + path := writeTempFile(t, ".pdf", pdfBytes) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: path, MimeType: "application/pdf"}, + }) + require.NoError(t, err) + assert.Equal(t, "application/pdf", doc.MimeType) + assert.Equal(t, pdfBytes, doc.Source.InlineData) + assert.Empty(t, doc.Source.InlineText) +} + +func TestProcessAttachment_BinaryFileTooLarge_Error(t *testing.T) { + // Write a file whose Stat.Size exceeds MaxInlineBinarySize. + // We use a sparse file (truncate to the target size) so the test is fast. + path := writeTempFile(t, ".pdf", nil) + f, err := os.OpenFile(path, os.O_WRONLY, 0o600) + require.NoError(t, err) + require.NoError(t, f.Truncate(chat.MaxInlineBinarySize+1)) + require.NoError(t, f.Close()) + + _, err = chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: path}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "too large") +} + +func TestProcessAttachment_TextFile_InlineText(t *testing.T) { + content := "Hello, this is a text file.\nLine 2." + path := writeTempFile(t, ".txt", []byte(content)) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: path, MimeType: "text/plain"}, + }) + require.NoError(t, err) + assert.Empty(t, doc.Source.InlineData) + assert.Contains(t, doc.Source.InlineText, content) + assert.Contains(t, doc.Source.InlineText, filepath.Base(path)) // ReadFileForInline wraps in a tag +} + +func TestProcessAttachment_MarkdownFile_InlineText(t *testing.T) { + content := "# Title\n\nBody paragraph." + path := writeTempFile(t, ".md", []byte(content)) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: path}, + }) + require.NoError(t, err) + assert.Empty(t, doc.Source.InlineData) + assert.Contains(t, doc.Source.InlineText, content) +} + +func TestProcessAttachment_MissingFile_Error(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: &chat.MessageFile{Path: "/nonexistent/path/file.jpg"}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot stat") +} + +func TestProcessAttachment_NilFile_Error(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeFile, + File: nil, + }) + require.Error(t, err) +} + +// ────────────────────────────────────────────────────────────────────────────── +// ProcessAttachment — MessagePartTypeImageURL +// ────────────────────────────────────────────────────────────────────────────── + +func TestProcessAttachment_DataURI_JPEG(t *testing.T) { + jpegData := encodeJPEGBytes(50, 50) + dataURI := "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(jpegData) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: dataURI}, + }) + require.NoError(t, err) + assert.NotEmpty(t, doc.Source.InlineData) + assert.True(t, doc.MimeType == "image/jpeg" || doc.MimeType == "image/png") +} + +func TestProcessAttachment_DataURI_PNG(t *testing.T) { + pngData := encodePNGBytes(50, 50, false) + dataURI := "data:image/png;base64," + base64.StdEncoding.EncodeToString(pngData) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: dataURI}, + }) + require.NoError(t, err) + assert.NotEmpty(t, doc.Source.InlineData) +} + +func TestProcessAttachment_DataURI_NonBase64_Error(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: "data:text/plain,hello"}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "not base64") +} + +func TestProcessAttachment_UnsupportedScheme_Error(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: "ftp://example.com/image.jpg"}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported image URL scheme") +} + +func TestProcessAttachment_HTTPS_Image_Via_HTTPTestServer(t *testing.T) { + jpegData := encodeJPEGBytes(60, 60) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "image/jpeg") + _, _ = w.Write(jpegData) + })) + t.Cleanup(srv.Close) + + // The httptest.Server binds to 127.0.0.1 (loopback), which is blocked by + // the SSRF filter in production. We use a plain http.Client without the + // SSRF-filtering transport so the test can reach the local server. + restore := chat.SetFetchHTTPClientForTest(srv.Client()) + t.Cleanup(restore) + + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: srv.URL + "/photo.jpg"}, + }) + require.NoError(t, err) + assert.NotEmpty(t, doc.Source.InlineData) + assert.True(t, doc.MimeType == "image/jpeg" || doc.MimeType == "image/png", + "unexpected MIME: %q", doc.MimeType) + assert.Equal(t, "photo.jpg", doc.Name) +} + +func TestProcessAttachment_SSRF_LoopbackBlocked(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: "http://127.0.0.1:9999/secret"}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "private/reserved") +} + +func TestProcessAttachment_SSRF_MetadataEndpointBlocked(t *testing.T) { + // 169.254.169.254 is the AWS/GCP/Azure instance metadata endpoint. + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: "http://169.254.169.254/latest/meta-data/"}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "private/reserved") +} + +func TestProcessAttachment_HTTP_Non200_Error(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "not found", http.StatusNotFound) + })) + t.Cleanup(srv.Close) + + // Bypass SSRF filter so we can reach the loopback-bound httptest server. + restore := chat.SetFetchHTTPClientForTest(srv.Client()) + t.Cleanup(restore) + + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: srv.URL + "/missing.jpg"}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "404") +} + +func TestProcessAttachment_NilImageURL_Error(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: nil, + }) + require.Error(t, err) +} + +// ────────────────────────────────────────────────────────────────────────────── +// ProcessAttachment — MessagePartTypeDocument +// ────────────────────────────────────────────────────────────────────────────── + +func TestProcessAttachment_Document_WithInlineData_Passthrough(t *testing.T) { + pdfBytes := []byte("%PDF-1.4 test") + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: "spec.pdf", + MimeType: "application/pdf", + Source: chat.DocumentSource{InlineData: pdfBytes}, + }, + }) + require.NoError(t, err) + assert.Equal(t, pdfBytes, doc.Source.InlineData) + assert.Equal(t, "application/pdf", doc.MimeType) +} + +func TestProcessAttachment_Document_WithInlineText_Passthrough(t *testing.T) { + text := "# Markdown content" + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: "readme.md", + MimeType: "text/markdown", + Source: chat.DocumentSource{InlineText: text}, + }, + }) + require.NoError(t, err) + assert.Equal(t, text, doc.Source.InlineText) + assert.Empty(t, doc.Source.InlineData) +} + +func TestProcessAttachment_Document_ImageInlineData_Transcoded(t *testing.T) { + jpegData := encodeJPEGBytes(40, 40) + doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: "photo.jpg", + MimeType: "image/jpeg", + Source: chat.DocumentSource{InlineData: jpegData}, + }, + }) + require.NoError(t, err) + assert.NotEmpty(t, doc.Source.InlineData) + assert.True(t, doc.MimeType == "image/jpeg" || doc.MimeType == "image/png") +} + +func TestProcessAttachment_Document_NoContent_Error(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: "empty.md", + MimeType: "text/markdown", + Source: chat.DocumentSource{}, + }, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "no inline content") +} + +func TestProcessAttachment_Document_NilDocument_Error(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: nil, + }) + require.Error(t, err) +} + +// ────────────────────────────────────────────────────────────────────────────── +// ProcessAttachment — unsupported type +// ────────────────────────────────────────────────────────────────────────────── + +func TestProcessAttachment_UnsupportedType_Error(t *testing.T) { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeText, + Text: "hello", + }) + require.Error(t, err) + assert.Contains(t, strings.ToLower(err.Error()), "unsupported") +} From 0472016240f4902bbccf64cf37d2c0cd2db0c013 Mon Sep 17 00:00:00 2001 From: Simon Ferquel's Clanker Date: Thu, 7 May 2026 09:42:17 +0000 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20remove=20URL=20fetching=20from=20Pro?= =?UTF-8?q?cessAttachment=20=E2=80=93=20local=20files=20and=20data:=20URIs?= =?UTF-8?q?=20only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per rumpl's review: remove all http(s):// fetching from the attach-time pipeline. The library should not download things from the internet. Changes: - Remove fetchRemoteImage function entirely - Remove safeHTTPClient, newSafeHTTPClient, privateIPNets, isPrivateIP, checkSafeHostCtx, SetFetchHTTPClientForTest (whole SSRF mitigation is now unnecessary) - Remove time, io, mime, net, net/http imports - processImageURLPart: data: URIs still decoded and transcoded as before; http(s):// URLs now return errors.New(...remote URLs are not supported...) - ProcessAttachment / ProcessAttachmentWithMetadata: context parameter kept on the public API (for forward compat) but marked _ since no I/O needs it - Remove network tests: TestProcessAttachment_HTTPS_Image_Via_HTTPTestServer, TestProcessAttachment_HTTP_Non200_Error, TestProcessAttachment_SSRF_* - Add TestProcessAttachment_RemoteURL_Error covering both http:// and https:// - MaxInlineBinarySize guard retained (local files, still valid) Assisted-By: docker-agent --- pkg/chat/attach.go | 226 +++++----------------------------------- pkg/chat/attach_test.go | 89 +++------------- 2 files changed, 37 insertions(+), 278 deletions(-) diff --git a/pkg/chat/attach.go b/pkg/chat/attach.go index d63339152..b0f98bd4c 100644 --- a/pkg/chat/attach.go +++ b/pkg/chat/attach.go @@ -17,128 +17,17 @@ import ( "encoding/base64" "errors" "fmt" - "io" - "mime" - "net" - "net/http" "os" "path/filepath" "strings" - "time" ) const ( - // attachHTTPTimeout is the maximum time allowed to fetch a remote image URL. - attachHTTPTimeout = 10 * time.Second - - // attachMaxRemoteBytes is the maximum number of bytes read from a remote URL. - attachMaxRemoteBytes = 20 * 1024 * 1024 // 20 MB - - // MaxInlineBinarySize is the maximum byte size of a binary file (PDF, etc.) - // that will be read into memory for inline attachment. Matches the remote - // fetch cap so local and remote paths behave consistently. + // MaxInlineBinarySize is the maximum byte size of a binary file (PDF, image, + // etc.) that will be read into memory for inline attachment. MaxInlineBinarySize = 20 * 1024 * 1024 // 20 MB ) -// privateIPNets lists address ranges that must not be dialled by the -// attach-time URL fetcher. Blocking these prevents SSRF attacks against -// cloud metadata services, internal APIs, and loopback services. -var privateIPNets = func() []*net.IPNet { - blocks := []string{ - // Loopback - "127.0.0.0/8", - "::1/128", - // Link-local — covers AWS/GCP/Azure metadata endpoints (169.254.169.254) - "169.254.0.0/16", - "fe80::/10", - // RFC-1918 private ranges - "10.0.0.0/8", - "172.16.0.0/12", - "192.168.0.0/16", - // IPv6 unique-local - "fc00::/7", - } - nets := make([]*net.IPNet, 0, len(blocks)) - for _, b := range blocks { - _, ipNet, err := net.ParseCIDR(b) - if err != nil { - panic("attachment: invalid built-in CIDR " + b + ": " + err.Error()) - } - nets = append(nets, ipNet) - } - return nets -}() - -// isPrivateIP reports whether ip falls in any of the blocked address ranges. -func isPrivateIP(ip net.IP) bool { - for _, block := range privateIPNets { - if block.Contains(ip) { - return true - } - } - return false -} - -// checkSafeHostCtx resolves host to IP addresses and returns an error if any -// resolved address is in a private/reserved range. This is called both at -// dial time (where a context is available) and on redirect destinations. -func checkSafeHostCtx(ctx context.Context, host string) error { - addrs, err := net.DefaultResolver.LookupHost(ctx, host) - if err != nil { - return fmt.Errorf("attachment: cannot resolve host %q: %w", host, err) - } - for _, addr := range addrs { - ip := net.ParseIP(addr) - if ip == nil { - continue - } - if isPrivateIP(ip) { - return fmt.Errorf("attachment: URL resolves to private/reserved address %s (SSRF protection)", addr) - } - } - return nil -} - -// safeHTTPClient is a shared HTTP client used by fetchRemoteImage. -// It enforces SSRF protection by refusing connections to private and -// reserved IP ranges at both dial time and on each redirect hop. -// -// Tests may override this variable to inject a custom client that bypasses -// SSRF filtering (e.g. to reach httptest.Server on loopback). Only do this -// in test code guarded by a t.Cleanup restore. -var safeHTTPClient = newSafeHTTPClient() - -func newSafeHTTPClient() *http.Client { - return &http.Client{ - Timeout: attachHTTPTimeout, - Transport: &http.Transport{ - DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { - host, _, err := net.SplitHostPort(addr) - if err != nil { - return nil, fmt.Errorf("attachment: malformed dial address %q: %w", addr, err) - } - if err := checkSafeHostCtx(ctx, host); err != nil { - return nil, err - } - return (&net.Dialer{}).DialContext(ctx, network, addr) - }, - }, - CheckRedirect: func(req *http.Request, _ []*http.Request) error { - return checkSafeHostCtx(req.Context(), req.URL.Hostname()) - }, - } -} - -// SetFetchHTTPClientForTest replaces the HTTP client used by fetchRemoteImage -// and returns a restore function. Only for use in tests. -// -// defer chat.SetFetchHTTPClientForTest(t, myClient)() -func SetFetchHTTPClientForTest(client *http.Client) (restore func()) { - prev := safeHTTPClient - safeHTTPClient = client - return func() { safeHTTPClient = prev } -} - // ProcessAttachment converts a raw [MessagePart] into a [Document] with fully // resolved Source.InlineData or Source.InlineText. It is called once when a // message is assembled — never at inference time. @@ -150,19 +39,16 @@ func SetFetchHTTPClientForTest(client *http.Client) (restore func()) { // binary bytes (images are transcoded+resized via [ResizeImage]; PDFs and // other supported types are read verbatim). // -// - [MessagePartTypeImageURL]: handles data: URIs (decoded inline) and -// http(s):// URLs (fetched with a 10-second timeout). The image bytes are -// then passed through [ResizeImage] to normalise to JPEG or PNG. +// - [MessagePartTypeImageURL]: handles data: URIs (decoded inline). Remote +// http(s):// URLs are not supported; callers should download the file +// locally first and pass it as a [MessagePartTypeFile] instead. // // - [MessagePartTypeDocument]: if Source.InlineData or Source.InlineText is // already set, the document is returned as-is after applying image // transcoding to any image/* InlineData. A Document with no inline content // is an error. -// -// The context is forwarded to any network operations; filesystem and image -// decoding operations are not yet context-aware. -func ProcessAttachment(ctx context.Context, part MessagePart) (Document, error) { - doc, _, err := ProcessAttachmentWithMetadata(ctx, part) +func ProcessAttachment(_ context.Context, part MessagePart) (Document, error) { + doc, _, err := ProcessAttachmentWithMetadata(part) return doc, err } @@ -173,12 +59,12 @@ func ProcessAttachment(ctx context.Context, part MessagePart) (Document, error) // Callers that need to emit a dimension note (for model coordinate-mapping) // should use this variant and call [FormatDimensionNote] on the returned // metadata. -func ProcessAttachmentWithMetadata(ctx context.Context, part MessagePart) (Document, *ImageResizeResult, error) { +func ProcessAttachmentWithMetadata(part MessagePart) (Document, *ImageResizeResult, error) { switch part.Type { case MessagePartTypeFile: return processFilePart(part) case MessagePartTypeImageURL: - return processImageURLPart(ctx, part) + return processImageURLPart(part) case MessagePartTypeDocument: return processDocumentPart(part) default: @@ -271,48 +157,34 @@ func processFilePart(part MessagePart) (Document, *ImageResizeResult, error) { } // processImageURLPart handles MessagePartTypeImageURL. -// Supports: -// - data: URIs (data:;base64,) -// - http:// and https:// URLs (fetched with attachHTTPTimeout) -func processImageURLPart(ctx context.Context, part MessagePart) (Document, *ImageResizeResult, error) { +// Only data: URIs are supported; remote http(s):// URLs are rejected. +// Callers with a remote URL should download the file locally first and +// pass it as a MessagePartTypeFile instead. +func processImageURLPart(part MessagePart) (Document, *ImageResizeResult, error) { if part.ImageURL == nil { return Document{}, nil, errors.New("ProcessAttachment: image-url part has nil ImageURL field") } rawURL := part.ImageURL.URL - var ( - data []byte - mimeType string - name string - err error - ) - switch { case strings.HasPrefix(rawURL, "data:"): - mimeType, data, err = parseDataURI(rawURL) + mimeType, data, err := parseDataURI(rawURL) if err != nil { return Document{}, nil, fmt.Errorf("ProcessAttachment: parse data URI: %w", err) } - name = "image" + // When content detection returns an image type, prefer it over the + // declared MIME. (Only image types are trusted from the sniffer.) + if detected := DetectMimeTypeByContent(data); IsImageMimeType(detected) { + mimeType = detected + } + return transcodeImageWithMeta("image", data, mimeType) case strings.HasPrefix(rawURL, "http://") || strings.HasPrefix(rawURL, "https://"): - data, mimeType, name, err = fetchRemoteImage(ctx, rawURL) - if err != nil { - return Document{}, nil, fmt.Errorf("ProcessAttachment: fetch remote image %q: %w", rawURL, err) - } + return Document{}, nil, errors.New("attachment: remote URLs are not supported; download the file locally first") default: - return Document{}, nil, fmt.Errorf("ProcessAttachment: unsupported image URL scheme (must be data: or http(s)://): %q", rawURL) + return Document{}, nil, fmt.Errorf("attachment: unsupported image URL scheme: %q", rawURL) } - - // When content detection returns an image type, prefer it over whatever - // the URI or Content-Type header said. (Only image types are trusted from - // the sniffer to avoid promoting an unknown binary to a valid MIME.) - if detected := DetectMimeTypeByContent(data); IsImageMimeType(detected) { - mimeType = detected - } - - return transcodeImageWithMeta(name, data, mimeType) } // processDocumentPart handles MessagePartTypeDocument. @@ -325,15 +197,12 @@ func processDocumentPart(part MessagePart) (Document, *ImageResizeResult, error) if len(doc.Source.InlineData) > 0 { if IsImageMimeType(doc.MimeType) { - // Apply transcoding/resizing to normalise the image. return transcodeImageWithMeta(doc.Name, doc.Source.InlineData, doc.MimeType) } - // Non-image binary — pass through unchanged. return doc, nil, nil } if doc.Source.InlineText != "" { - // Already resolved text — pass through unchanged. return doc, nil, nil } @@ -348,13 +217,12 @@ func transcodeImageWithMeta(name string, data []byte, mimeType string) (Document if err != nil { return Document{}, nil, fmt.Errorf("ProcessAttachment: transcode image %q: %w", name, err) } - doc := Document{ + return Document{ Name: name, MimeType: result.MimeType, Size: int64(len(result.Data)), Source: DocumentSource{InlineData: result.Data}, - } - return doc, result, nil + }, result, nil } // parseDataURI parses a data URI of the form "data:;base64,". @@ -391,49 +259,3 @@ func parseDataURI(uri string) (mimeType string, data []byte, err error) { } return mimeType, data, nil } - -// fetchRemoteImage fetches an image from an http(s) URL. -// The safeHTTPClient enforces a timeout and blocks connections to -// private/reserved IP ranges (SSRF protection). -// Returns the image bytes, detected MIME type, and a filename derived from the URL. -func fetchRemoteImage(ctx context.Context, rawURL string) (data []byte, mimeType, name string, err error) { - req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, http.NoBody) - if err != nil { - return nil, "", "", fmt.Errorf("create request: %w", err) - } - - resp, err := safeHTTPClient.Do(req) - if err != nil { - return nil, "", "", fmt.Errorf("HTTP GET: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return nil, "", "", fmt.Errorf("HTTP %d for %q", resp.StatusCode, rawURL) - } - - data, err = io.ReadAll(io.LimitReader(resp.Body, attachMaxRemoteBytes+1)) - if err != nil { - return nil, "", "", fmt.Errorf("read response body: %w", err) - } - if int64(len(data)) > attachMaxRemoteBytes { - return nil, "", "", fmt.Errorf("remote image too large (max %d bytes)", attachMaxRemoteBytes) - } - - // Determine MIME type: prefer Content-Type header, fall back to content sniff. - mimeType = resp.Header.Get("Content-Type") - if ct, _, parseErr := mime.ParseMediaType(mimeType); parseErr == nil { - mimeType = ct - } - if mimeType == "" || mimeType == "application/octet-stream" { - mimeType = DetectMimeTypeByContent(data) - } - - // Derive a filename from the URL path. - name = filepath.Base(req.URL.Path) - if name == "" || name == "." || name == "/" { - name = "image" - } - - return data, mimeType, name, nil -} diff --git a/pkg/chat/attach_test.go b/pkg/chat/attach_test.go index c2f493003..f2a484b95 100644 --- a/pkg/chat/attach_test.go +++ b/pkg/chat/attach_test.go @@ -7,8 +7,6 @@ import ( "image/color" "image/jpeg" "image/png" - "net/http" - "net/http/httptest" "os" "path/filepath" "strings" @@ -24,7 +22,6 @@ import ( // Helpers // ────────────────────────────────────────────────────────────────────────────── -// encodeJPEGBytes returns a minimal JPEG image as raw bytes. func encodeJPEGBytes(w, h int) []byte { img := image.NewRGBA(image.Rect(0, 0, w, h)) for y := range h { @@ -39,7 +36,6 @@ func encodeJPEGBytes(w, h int) []byte { return buf.Bytes() } -// encodePNGBytes returns a minimal PNG image as raw bytes. func encodePNGBytes(w, h int, alpha bool) []byte { if alpha { img := image.NewNRGBA(image.Rect(0, 0, w, h)) @@ -67,7 +63,6 @@ func encodePNGBytes(w, h int, alpha bool) []byte { return buf.Bytes() } -// writeTempFile writes data to a temp file with the given extension and returns its path. func writeTempFile(t *testing.T, ext string, data []byte) string { t.Helper() f, err := os.CreateTemp(t.TempDir(), "attach-*"+ext) @@ -111,7 +106,6 @@ func TestProcessAttachment_PNG_Passthrough(t *testing.T) { } func TestProcessAttachment_PNG_WithAlpha_StaysPNG(t *testing.T) { - // A PNG with alpha channel must not be converted to JPEG (lossy strips alpha). data := encodePNGBytes(100, 100, true) path := writeTempFile(t, ".png", data) @@ -120,15 +114,12 @@ func TestProcessAttachment_PNG_WithAlpha_StaysPNG(t *testing.T) { File: &chat.MessageFile{Path: path, MimeType: "image/png"}, }) require.NoError(t, err) - // ResizeImage picks the smaller of PNG/JPEG; for small images PNG is usually smaller - // but what matters is that the output is a valid image. assert.True(t, doc.MimeType == "image/png" || doc.MimeType == "image/jpeg", "expected png or jpeg, got %q", doc.MimeType) assert.NotEmpty(t, doc.Source.InlineData) } func TestProcessAttachment_ImageTooLarge_Resized(t *testing.T) { - // An image larger than MaxImageDimension should be resized. bigData := encodeJPEGBytes(chat.MaxImageDimension+200, chat.MaxImageDimension+200) path := writeTempFile(t, ".jpg", bigData) @@ -139,7 +130,6 @@ func TestProcessAttachment_ImageTooLarge_Resized(t *testing.T) { require.NoError(t, err) assert.NotEmpty(t, doc.Source.InlineData) - // Verify the output image dimensions fit within limits. img, _, decErr := image.Decode(bytes.NewReader(doc.Source.InlineData)) require.NoError(t, decErr) b := img.Bounds() @@ -162,8 +152,7 @@ func TestProcessAttachment_PDF_Passthrough(t *testing.T) { } func TestProcessAttachment_BinaryFileTooLarge_Error(t *testing.T) { - // Write a file whose Stat.Size exceeds MaxInlineBinarySize. - // We use a sparse file (truncate to the target size) so the test is fast. + // Sparse file: Stat.Size > MaxInlineBinarySize without allocating memory. path := writeTempFile(t, ".pdf", nil) f, err := os.OpenFile(path, os.O_WRONLY, 0o600) require.NoError(t, err) @@ -189,7 +178,6 @@ func TestProcessAttachment_TextFile_InlineText(t *testing.T) { require.NoError(t, err) assert.Empty(t, doc.Source.InlineData) assert.Contains(t, doc.Source.InlineText, content) - assert.Contains(t, doc.Source.InlineText, filepath.Base(path)) // ReadFileForInline wraps in a tag } func TestProcessAttachment_MarkdownFile_InlineText(t *testing.T) { @@ -260,6 +248,18 @@ func TestProcessAttachment_DataURI_NonBase64_Error(t *testing.T) { assert.Contains(t, err.Error(), "not base64") } +func TestProcessAttachment_RemoteURL_Error(t *testing.T) { + // Remote http(s):// URLs are not supported; callers must download locally. + for _, scheme := range []string{"http://", "https://"} { + _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ + Type: chat.MessagePartTypeImageURL, + ImageURL: &chat.MessageImageURL{URL: scheme + "example.com/photo.jpg"}, + }) + require.Error(t, err, "expected error for scheme %s", scheme) + assert.Contains(t, err.Error(), "remote URLs are not supported") + } +} + func TestProcessAttachment_UnsupportedScheme_Error(t *testing.T) { _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ Type: chat.MessagePartTypeImageURL, @@ -269,69 +269,6 @@ func TestProcessAttachment_UnsupportedScheme_Error(t *testing.T) { assert.Contains(t, err.Error(), "unsupported image URL scheme") } -func TestProcessAttachment_HTTPS_Image_Via_HTTPTestServer(t *testing.T) { - jpegData := encodeJPEGBytes(60, 60) - - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "image/jpeg") - _, _ = w.Write(jpegData) - })) - t.Cleanup(srv.Close) - - // The httptest.Server binds to 127.0.0.1 (loopback), which is blocked by - // the SSRF filter in production. We use a plain http.Client without the - // SSRF-filtering transport so the test can reach the local server. - restore := chat.SetFetchHTTPClientForTest(srv.Client()) - t.Cleanup(restore) - - doc, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ - Type: chat.MessagePartTypeImageURL, - ImageURL: &chat.MessageImageURL{URL: srv.URL + "/photo.jpg"}, - }) - require.NoError(t, err) - assert.NotEmpty(t, doc.Source.InlineData) - assert.True(t, doc.MimeType == "image/jpeg" || doc.MimeType == "image/png", - "unexpected MIME: %q", doc.MimeType) - assert.Equal(t, "photo.jpg", doc.Name) -} - -func TestProcessAttachment_SSRF_LoopbackBlocked(t *testing.T) { - _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ - Type: chat.MessagePartTypeImageURL, - ImageURL: &chat.MessageImageURL{URL: "http://127.0.0.1:9999/secret"}, - }) - require.Error(t, err) - assert.Contains(t, err.Error(), "private/reserved") -} - -func TestProcessAttachment_SSRF_MetadataEndpointBlocked(t *testing.T) { - // 169.254.169.254 is the AWS/GCP/Azure instance metadata endpoint. - _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ - Type: chat.MessagePartTypeImageURL, - ImageURL: &chat.MessageImageURL{URL: "http://169.254.169.254/latest/meta-data/"}, - }) - require.Error(t, err) - assert.Contains(t, err.Error(), "private/reserved") -} - -func TestProcessAttachment_HTTP_Non200_Error(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - http.Error(w, "not found", http.StatusNotFound) - })) - t.Cleanup(srv.Close) - - // Bypass SSRF filter so we can reach the loopback-bound httptest server. - restore := chat.SetFetchHTTPClientForTest(srv.Client()) - t.Cleanup(restore) - - _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ - Type: chat.MessagePartTypeImageURL, - ImageURL: &chat.MessageImageURL{URL: srv.URL + "/missing.jpg"}, - }) - require.Error(t, err) - assert.Contains(t, err.Error(), "404") -} - func TestProcessAttachment_NilImageURL_Error(t *testing.T) { _, err := chat.ProcessAttachment(t.Context(), chat.MessagePart{ Type: chat.MessagePartTypeImageURL,