package vision import ( "context" "encoding/base64" "fmt" "io" "net/http" "os" "strings" ) // 读取图像字节与响应体的体积上限,使恶意或 // 畸形来源无法耗尽内存。50MB 上限与 // DevDocs/02-development-constraints.md §11 的文件上传限制一致。 const ( maxImageBytes = 50 << 20 maxFetchedBody = 50 << 20 ) // 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 类型。 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 返回 base64/data: 载荷解码后的字节及其 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 通过 http(s) 下载一张图像;只允许 http/https 协议。 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 解析 "data:;base64," 前缀。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 } // mimeByExt 从文件路径扩展名推断 MIME 类型。 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 从 HTTP Content-Type 头返回 MIME 类型, // 缺失或非图像时默认为 image/png。 func mimeByContentType(ct string) string { if ct == "" || !strings.HasPrefix(ct, "image/") { return "image/png" } return strings.TrimSpace(strings.Split(ct, ";")[0]) } // pathExtension 返回带点号的文件扩展名,没有则返回空串。 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 报告 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} }