tamabox/internal/storage/s3.go

// Copyright 2024 E99p1ant. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package storage

import (
	"bytes"
	"context"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"mime/multipart"
	"os"
	"os/exec"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/chai2010/webp"
	"github.com/disintegration/imaging"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"

	"github.com/tamakyi/TamaBox/internal/conf"
	"github.com/tamakyi/TamaBox/internal/imagecheck"
)

// MaxUploadBytes 是单次图片上传的内存读取上限（8MB）。
// 调用方虽在 HTTP 层对 fileHeader.Size 做了校验，但该值可被客户端伪造，
// 因此在读取前用 io.LimitReader 兜底，避免超大文件将内存撑爆（DoS）。
const MaxUploadBytes = 8 * 1024 * 1024

// aspect 约束一次上传的宽高比（wantW:wantH，容差 DefaultAspectTolerance）。
// 头像固定 1:1、页面背景固定 16:9：客户端裁剪后即为精确比例，
// 服务端复校验用于拦截绕过前端直接 POST 的越界图片。
type aspect struct {
	W, H float64
}

// check 校验给定尺寸是否满足本用途的宽高比约束。
// orientation 是 EXIF Orientation：手机竖拍照片解出的像素是横躺的（5/6/7/8 会
// 交换宽高），若直接用原始像素宽高比对，合规的竖图会被误判为比例不符而拒。
// 修正只交换宽高、不解码像素，大图内存不受影响。
func (a *aspect) check(cfg image.Config, orientation int) error {
	if a == nil {
		return nil
	}
	w, h := cfg.Width, cfg.Height
	if orientationSwapsDimensions(orientation) {
		w, h = h, w
	}
	return imagecheck.CheckAspectRatioSize(w, h, a.W, a.H, imagecheck.DefaultAspectTolerance)
}

// EncodeSpec 描述一次上传的 WebP 编码档位。
//
// 此前所有用途共用「最长边 1920 + q80」，头像这种实际只展示 120px 的用途
// 也要按 1920 存盘，实测体积差约 7 倍（83KB vs 11KB）。故按用途分档：
// 头像跟随前端 cropper 的输出上限，背景铺满宽度，通用图取折中值。
type EncodeSpec struct {
	// MaxSide 是允许的最大边长（宽高均不超过），<=0 表示不缩放。
	MaxSide int
	// Quality 是有损质量（0~100）。
	Quality float32
}

var (
	// avatarEncodeSpec 头像 1:1：模板中最大展示 120px，
	// 前端 cropper 输出上限为 512，服务端取同值兜底直传的大图。
	avatarEncodeSpec = EncodeSpec{MaxSide: 512, Quality: 78}
	// backgroundEncodeSpec 页面背景 16:9：作为整页头部铺满宽度，保留 1920。
	backgroundEncodeSpec = EncodeSpec{MaxSide: 1920, Quality: 80}
	// GeneralEncodeSpec 通用图（提问配图、私聊图片）：正文内展示，取 1600 折中。
	GeneralEncodeSpec = EncodeSpec{MaxSide: 1600, Quality: 80}
)

// UploadAvatarToS3 上传头像：除通用校验/转换外强制 1:1 宽高比。
func UploadAvatarToS3(file multipart.File, fileHeader *multipart.FileHeader) (string, error) {
	return uploadPictureToS3(file, fileHeader, &aspect{W: 1, H: 1}, avatarEncodeSpec)
}

// UploadBackgroundToS3 上传页面背景：除通用校验/转换外强制 16:9 宽高比。
func UploadBackgroundToS3(file multipart.File, fileHeader *multipart.FileHeader) (string, error) {
	return uploadPictureToS3(file, fileHeader, &aspect{W: 16, H: 9}, backgroundEncodeSpec)
}

// UploadPictureToS3 上传通用图片到 Cloudflare R2，自动转换为 WebP（动画 GIF 可选转换）
func UploadPictureToS3(file multipart.File, fileHeader *multipart.FileHeader) (string, error) {
	return uploadPictureToS3(file, fileHeader, nil, GeneralEncodeSpec)
}

// uploadPictureToS3 是上传管线实现：读入 → 魔数 → （可选）比例校验 → 像素总量校验
// → 解码 → WebP 转换 → 落盘 R2。avatar/background 的固定比例校验与各自的
// 编码档位复用同一管线，仅通过 asp/spec 参数区分，避免为一种用途复制一套上传代码。
func uploadPictureToS3(file multipart.File, fileHeader *multipart.FileHeader, asp *aspect, spec EncodeSpec) (string, error) {
	// 1. 限制读取大小，防止超大上传撑爆内存（DoS）。
	//    先用 LimitReader 截断到 MaxUploadBytes+1 字节；若实际读到的字节数超过上限则直接拒绝。
	limited := io.LimitReader(file, MaxUploadBytes+1)
	fileData, err := io.ReadAll(limited)
	if err != nil {
		return "", errors.Wrap(err, "read file data")
	}
	if int64(len(fileData)) > MaxUploadBytes {
		return "", errors.New("uploaded file is too large")
	}

	// 2. 校验文件魔数，仅允许 JPEG / PNG / GIF / WebP
	if !isSupportedImage(fileData) {
		return "", errors.New("unsupported file type: not a valid image")
	}

	// 3. 只解析文件头：一次拿到格式与宽高，供像素上限校验、比例校验、
	//    WebP 透传判断共用。image.DecodeConfig 不解像素，因此这里不产生解码开销。
	//    （此前为拿 format 先 image.Decode 一遍、转换时再解一遍，同一张图解码两次，
	//    4096×4096 上限下峰值内存翻倍，1C1G 机器并发上传有 OOM 风险。）
	cfg, format, err := image.DecodeConfig(bytes.NewReader(fileData))
	if err != nil {
		return "", errors.Wrap(err, "image decode failed, file may be corrupted")
	}

	// 3.5 像素总量校验（防解压炸弹），必须在任何完整解码之前。
	if cfg.Width > 0 && cfg.Height > 0 &&
		int64(cfg.Width)*int64(cfg.Height) > imagecheck.MaxImagePixels {
		return "", imagecheck.ErrImageTooLarge
	}

	// 3.6 用途宽高比校验（头像 1:1 / 背景 16:9），绕过前端裁剪的直传在此被拦截。
	//     比例校验必须与转码时的 EXIF 归一保持同一套方向，否则手机竖拍照片
	//     （像素横躺 + Orientation=6）会被误判为比例不符。
	if err := asp.check(cfg, jpegOrientation(fileData)); err != nil {
		return "", err
	}

	var uploadData []byte
	var contentType string

	// 4. 按格式转换或直接使用原图
	switch format {
	case "gif":
		webpData, err := ConvertGifToAnimatedWebP(fileData)
		if err != nil || len(webpData) >= len(fileData) {
			uploadData = fileData // 转换失败、或转换后并不更小 → 回退原始 GIF
			contentType = "image/gif"
		} else {
			uploadData = webpData
			contentType = "image/webp"
		}
	case "webp":
		// 已是 WebP：仅当超尺寸上限时才重编（重编会再丢一代画质），
		// 尺寸合规则原样上传。
		if fitsWithin(cfg.Width, cfg.Height, spec.MaxSide) {
			uploadData = fileData
			contentType = "image/webp"
			break
		}
		webpData, err := ConvertStaticImageToWebP(fileData, spec)
		if err != nil || len(webpData) >= len(fileData) {
			uploadData = fileData
			contentType = mimeTypeFromFormat(format)
		} else {
			uploadData = webpData
			contentType = "image/webp"
		}
	default:
		// 静态图片（JPEG / PNG 等）转静态 WebP
		webpData, err := ConvertStaticImageToWebP(fileData, spec)
		if err != nil || len(webpData) >= len(fileData) {
			// 转换失败、或转换后并不比原图小（小图标/线框图/已优化 JPEG
			// 上有损 WebP 可能反而更大）→ 一律沿用原图。
			uploadData = fileData
			contentType = mimeTypeFromFormat(format)
		} else {
			uploadData = webpData
			contentType = "image/webp"
		}
	}

	// 5. 生成存储路径：`YYYY/MM/<UnixNano><ext>`（生产口径，无前缀、无 day 层）。
	//    扩展名跟随最终落盘内容类型，保证 key 扩展名与 ContentType 严格一致。
	key := uploadKey(extFromContentType(contentType))

	// 6. 获取共享 S3 客户端
	client, err := GetS3Client()
	if err != nil {
		return "", errors.Wrap(err, "get S3 client")
	}

	// 7. 上传到 R2。现有上传 API 不接收调用方 context，故显式标记为 TODO；
	// PutObject 内部会为每次尝试设置独立超时，并对可取消 context 做退避。
	if err := PutObject(context.TODO(), client, &s3.PutObjectInput{
		Bucket:        aws.String(conf.Upload.ImageBucket),
		Key:           aws.String(key),
		Body:          bytes.NewReader(uploadData),
		ContentLength: aws.Int64(int64(len(uploadData))),
		ContentType:   aws.String(contentType),
		// key 含纳秒时间戳且内容不可变 → 可安全启用 immutable 强缓存，
		// 避免每次页面加载都回源校验（曾导致头像/二维码 logo 每次重新下载）。
		CacheControl: aws.String("public, max-age=31536000, immutable"),
	}); err != nil {
		return "", errors.Wrap(err, "put object after retry")
	}

	// 8. 返回 CDN URL
	return fmt.Sprintf("https://%s/%s", conf.Upload.ImageBucketCDNHost, key), nil
}

// ConvertStaticImageToWebP 将静态图片（JPEG/PNG 等）按给定档位转换为 WebP 格式。
// spec.MaxSide 同时约束宽与高（此前只限宽，竖图长边不受约束）。
// 调用方负责比较转换前后的体积，决定是否采用转换结果。
//
// 接收已读入的字节而非 io.Reader：调用方手里已有原始字节，避免再包一层
// bytes.Reader 重复持有整份数据。
func ConvertStaticImageToWebP(data []byte, spec EncodeSpec) ([]byte, error) {
	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, errors.Wrap(err, "decode image")
	}

	// 必须先按 EXIF 归一方向再限边缩放：手机竖拍照片解出来的像素是横躺的，
	// 归一过程会交换宽高，顺序反了就会按错误的边长做缩放判断。
	img = applyJPEGOrientation(img, jpegOrientation(data))

	b := img.Bounds()
	if !fitsWithin(b.Dx(), b.Dy(), spec.MaxSide) {
		img = imaging.Fit(img, spec.MaxSide, spec.MaxSide, imaging.Lanczos)
	}

	buf := new(bytes.Buffer)
	if err := webp.Encode(buf, img, &webp.Options{Quality: spec.Quality}); err != nil {
		return nil, errors.Wrap(err, "webp encode")
	}
	return buf.Bytes(), nil
}

// fitsWithin 判断 w×h 是否已在 maxSide 的边长限制内（maxSide <= 0 时恒为 true）。
func fitsWithin(w, h, maxSide int) bool {
	if maxSide <= 0 {
		return true
	}
	return w <= maxSide && h <= maxSide
}

// maxGIFFrames 是允许转换的 GIF 最大帧数。imagecheck.CheckPixelLimit 只看 GIF 的
// 逻辑屏幕尺寸（不解码帧），一个 2000×2000、500 帧的 GIF 体积可能很小却会让
// cwebp 长时间占用 CPU 与内存，故在此显式设限。
const maxGIFFrames = 200

// gifConvertTimeout 是单次 cwebp 转换的硬上限。
// 帧数限制挡不住「帧少但每帧极大」或畸形数据把 cwebp 拖住的情况；子进程一旦挂住，
// 调用它的上传 goroutine 会一直等，最终拖死整个上传路径。超时后 exec 会杀掉子进程，
// 调用方按错误回退到原 GIF，不影响上传本身成功。
const gifConvertTimeout = 30 * time.Second

// ConvertGifToAnimatedWebP 将 GIF 转换为动画 WebP（需要系统安装 cwebp）。
// cwebp 不可用、GIF 损坏或帧数超限时返回错误，调用方应回退到原始 GIF。
func ConvertGifToAnimatedWebP(gifData []byte) ([]byte, error) {
	if _, err := exec.LookPath("cwebp"); err != nil {
		// 回退是静默的（线上表现为动画图一直以原 GIF 存盘、体积异常偏高），
		// 故在此显式告警，否则线上无法察觉 cwebp 缺失。
		logrus.Warn("cwebp not found in PATH, animated GIF will be stored as-is (install the webp package to enable animated WebP)")
		return nil, errors.New("cwebp not found in PATH")
	}

	if n := gifFrameCount(gifData); n > maxGIFFrames {
		return nil, errors.Errorf("too many GIF frames: %d > %d", n, maxGIFFrames)
	}

	tmpGif, err := os.CreateTemp("", "input_*.gif")
	if err != nil {
		return nil, errors.Wrap(err, "create temp gif")
	}
	defer os.Remove(tmpGif.Name())
	if _, err := tmpGif.Write(gifData); err != nil {
		return nil, errors.Wrap(err, "write temp gif")
	}
	tmpGif.Close()

	tmpWebp, err := os.CreateTemp("", "output_*.webp")
	if err != nil {
		return nil, errors.Wrap(err, "create temp webp")
	}
	defer os.Remove(tmpWebp.Name())
	tmpWebp.Close()

	// -mixed：对透明帧用无损、其余有损，是动画的推荐组合。
	// -metadata none：剥离调色板外的元数据，进一步减小体积。
	// 用 CommandContext 而非裸 Command：超时会杀进程，不会让上传 goroutine 无限等待。
	ctx, cancel := context.WithTimeout(context.Background(), gifConvertTimeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, "cwebp", "-q", "80", "-mixed", "-metadata", "none",
		"-mt", "-quiet", tmpGif.Name(), "-o", tmpWebp.Name())
	if err := cmd.Run(); err != nil {
		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
			logrus.WithField("timeout", gifConvertTimeout).Warn("cwebp 转换超时，回退为原 GIF")
			return nil, errors.Errorf("cwebp timed out after %s", gifConvertTimeout)
		}
		return nil, errors.Wrap(err, "run cwebp")
	}

	webpData, err := os.ReadFile(tmpWebp.Name())
	if err != nil {
		return nil, errors.Wrap(err, "read webp output")
	}
	return webpData, nil
}

// gifFrameCount 通过遍历 GIF 数据块统计 Image Descriptor 数量，
// 不解码像素，因此对超大 GIF 也是常数级开销。
func gifFrameCount(data []byte) int {
	if len(data) < 13 || string(data[:3]) != "GIF" {
		return 0
	}
	pos := 6
	// Logical Screen Descriptor 固定 7 字节，其 packed 字段在 data[10]（即 pos+4）。
	// 无条件跳过这 7 字节后再按 Global Color Table Flag 追加调色板长度——
	// 漏跳会让 pos 停在宽度低字节上，块遍历立即失效。
	if pos+7 > len(data) {
		return 0
	}
	packed := data[pos+4]
	pos += 7
	if packed&0x80 != 0 {
		pos += 3 * (1 << ((packed & 0x07) + 1))
	}
	frames := 0
	for pos < len(data) {
		switch data[pos] {
		case 0x21: // Extension Introducer：跳过标签与子块
			pos += 2
			pos = skipSubBlocks(data, pos)
		case 0x2C: // Image Descriptor：计一帧，再跳过局部色表与 LZW 子块
			if pos+10 > len(data) {
				return frames
			}
			frames++
			if data[pos+9]&0x80 != 0 { // Local Color Table Flag
				pos += 3 * (1 << ((data[pos+9] & 0x07) + 1))
			}
			pos += 10 // Image Descriptor 固定 10 字节
			if pos < len(data) {
				pos++ // LZW minimum code size
				pos = skipSubBlocks(data, pos)
			}
		case 0x3B: // Trailer
			return frames
		default:
			return frames
		}
	}
	return frames
}

// skipSubBlocks 返回子块链结束后的偏移（terminator 0x00 之后的位置）。
func skipSubBlocks(data []byte, pos int) int {
	for pos < len(data) {
		size := int(data[pos])
		pos++
		if size == 0 {
			return pos
		}
		pos += size
	}
	return pos
}

// mimeTypeFromFormat 根据图片格式返回 MIME 类型
func mimeTypeFromFormat(format string) string {
	switch format {
	case "jpeg":
		return "image/jpeg"
	case "png":
		return "image/png"
	case "gif":
		return "image/gif"
	case "webp":
		return "image/webp"
	default:
		return "application/octet-stream"
	}
}

// isSupportedImage 通过文件头魔数判断是否为支持的图片格式
func isSupportedImage(data []byte) bool {
	if len(data) < 4 {
		return false
	}

	// JPEG: FF D8 FF
	if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
		return true
	}

	// PNG: 89 50 4E 47
	if bytes.Equal(data[:4], []byte{0x89, 0x50, 0x4E, 0x47}) {
		return true
	}

	// GIF: GIF87a / GIF89a
	if len(data) >= 6 {
		if bytes.Equal(data[:6], []byte("GIF87a")) || bytes.Equal(data[:6], []byte("GIF89a")) {
			return true
		}
	}

	// WebP: RIFF ... WEBP
	if len(data) >= 12 && bytes.Equal(data[:4], []byte("RIFF")) && bytes.Equal(data[8:12], []byte("WEBP")) {
		return true
	}

	return false
}