Cyrene_v2/internal/runtime/multimodal/vision/image.go

package vision

import (
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
)

// Size caps for reading image bytes and response bodies, so a hostile or
// malformed source cannot exhaust memory. The 50MB cap matches the file-upload
// limit in DevDocs/02-development-constraints.md §11.
const (
	maxImageBytes  = 50 << 20
	maxFetchedBody = 50 << 20
)

// DataURI builds a base64 data: image URI from mime and raw bytes.
func DataURI(mime string, data []byte) string {
	return "data:" + mime + ";base64," + base64.StdEncoding.EncodeToString(data)
}

// Resolve loads the bytes for an ImageRef. Paths and base64 payloads are read
// locally; URLs are fetched with an http.Client (the provided client, or a
// default one when nil). For base64 input the Data field is returned verbatim as
// the payload and the MIME type is guessed from a data: prefix.
func Resolve(ctx context.Context, ref ImageRef, client *http.Client) (mime string, data []byte, err error) {
	switch ref.Source {
	case SourcePath:
		return readFile(ref.Data)
	case SourceBase64:
		return decodeBase64(ref.Data)
	case SourceURL:
		return fetchURL(ctx, ref.Data, client)
	default:
		return "", nil, fmt.Errorf("vision: %w: %q", ErrUnsupportedSource, ref.Source)
	}
}

// readFile reads an image file and guesses its MIME type from the extension.
func readFile(path string) (string, []byte, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", nil, fmt.Errorf("vision: %w: %v", ErrBadImage, err)
	}
	defer func() { _ = f.Close() }()
	data, err := io.ReadAll(io.LimitReader(f, maxImageBytes))
	if err != nil {
		return "", nil, fmt.Errorf("vision: %w: %v", ErrBadImage, err)
	}
	return mimeByExt(path), data, nil
}

// decodeBase64 returns the decoded bytes of a base64/data: payload and its MIME.
func decodeBase64(src string) (string, []byte, error) {
	mime, encoded, isDataURI := splitDataURI(src)
	if isDataURI {
		decoded, err := base64.StdEncoding.DecodeString(encoded)
		if err != nil {
			return "", nil, fmt.Errorf("vision: %w: %v", ErrBadImage, err)
		}
		return mime, decoded, nil
	}
	decoded, err := base64.StdEncoding.DecodeString(src)
	if err != nil {
		return "", nil, fmt.Errorf("vision: %w: %v", ErrBadImage, err)
	}
	return "image/png", decoded, nil
}

// fetchURL downloads an image over http(s); only http/https schemes are allowed.
func fetchURL(ctx context.Context, rawURL string, client *http.Client) (string, []byte, error) {
	if !strings.HasPrefix(rawURL, "http://") && !strings.HasPrefix(rawURL, "https://") {
		return "", nil, fmt.Errorf("vision: %w: %q", ErrUnsupportedSource, rawURL)
	}
	if client == nil {
		client = defaultHTTPClient()
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		return "", nil, fmt.Errorf("vision: %w: %v", ErrFetch, err)
	}
	resp, err := client.Do(req)
	if err != nil {
		return "", nil, fmt.Errorf("vision: %w: %v", ErrFetch, err)
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		return "", nil, fmt.Errorf("vision: %w: status %d", ErrFetch, resp.StatusCode)
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, maxFetchedBody))
	if err != nil {
		return "", nil, fmt.Errorf("vision: %w: %v", ErrFetch, err)
	}
	return mimeByContentType(resp.Header.Get("Content-Type")), data, nil
}

// splitDataURI parses a "data:<mime>;base64,<payload>" prefix. The second return
// is false when src is not a data: URI.
func splitDataURI(src string) (mime, payload string, ok bool) {
	if !strings.HasPrefix(src, "data:") {
		return "", "", false
	}
	rest := strings.TrimPrefix(src, "data:")
	comma := strings.IndexByte(rest, ',')
	if comma < 0 {
		return "", "", false
	}
	meta := rest[:comma]
	payload = rest[comma+1:]
	mime = "image/png"
	if strings.HasPrefix(meta, ";base64") {
		return mime, payload, true
	}
	if idx := strings.Index(meta, ";base64"); idx >= 0 {
		if t := strings.TrimSuffix(meta[:idx], ";base64"); t != "" {
			mime = t
		}
		return mime, payload, true
	}
	return mime, payload, true
}

// mimeByExt guesses a MIME type from a file path extension.
func mimeByExt(path string) string {
	switch strings.ToLower(pathExtension(path)) {
	case ".png":
		return "image/png"
	case ".jpg", ".jpeg":
		return "image/jpeg"
	case ".gif":
		return "image/gif"
	case ".webp":
		return "image/webp"
	default:
		return "image/png"
	}
}

// mimeByContentType returns a MIME type from an HTTP Content-Type header,
// defaulting to image/png when absent or non-image.
func mimeByContentType(ct string) string {
	if ct == "" || !strings.HasPrefix(ct, "image/") {
		return "image/png"
	}
	return strings.TrimSpace(strings.Split(ct, ";")[0])
}

// pathExtension returns the file extension with its dot, or empty when none.
func pathExtension(path string) string {
	base := path
	if idx := strings.LastIndexByte(path, '/'); idx >= 0 {
		base = path[idx+1:]
	} else if idx := strings.LastIndexByte(path, '\\'); idx >= 0 {
		base = path[idx+1:]
	}
	dot := strings.LastIndexByte(base, '.')
	if dot < 0 {
		return ""
	}
	return base[dot:]
}

// fileExists reports whether path exists on disk and is a regular file.
func fileExists(path string) bool {
	info, err := os.Stat(path)
	return err == nil && info.Mode().IsRegular()
}

// defaultHTTPClient builds the HTTP client described in
// DevDocs/02-development-constraints.md §2.3. A fresh client is returned per
// call rather than a package global, which the constraints forbid.
func defaultHTTPClient() *http.Client {
	return &http.Client{Transport: http.DefaultTransport}
}