Files
NianAIGC/backend/internal/assets/import_image.go
2026-09-30 13:22:53 +08:00

226 lines
6.8 KiB
Go

package assets
import (
"bytes"
"context"
"encoding/binary"
"errors"
"image"
_ "image/gif"
_ "image/jpeg"
_ "image/png"
"io"
"net/url"
"strings"
)
var (
ErrImportImageURL = errors.New("invalid image URL")
ErrImportImageUnavailable = errors.New("image URL is unavailable")
ErrImportImageUnsupported = errors.New("unsupported image format")
)
const (
maxImportImageURLBytes = 8192
maxImportImageBytes = 20 << 20
maxImportImagePixels = 25_000_000
)
// ImportImage fetches an untrusted image into the configured blob store using
// the same write ordering and catalog compensation as a regular upload.
func (s *Service) ImportImage(ctx context.Context, scope Scope, rawURL string) (Asset, error) {
if err := validScope(scope); err != nil {
return Asset{}, err
}
if scope.kind != platformScope {
return Asset{}, ErrImportImageURL
}
if len(rawURL) == 0 || len(rawURL) > maxImportImageURLBytes || strings.TrimSpace(rawURL) != rawURL {
return Asset{}, ErrImportImageURL
}
u, err := url.Parse(rawURL)
if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil {
return Asset{}, ErrImportImageURL
}
if s.remote == nil {
return Asset{}, ErrImportImageUnavailable
}
blob, err := s.remote.Fetch(ctx, rawURL)
if err != nil {
if errors.Is(err, ErrRemoteTooLarge) {
return Asset{}, ErrRemoteTooLarge
}
return Asset{}, ErrImportImageUnavailable
}
if blob.Body == nil {
return Asset{}, ErrImportImageUnavailable
}
defer blob.Body.Close()
if blob.Size > maxImportImageBytes {
return Asset{}, ErrRemoteTooLarge
}
content, err := io.ReadAll(io.LimitReader(blob.Body, maxImportImageBytes+1))
if err != nil {
return Asset{}, ErrImportImageUnavailable
}
if len(content) > maxImportImageBytes {
return Asset{}, ErrRemoteTooLarge
}
extension, mime, ok := inspectImportedImage(content)
if !ok {
return Asset{}, ErrImportImageUnsupported
}
return s.Upload(ctx, scope, UploadCommand{Bytes: content, FileName: "导入图片" + extension, ContentType: mime, Kind: KindImage})
}
func inspectImportedImage(content []byte) (string, string, bool) {
if len(content) < 12 {
return "", "", false
}
if bytes.Equal(content[:4], []byte("RIFF")) && bytes.Equal(content[8:12], []byte("WEBP")) {
if validWebP(content) {
return ".webp", "image/webp", true
}
return "", "", false
}
if content[0] == 'B' && content[1] == 'M' {
if validBMP(content) {
return ".bmp", "image/bmp", true
}
return "", "", false
}
config, format, err := image.DecodeConfig(bytes.NewReader(content))
if err != nil || !validImageDimensions(config.Width, config.Height) {
return "", "", false
}
// Decode the first frame so that a forged or truncated header cannot be
// stored as a supported image. DecodeConfig alone checks only the header.
if _, _, err = image.Decode(bytes.NewReader(content)); err != nil {
return "", "", false
}
switch format {
case "png":
return ".png", "image/png", true
case "jpeg":
return ".jpg", "image/jpeg", true
case "gif":
return ".gif", "image/gif", true
default:
return "", "", false
}
}
func validImageDimensions(width, height int) bool {
return width > 0 && height > 0 && int64(width)*int64(height) <= maxImportImagePixels
}
// The standard library does not decode WebP or BMP. Validate their containers,
// image headers and dimensions before accepting the original raster bytes.
func validWebP(b []byte) bool {
if len(b) < 30 || uint64(binary.LittleEndian.Uint32(b[4:8]))+8 != uint64(len(b)) {
return false
}
first, payload, next, ok := webPChunk(b, 12)
if !ok {
return false
}
if first == "VP8 " || first == "VP8L" {
_, _, valid := webPImageDimensions(first, payload)
return valid && next == len(b)
}
if first != "VP8X" || len(payload) != 10 || payload[0]&^byte(0x3e) != 0 || payload[1] != 0 || payload[2] != 0 || payload[3] != 0 {
return false
}
// Animation requires ANIM/ANMF frame handling; reject it in this version.
if payload[0]&0x02 != 0 {
return false
}
canvasWidth := 1 + int(payload[4]) + (int(payload[5]) << 8) + (int(payload[6]) << 16)
canvasHeight := 1 + int(payload[7]) + (int(payload[8]) << 8) + (int(payload[9]) << 16)
if !validImageDimensions(canvasWidth, canvasHeight) {
return false
}
images := 0
for next < len(b) {
chunk, data, after, ok := webPChunk(b, next)
if !ok {
return false
}
switch chunk {
case "VP8 ", "VP8L":
width, height, valid := webPImageDimensions(chunk, data)
if !valid || width != canvasWidth || height != canvasHeight || images != 0 {
return false
}
images++
case "ALPH", "ICCP", "EXIF", "XMP ":
if len(data) == 0 {
return false
}
default:
return false
}
next = after
}
return images == 1
}
func webPChunk(b []byte, offset int) (string, []byte, int, bool) {
if offset > len(b)-8 {
return "", nil, 0, false
}
size := uint64(binary.LittleEndian.Uint32(b[offset+4 : offset+8]))
end := uint64(offset) + 8 + size + (size & 1)
if size == 0 || end > uint64(len(b)) {
return "", nil, 0, false
}
return string(b[offset : offset+4]), b[offset+8 : offset+8+int(size)], int(end), true
}
func webPImageDimensions(chunk string, payload []byte) (int, int, bool) {
switch chunk {
case "VP8 ":
if len(payload) < 10 || payload[0]&1 != 0 || !bytes.Equal(payload[3:6], []byte{0x9d, 0x01, 0x2a}) {
return 0, 0, false
}
partitionBytes := int(payload[0]) | int(payload[1])<<8 | int(payload[2])<<16
if partitionBytes>>5 > len(payload)-10 {
return 0, 0, false
}
width := int(binary.LittleEndian.Uint16(payload[6:8]) & 0x3fff)
height := int(binary.LittleEndian.Uint16(payload[8:10]) & 0x3fff)
return width, height, validImageDimensions(width, height)
case "VP8L":
if len(payload) <= 5 || payload[0] != 0x2f || payload[4]&0xe0 != 0 {
return 0, 0, false
}
width := 1 + int(payload[1]) + (int(payload[2]&0x3f) << 8)
height := 1 + int(payload[2]>>6) + (int(payload[3]) << 2) + (int(payload[4]&0x0f) << 10)
return width, height, validImageDimensions(width, height)
}
return 0, 0, false
}
func validBMP(b []byte) bool {
if len(b) < 54 || int(binary.LittleEndian.Uint32(b[2:6])) != len(b) {
return false
}
offset := int(binary.LittleEndian.Uint32(b[10:14]))
header := binary.LittleEndian.Uint32(b[14:18])
if header < 40 || int(header) > len(b)-14 || offset < 14+int(header) || offset >= len(b) {
return false
}
width := int(int32(binary.LittleEndian.Uint32(b[18:22])))
height := int(int32(binary.LittleEndian.Uint32(b[22:26])))
if height < 0 {
height = -height
}
depth := binary.LittleEndian.Uint16(b[28:30])
compression := binary.LittleEndian.Uint32(b[30:34])
if !validImageDimensions(width, height) || binary.LittleEndian.Uint16(b[26:28]) != 1 || compression != 0 || (depth != 24 && depth != 32) {
return false
}
stride := (int64(width)*int64(depth) + 31) / 32 * 4
return stride*int64(height) <= int64(len(b)-offset)
}