0904
This commit is contained in:
tamakyi
@@ -0,0 +1,81 @@
|
||||
// Package imagecheck 提供图片尺寸校验，用于在上传/解码前拦截"解压炸弹"。
|
||||
package imagecheck
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"image"
|
||||
_ "image/gif"
|
||||
_ "image/jpeg"
|
||||
_ "image/png"
|
||||
"io"
|
||||
"math"
|
||||
)
|
||||
|
||||
// maxImageSide 是单边最大像素，MaxImagePixels 为其平方。
|
||||
const maxImageSide = 4096
|
||||
|
||||
// MaxImageSide 导出单边最大像素，供上层在错误提示中向用户展示具体限制。
|
||||
const MaxImageSide = maxImageSide
|
||||
|
||||
// MaxImagePixels 限制解码后的最大像素总量（4096×4096），用于拦截"解压炸弹"：
|
||||
// 超大纯色 PNG/GIF 压缩后体积很小，但 image.Decode 会按宽×高分配内存导致 OOM。
|
||||
const MaxImagePixels = maxImageSide * maxImageSide
|
||||
|
||||
// ErrImageTooLarge 表示图片像素总量超过 MaxImagePixels，
|
||||
// 是所有"图片尺寸超限"的统一错误，供各上传入口识别并展示统一文案。
|
||||
var ErrImageTooLarge = errors.New(fmt.Sprintf(
|
||||
"图片尺寸过大，最大支持 %d×%d 像素，请压缩后重新上传",
|
||||
maxImageSide, maxImageSide,
|
||||
))
|
||||
|
||||
// CheckPixelLimit 仅读取图片头部（image.DecodeConfig）校验宽高与像素总量，
|
||||
// 在完整解码前拦截超大图片，避免内存耗尽。
|
||||
//
|
||||
// 说明：对 GIF 同样适用——image.DecodeConfig 对 GIF 只读取逻辑屏幕宽高（不解码帧），
|
||||
// 成本极低；但大逻辑屏幕的 GIF 在后续 image.Decode 时仍会按屏幕尺寸分配内存，故必须校验。
|
||||
//
|
||||
// 注意：该方法会读取（消耗）传入的 reader，调用方需自行重建/复位 reader。
|
||||
func CheckPixelLimit(r io.Reader) error {
|
||||
cfg, _, err := image.DecodeConfig(r)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if cfg.Width <= 0 || cfg.Height <= 0 {
|
||||
return nil
|
||||
}
|
||||
if int64(cfg.Width)*int64(cfg.Height) > MaxImagePixels {
|
||||
return ErrImageTooLarge
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ErrBadAspectRatio 表示图片宽高比不符合用途要求（如头像必须 1:1、背景必须 16:9）。
|
||||
// 由 storage 在头像/背景上传入口校验后返回，route 层负责映射为对应语言的提示文案。
|
||||
var ErrBadAspectRatio = errors.New("image aspect ratio does not meet the requirement")
|
||||
|
||||
// DefaultAspectTolerance 是宽高比允许的绝对偏差。
|
||||
// 客户端裁剪输出即为精确比例，此处仅用于拦截绕过裁剪直传的明显不符图片，
|
||||
// 因此使用窄容差（±0.01）；宽高比越界（如 4:3 传成 16:9）会稳定触发拒绝。
|
||||
const DefaultAspectTolerance = 0.01
|
||||
|
||||
// CheckAspectRatio 仅读取图片头部校验宽高比是否接近 wantW:wantH（绝对偏差 ≤ tolerance）。
|
||||
// 与 CheckPixelLimit 一致：只做 DecodeConfig 级轻量检查，不解码全图。
|
||||
// 注意：该方法会读取（消耗）传入的 reader，调用方需自行重建/复位 reader。
|
||||
func CheckAspectRatio(r io.Reader, wantW, wantH float64, tolerance float64) error {
|
||||
if wantW <= 0 || wantH <= 0 {
|
||||
return nil
|
||||
}
|
||||
cfg, _, err := image.DecodeConfig(r)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if cfg.Width <= 0 || cfg.Height <= 0 {
|
||||
return nil
|
||||
}
|
||||
actual := float64(cfg.Width) / float64(cfg.Height)
|
||||
if math.Abs(actual-wantW/wantH) > tolerance {
|
||||
return ErrBadAspectRatio
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package imagecheck
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/png"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// buildGIFHeader 构造仅含文件头与逻辑屏幕描述符的最小 GIF（无像素数据），
|
||||
// 用于让 image.DecodeConfig 读取出宽高。GIF 无 CRC，构造简单可靠。
|
||||
func buildGIFHeader(t *testing.T, width, height uint16) []byte {
|
||||
t.Helper()
|
||||
var buf bytes.Buffer
|
||||
buf.WriteString("GIF89a")
|
||||
// Logical Screen Descriptor: width(2 LE) + height(2 LE) + packed(1) + bgcolor(1) + aspect(1)
|
||||
if err := binary.Write(&buf, binary.LittleEndian, width); err != nil {
|
||||
t.Fatalf("write width: %v", err)
|
||||
}
|
||||
if err := binary.Write(&buf, binary.LittleEndian, height); err != nil {
|
||||
t.Fatalf("write height: %v", err)
|
||||
}
|
||||
buf.Write([]byte{0x00, 0x00, 0x00}) // 无全局颜色表
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func TestCheckPixelLimit_NormalPNG(t *testing.T) {
|
||||
img := image.NewRGBA(image.Rect(0, 0, 100, 100))
|
||||
for x := 0; x < 100; x++ {
|
||||
for y := 0; y < 100; y++ {
|
||||
img.Set(x, y, color.RGBA{R: 255, A: 255})
|
||||
}
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
if err := png.Encode(&buf, img); err != nil {
|
||||
t.Fatalf("encode png: %v", err)
|
||||
}
|
||||
if err := CheckPixelLimit(bytes.NewReader(buf.Bytes())); err != nil {
|
||||
t.Fatalf("normal image should pass, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckPixelLimit_AtLimit(t *testing.T) {
|
||||
// 恰好等于上限 4096×4096，应通过（校验为严格大于才拒绝）。
|
||||
data := buildGIFHeader(t, 4096, 4096)
|
||||
if err := CheckPixelLimit(bytes.NewReader(data)); err != nil {
|
||||
t.Fatalf("4096×4096 should pass, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckPixelLimit_TooLarge(t *testing.T) {
|
||||
// 4096×4097 超过上限，应拒绝并返回统一错误。
|
||||
data := buildGIFHeader(t, 4096, 4097)
|
||||
err := CheckPixelLimit(bytes.NewReader(data))
|
||||
if err == nil {
|
||||
t.Fatal("oversized image should be rejected, got nil")
|
||||
}
|
||||
if !errors.Is(err, ErrImageTooLarge) {
|
||||
t.Fatalf("expected ErrImageTooLarge, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckPixelLimit_InvalidData(t *testing.T) {
|
||||
err := CheckPixelLimit(bytes.NewReader([]byte("this is not an image")))
|
||||
if err == nil {
|
||||
t.Fatal("invalid data should error, got nil")
|
||||
}
|
||||
if errors.Is(err, ErrImageTooLarge) {
|
||||
t.Fatal("invalid data should not be classified as ErrImageTooLarge")
|
||||
}
|
||||
}
|
||||
|
||||
// encodePNG 生成指定宽高的纯色 PNG，供比例校验测试使用。
|
||||
func encodePNG(t *testing.T, width, height int) []byte {
|
||||
t.Helper()
|
||||
img := image.NewRGBA(image.Rect(0, 0, width, height))
|
||||
var buf bytes.Buffer
|
||||
if err := png.Encode(&buf, img); err != nil {
|
||||
t.Fatalf("encode png: %v", err)
|
||||
}
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func TestCheckAspectRatio_Pass(t *testing.T) {
|
||||
// 1920×1080 精确 16:9
|
||||
if err := CheckAspectRatio(bytes.NewReader(encodePNG(t, 1920, 1080)), 16, 9, DefaultAspectTolerance); err != nil {
|
||||
t.Fatalf("1920×1080 should pass 16:9, got %v", err)
|
||||
}
|
||||
// 1:1 正方形
|
||||
if err := CheckAspectRatio(bytes.NewReader(encodePNG(t, 512, 512)), 1, 1, DefaultAspectTolerance); err != nil {
|
||||
t.Fatalf("512×512 should pass 1:1, got %v", err)
|
||||
}
|
||||
// 容差内的微小偏差（1281×720 ≈ 1.779，仍属 16:9）
|
||||
if err := CheckAspectRatio(bytes.NewReader(encodePNG(t, 1281, 720)), 16, 9, DefaultAspectTolerance); err != nil {
|
||||
t.Fatalf("1281×720 should pass within tolerance, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckAspectRatio_Reject(t *testing.T) {
|
||||
// 4:3 直传背景图应被拒绝
|
||||
err := CheckAspectRatio(bytes.NewReader(encodePNG(t, 800, 600)), 16, 9, DefaultAspectTolerance)
|
||||
if err == nil {
|
||||
t.Fatal("800×600 should be rejected as 16:9")
|
||||
}
|
||||
if !errors.Is(err, ErrBadAspectRatio) {
|
||||
t.Fatalf("expected ErrBadAspectRatio, got %v", err)
|
||||
}
|
||||
// 竖长条不能当头像
|
||||
if err := CheckAspectRatio(bytes.NewReader(encodePNG(t, 300, 500)), 1, 1, DefaultAspectTolerance); err == nil {
|
||||
t.Fatal("300×500 should be rejected as 1:1")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckAspectRatio_InvalidData(t *testing.T) {
|
||||
// 无法解码的数据不应误判为比例不符
|
||||
err := CheckAspectRatio(bytes.NewReader([]byte("this is not an image")), 1, 1, DefaultAspectTolerance)
|
||||
if err == nil {
|
||||
t.Fatal("invalid data should error, got nil")
|
||||
}
|
||||
if errors.Is(err, ErrBadAspectRatio) {
|
||||
t.Fatal("invalid data should not be classified as ErrBadAspectRatio")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user