tamakyi
4e8471deea
审查发现：CDN 的响应大小是外部输入，而下载头像时是整份读进内存的。 1. 体积上限 downloadAvatar 用的是裸 io.ReadAll，同文件里另一处请求头像链接的代码 （fetchAvatarURL）却规规矩矩用了 LimitReader(256K)，两处不对称。 实测：让假服务端无限吐数据，一个头像就吃掉 512MB 内存才去解码。 现在按 avatarMaxBytes（2MB）截断 —— 真实头像最大也就几百 KB，余量很足。 多读一个字节才能分辨「刚好到上限」和「超过了上限」。 2. 像素上限（压缩炸弹） 光限制文件体积不够：几十 KB 的 PNG/GIF 能解压出上亿像素，几 MB 的文件 就能吃掉几 GB。decodeAvatar 解码后按 avatarMaxPixels（2048x2048）再查一次。 头像最终只会被缩到几十像素（跟随字号），根本用不到那么大的原图。 3. 复用连接 downloadAvatar 和 fetchAvatarURL 原来每次调用都新建一个 http.Client， 等于每个头像都重新握手一次，一场直播几百个头像就是几百次。 同项目的 notifyClient / meowClient 早就是包级复用了，这里补上。 已验证： - 新增 3 个用例： TestDownloadAvatarCapsResponseBody（假服务端吐 10MB，断言整份没被读完且 报错要写清是体积超限、失败不留文件）、 TestDownloadAvatarAcceptsNormalSize（96x96 正常头像不受影响）、 TestDecodeAvatarCapsPixels（3000x3000 被拦，300x300 放行）。 - 反向验证：去掉 LimitReader 后报「整份响应都被读完了（10.0 MB）」； 去掉像素检查后报「3000x3000 应当被尺寸上限拦住」。 - 一处测试写法修正：最初断言的是服务端「发出」的字节数（含 socket 缓冲， 实测 5~7MB 抖动导致误报），改成断言「是否读满了服务端准备的量」才准确 —— 客户端读到上限就撒手，服务端根本写不完整个响应。连跑三次确认稳定。 - gofmt / go vet / go build / GOOS=linux go build / go test / go test -race 全过。
269 lines
8.7 KiB
Go
269 lines
8.7 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"crypto/sha1"
|
||
"encoding/base64"
|
||
"fmt"
|
||
"image"
|
||
"image/color"
|
||
"image/gif"
|
||
"image/jpeg"
|
||
"image/png"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
)
|
||
|
||
// 2x2 的静态 webp（纯红），用来确认 webp 解码器确实注册上了
|
||
const tinyWebPBase64 = "UklGRjwAAABXRUJQVlA4IDAAAADQAQCdASoCAAIAAgA0JaACdLoB+AADsAD+8Oj3/yC5YXXI1/8gP+QH/ID/+PIAAAA="
|
||
|
||
// 多帧 gif 解出来应该是第一帧（AcFun 上有动图头像，压成静态就用第一帧）
|
||
func TestDecodeAvatarGifTakesFirstFrame(t *testing.T) {
|
||
pal := color.Palette{color.RGBA{R: 255, A: 255}, color.RGBA{B: 255, A: 255}}
|
||
first := image.NewPaletted(image.Rect(0, 0, 8, 8), pal)
|
||
second := image.NewPaletted(image.Rect(0, 0, 8, 8), pal)
|
||
for i := range second.Pix {
|
||
second.Pix[i] = 1
|
||
}
|
||
var buf bytes.Buffer
|
||
if err := gif.EncodeAll(&buf, &gif.GIF{
|
||
Image: []*image.Paletted{first, second},
|
||
Delay: []int{10, 10},
|
||
}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
src, err := decodeAvatar(buf.Bytes(), "image/gif")
|
||
if err != nil {
|
||
t.Fatalf("gif 解码失败：%v", err)
|
||
}
|
||
r, g, b, _ := src.At(4, 4).RGBA()
|
||
if r>>8 != 255 || g>>8 != 0 || b>>8 != 0 {
|
||
t.Errorf("应该取第一帧（红），实际拿到 r=%d g=%d b=%d", r>>8, g>>8, b>>8)
|
||
}
|
||
}
|
||
|
||
// jpg / png / 静态 webp 都能解，解不出来的要给中文原因
|
||
func TestDecodeAvatarFormats(t *testing.T) {
|
||
var pngBuf, jpgBuf bytes.Buffer
|
||
img := image.NewRGBA(image.Rect(0, 0, 8, 8))
|
||
if err := png.Encode(&pngBuf, img); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if err := jpeg.Encode(&jpgBuf, img, nil); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
webpData, err := base64.StdEncoding.DecodeString(tinyWebPBase64)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
cases := []struct {
|
||
name string
|
||
data []byte
|
||
contentType string
|
||
wantErr bool
|
||
}{
|
||
{"png", pngBuf.Bytes(), "image/png", false},
|
||
{"jpeg", jpgBuf.Bytes(), "image/jpeg", false},
|
||
{"webp", webpData, "image/webp", false},
|
||
{"不是图片", []byte("这不是图片"), "text/html", true},
|
||
}
|
||
for _, c := range cases {
|
||
t.Run(c.name, func(t *testing.T) {
|
||
src, err := decodeAvatar(c.data, c.contentType)
|
||
if c.wantErr {
|
||
if err == nil {
|
||
t.Fatal("应该报错")
|
||
}
|
||
if !bytes.Contains([]byte(err.Error()), []byte(c.contentType)) {
|
||
t.Errorf("错误信息里应该带上 Content-Type 方便排查：%v", err)
|
||
}
|
||
return
|
||
}
|
||
if err != nil {
|
||
t.Fatalf("解码失败：%v", err)
|
||
}
|
||
if src == nil {
|
||
t.Fatal("解出来是空的")
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// 解出来的头像要能走通圆形裁切 + 缓存用的 PNG 编码（gif 是 *image.Paletted，和 jpg/png 的底层类型都不同）
|
||
func TestDecodeAvatarGifGoesThroughCircleAvatar(t *testing.T) {
|
||
pal := color.Palette{color.RGBA{G: 255, A: 255}, color.RGBA{R: 255, B: 255, A: 255}}
|
||
frame := image.NewPaletted(image.Rect(0, 0, 16, 16), pal)
|
||
var buf bytes.Buffer
|
||
if err := gif.EncodeAll(&buf, &gif.GIF{Image: []*image.Paletted{frame}, Delay: []int{10}}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
src, err := decodeAvatar(buf.Bytes(), "image/gif")
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
var out bytes.Buffer
|
||
if err := png.Encode(&out, circleAvatar(src, avatarCacheSize)); err != nil {
|
||
t.Fatalf("圆形裁切后的 PNG 编码失败：%v", err)
|
||
}
|
||
if out.Len() == 0 {
|
||
t.Fatal("编码出来是空的")
|
||
}
|
||
}
|
||
|
||
// 头像缓存在程序所在目录，不再跟着保存目录跑
|
||
func TestAvatarCachePathUnderExeDir(t *testing.T) {
|
||
old := exeDir
|
||
defer func() { exeDir = old }()
|
||
|
||
exeDir = filepath.Join("C:", "app", "acfunlive")
|
||
want := filepath.Join("C:", "app", "acfunlive", avatarCacheDir)
|
||
if got := avatarCachePath(); got != want {
|
||
t.Errorf("应缓存到 %q，实际 %q", want, got)
|
||
}
|
||
|
||
// 取不到程序目录时退回相对路径，别拼出 /avatar_cache 这种绝对路径
|
||
exeDir = ""
|
||
if got := avatarCachePath(); got != avatarCacheDir {
|
||
t.Errorf("没有程序目录时应退回 %q，实际 %q", avatarCacheDir, got)
|
||
}
|
||
}
|
||
|
||
// 旧版本把头像存在录播目录里，换保存目录后要能直接搬过来用，不重新下载
|
||
func TestEnsureAvatarMigratesLegacyCache(t *testing.T) {
|
||
root := t.TempDir()
|
||
cacheDir := filepath.Join(root, "program", avatarCacheDir)
|
||
legacyDir := filepath.Join(root, "record", avatarCacheDir)
|
||
if err := os.MkdirAll(legacyDir, 0755); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
// 这个地址连不通：命中旧缓存就不该去下载（下载会重试 3 次，很慢）
|
||
url := "http://127.0.0.1:9/never.png"
|
||
sum := sha1.Sum([]byte(url))
|
||
name := fmt.Sprintf("%d_%x.png", int64(7), sum[:6])
|
||
old := filepath.Join(legacyDir, name)
|
||
if err := os.WriteFile(old, []byte("cached-png"), 0644); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
got := ensureAvatar(7, url, cacheDir, legacyDir)
|
||
want := filepath.Join(cacheDir, name)
|
||
if got != want {
|
||
t.Fatalf("应搬到 %q，实际 %q", want, got)
|
||
}
|
||
data, err := os.ReadFile(want)
|
||
if err != nil || string(data) != "cached-png" {
|
||
t.Fatalf("搬过来的内容不对：%v %q", err, data)
|
||
}
|
||
if _, err := os.Stat(old); !os.IsNotExist(err) {
|
||
t.Errorf("旧位置的文件应被搬走，实际还在：%v", err)
|
||
}
|
||
|
||
// 已经在程序目录里的直接命中，不用管旧目录
|
||
if got := ensureAvatar(7, url, cacheDir, legacyDir); got != want {
|
||
t.Errorf("命中缓存应返回 %q，实际 %q", want, got)
|
||
}
|
||
}
|
||
|
||
// ==================== 头像下载的容量上限 ====================
|
||
// CDN 的响应大小和图片尺寸都是外部输入，必须有上限。
|
||
// 改之前 downloadAvatar 用的是裸 io.ReadAll：实测一个头像就能吃掉 512MB 内存。
|
||
|
||
// 异常大的响应必须在读完之前就被截住，不能整份读进内存
|
||
func TestDownloadAvatarCapsResponseBody(t *testing.T) {
|
||
var wanted int64 = avatarMaxBytes + 8<<20 // 比上限多 8MB
|
||
var served int64
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Header().Set("Content-Type", "image/png")
|
||
buf := make([]byte, 1<<20)
|
||
for atomic.AddInt64(&served, int64(len(buf))) < wanted {
|
||
if _, err := w.Write(buf); err != nil {
|
||
return
|
||
}
|
||
}
|
||
}))
|
||
defer srv.Close()
|
||
|
||
out := filepath.Join(t.TempDir(), "a.png")
|
||
done := make(chan error, 1)
|
||
go func() { done <- downloadAvatar(1, srv.URL, out) }()
|
||
|
||
select {
|
||
case err := <-done:
|
||
if err == nil {
|
||
t.Fatal("超过体积上限的头像应当报错，实际下载成功了")
|
||
}
|
||
if !strings.Contains(err.Error(), "超过") {
|
||
t.Errorf("报错要说清是体积超限，实际：%v", err)
|
||
}
|
||
// 客户端只读到上限就撒手，服务端这边根本写不完整个响应。
|
||
// 反过来，如果 served 达到了 wanted，说明整份都被读进内存了 —— 限制没生效
|
||
if got := atomic.LoadInt64(&served); got >= wanted {
|
||
t.Errorf("整份响应都被读完了（%.1f MB），说明没有截断（上限 %d MB）",
|
||
float64(got)/1048576, avatarMaxBytes>>20)
|
||
}
|
||
t.Logf(" 已接收 %.1f MB 时就被截住（上限 %d MB）",
|
||
float64(atomic.LoadInt64(&served))/1048576, avatarMaxBytes>>20)
|
||
case <-time.After(30 * time.Second):
|
||
t.Fatalf("超过 30 秒还没结束，限制没生效（已接收 %.1f MB）",
|
||
float64(atomic.LoadInt64(&served))/1048576)
|
||
}
|
||
if _, err := os.Stat(out); err == nil {
|
||
t.Error("下载失败不该留下文件")
|
||
}
|
||
}
|
||
|
||
// 正常的头像是能下下来的：上限不能误伤
|
||
func TestDownloadAvatarAcceptsNormalSize(t *testing.T) {
|
||
var buf bytes.Buffer
|
||
if err := png.Encode(&buf, image.NewRGBA(image.Rect(0, 0, 96, 96))); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Header().Set("Content-Type", "image/png")
|
||
_, _ = w.Write(buf.Bytes())
|
||
}))
|
||
defer srv.Close()
|
||
|
||
out := filepath.Join(t.TempDir(), "ok.png")
|
||
if err := downloadAvatar(1, srv.URL, out); err != nil {
|
||
t.Fatalf("正常大小的头像应当能下载，实际报错：%v", err)
|
||
}
|
||
if _, err := os.Stat(out); err != nil {
|
||
t.Errorf("下载成功却没生成文件：%v", err)
|
||
}
|
||
}
|
||
|
||
// 压缩炸弹：几十 KB 的图可以解压出上亿像素，解码后必须再检查一次尺寸
|
||
func TestDecodeAvatarCapsPixels(t *testing.T) {
|
||
side := 3000 // 900 万像素，刚过 2048x2048 的上限
|
||
big := image.NewGray(image.Rect(0, 0, side, side))
|
||
var buf bytes.Buffer
|
||
if err := png.Encode(&buf, big); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, err := decodeAvatar(buf.Bytes(), "image/png"); err == nil {
|
||
t.Fatalf("%dx%d 的头像应当被尺寸上限拦住", side, side)
|
||
} else if !strings.Contains(err.Error(), "超出上限") {
|
||
t.Errorf("报错要说清是尺寸超限，实际：%v", err)
|
||
}
|
||
|
||
// 常见的头像尺寸不受影响
|
||
ok := image.NewRGBA(image.Rect(0, 0, 300, 300))
|
||
var b2 bytes.Buffer
|
||
if err := png.Encode(&b2, ok); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if _, err := decodeAvatar(b2.Bytes(), "image/png"); err != nil {
|
||
t.Errorf("300x300 的普通头像不该被拦：%v", err)
|
||
}
|
||
}
|