402 lines
14 KiB
Go
402 lines
14 KiB
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"
|
||
|
||
"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
|
||
|
||
// 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：剥离调色板外的元数据，进一步减小体积。
|
||
cmd := exec.Command("cwebp", "-q", "80", "-mixed", "-metadata", "none",
|
||
"-mt", "-quiet", tmpGif.Name(), "-o", tmpWebp.Name())
|
||
if err := cmd.Run(); err != nil {
|
||
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
|
||
}
|