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

package vision

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

// 读取图像字节与响应体的体积上限，使恶意或
// 畸形来源无法耗尽内存。50MB 上限与
// DevDocs/02-development-constraints.md §11 的文件上传限制一致。
const (
	maxImageBytes  = 50 << 20
	maxFetchedBody = 50 << 20
)

// imageFetchTimeout 限制一次图像 URL 取回的时长。适配器与消息解析器会
// 为模型代取 URL，没有上限时一个卡住或恶意的地址会一直占着这一轮
// （fetchURL 经 guardedImageClient 给未设超时的 client 补上它，
// 拨号守卫也用同一个值作单次建连上限）。
const imageFetchTimeout = 30 * time.Second

// DataURI 由 mime 与原始字节构造一个 base64 data: 图像 URI。
func DataURI(mime string, data []byte) string {
	return "data:" + mime + ";base64," + base64.StdEncoding.EncodeToString(data)
}

// Resolve 加载一个 ImageRef 对应的字节。路径与 base64 载荷在
// 本地读取；URL 用 http.Client 抓取（传入的 client，为 nil 时
// 用默认 client）。对 base64 输入，Data 字段原样作为
// 载荷返回，MIME 类型从 data: 前缀推断。
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 读取一个图像文件。扩展名给出声明的 MIME，真实字节嗅探是否决权
// （见 imageMIMEFor）：改名的 HTML 不会被当成图片读进模型。
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)
	}
	mime, err := imageMIMEFor(mimeByExt(path), data)
	if err != nil {
		return "", nil, err
	}
	return mime, data, nil
}

// decodeBase64 返回 base64/data: 载荷解码后的字节及其 MIME。
// 长度预检先于解码与 data: 解析（见 checkInlineImageLength）：超限载荷在
// 分配之前就被拦下，而不是先按约 4/3 分配再报错。
func decodeBase64(src string) (string, []byte, error) {
	if err := checkInlineImageLength(src); err != nil {
		return "", nil, err
	}
	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 通过 http(s) 下载一张图像；只允许 http/https 协议。
// 取回前、重定向每一跳与真正建连前都有守卫（见 guard.go）：链路本地与
// 云元数据地址被拒，环回/私网/公网照旧可达（同源自取需要它们）。
// 声明的 Content-Type 只作参考，真实字节嗅探说了算（见 imageMIMEFor）：
// 网关/PROXY 返回的 HTML 错误页绝不会被当成 png 送进模型。
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)
	}
	u, err := url.Parse(rawURL)
	if err != nil {
		// 不回显 URL：解析失败的原文里带着整条地址（可能含凭据）。
		return "", nil, fmt.Errorf("vision: %w: URL 无法解析", ErrFetch)
	}
	if err := guardImageURL(u); err != nil {
		return "", nil, err
	}
	guarded := guardedImageClient(client)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		return "", nil, fmt.Errorf("vision: %w: 无法构造请求", ErrFetch)
	}
	resp, err := guarded.Do(req)
	if err != nil {
		// 被拒地址直接返回类型化错误：它只含主机与网段原因，
		// 不像 client.Do 的 *url.Error 那样把整条 URL（可能含凭据）带出来。
		if blocked, ok := blockedAddressOf(err); ok {
			return "", nil, blocked
		}
		return "", nil, fmt.Errorf("vision: %w: %w", ErrFetch, fetchCause(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)
	}
	mime, err := imageMIMEFor(resp.Header.Get("Content-Type"), data)
	if err != nil {
		return "", nil, err
	}
	return mime, data, nil
}

// inlineImageSrc 把图像来源字符串取回并内联为 data: URI，供主模型
// 直接消费。已是 data: URI 的原样返回（客户端现有的 base64 路径因此
// 逐字节不变），但**短路之前**也要过长度预检：这条路径不解码，不预检
// 就等于把体积上限完全交给入口层。URL 取回复用 Resolve→fetchURL 的既有
// 守卫（协议白名单 + 体积上限 + 地址守卫），并加取回超时。
func inlineImageSrc(ctx context.Context, src string) (string, error) {
	if strings.HasPrefix(src, "data:") {
		if err := checkInlineImageLength(src); err != nil {
			return "", err
		}
		return src, nil
	}
	ref, err := ParseImageRef(src)
	if err != nil {
		return "", err
	}
	if ref.Source == SourceURL {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, imageFetchTimeout)
		defer cancel()
	}
	mime, data, err := Resolve(ctx, ref, nil)
	if err != nil {
		return "", err
	}
	return DataURI(mime, data), nil
}

// splitDataURI 解析 "data:<mime>;base64,<payload>" 前缀。src 不是
// data: URI 时第二个返回值是 false。
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
}

// fileExists 报告 path 是否存在于磁盘且为常规文件。
func fileExists(path string) bool {
	info, err := os.Stat(path)
	return err == nil && info.Mode().IsRegular()
}

// defaultHTTPClient 构造 DevDocs/02-development-constraints.md §2.3
// 所述的 HTTP client。这里每次调用返回一个新 client，
// 而不是包级全局变量——后者为约束所禁止。
func defaultHTTPClient() *http.Client {
	return &http.Client{Transport: http.DefaultTransport}
}