From 4e8471deea64be977d3887366d5cdb45cb93c761 Mon Sep 17 00:00:00 2001 From: tamakyi Date: Sat, 26 Sep 2026 03:06:54 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A4=B4=E5=83=8F=E4=B8=8B=E8=BD=BD=E5=8A=A0?= =?UTF-8?q?=E4=B8=8A=E4=BD=93=E7=A7=AF=E4=B8=8E=E5=83=8F=E7=B4=A0=E4=B8=8A?= =?UTF-8?q?=E9=99=90=EF=BC=8C=E5=B9=B6=E4=B8=8D=E5=86=8D=E6=AF=8F=E6=AC=A1?= =?UTF-8?q?=E6=96=B0=E5=BB=BA=E8=BF=9E=E6=8E=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 审查发现：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 全过。 --- avatar_test.go | 99 ++++++++++++++++++++++++++++++++++++++++++++++++++ danmuburn.go | 31 +++++++++++++--- 2 files changed, 125 insertions(+), 5 deletions(-) diff --git a/avatar_test.go b/avatar_test.go index 5a389e7..9330c8d 100644 --- a/avatar_test.go +++ b/avatar_test.go @@ -10,9 +10,14 @@ import ( "image/gif" "image/jpeg" "image/png" + "net/http" + "net/http/httptest" "os" "path/filepath" + "strings" + "sync/atomic" "testing" + "time" ) // 2x2 的静态 webp（纯红），用来确认 webp 解码器确实注册上了 @@ -167,3 +172,97 @@ func TestEnsureAvatarMigratesLegacyCache(t *testing.T) { 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) + } +} diff --git a/danmuburn.go b/danmuburn.go index b4a605c..c1516c1 100644 --- a/danmuburn.go +++ b/danmuburn.go @@ -1100,14 +1100,24 @@ func ensureAvatar(uid int64, url, cacheDir string, legacyDirs ...string) string return path } +const ( + // 头像下载的体积上限。真实头像最大也就几百 KB，这里给个大方的余量。 + // 没有上限的话，CDN 一旦返回异常大的响应就能把内存吃光（实测单个头像吃掉 512MB） + avatarMaxBytes = 2 << 20 + // 解码后允许的像素数上限（长 × 宽）。2048×2048 对头像来说是夸张地宽裕 + avatarMaxPixels = 2048 * 2048 +) + +// 复用连接：一场直播要下几百个头像，每个都新建 client 就是几百次握手 +var avatarClient = &http.Client{Timeout: 30 * time.Second} + func downloadAvatar(uid int64, url, path string) error { - client := &http.Client{Timeout: 30 * time.Second} req, err := http.NewRequest("GET", url, nil) if err != nil { return err } req.Header.Set("User-Agent", "Mozilla/5.0") - resp, err := client.Do(req) + resp, err := avatarClient.Do(req) if err != nil { return err } @@ -1115,10 +1125,14 @@ func downloadAvatar(uid int64, url, path string) error { if resp.StatusCode != http.StatusOK { return fmt.Errorf("HTTP %d", resp.StatusCode) } - data, err := io.ReadAll(resp.Body) + // 多读一个字节才能分辨「刚好到上限」和「超过了上限」 + data, err := io.ReadAll(io.LimitReader(resp.Body, avatarMaxBytes+1)) if err != nil { return err } + if len(data) > avatarMaxBytes { + return fmt.Errorf("头像体积超过 %d MB，不下载", avatarMaxBytes>>20) + } src, err := decodeAvatar(data, resp.Header.Get("Content-Type")) if err != nil { return err @@ -1140,6 +1154,11 @@ func decodeAvatar(data []byte, contentType string) (image.Image, error) { if err != nil { return nil, fmt.Errorf("图片解析不了（Content-Type: %s）：%v", contentType, err) } + // 防压缩炸弹：几十 KB 的 PNG/GIF 能解压出上亿像素，几 MB 的文件就能吃掉几 GB 内存。 + // 头像最后只会被缩到几十像素（跟随字号），根本用不到这么大的原图 + if b := src.Bounds(); b.Dx()*b.Dy() > avatarMaxPixels { + return nil, fmt.Errorf("头像尺寸 %dx%d 超出上限（最多 %d 像素）", b.Dx(), b.Dy(), avatarMaxPixels) + } return src, nil } @@ -1175,16 +1194,18 @@ func avatarURLByUID(uid int64) string { return url } +// 查头像链接用的 client，同样复用连接 +var avatarAPIClient = &http.Client{Timeout: 15 * time.Second} + // 查 AcFun 的用户信息接口取头像地址 func fetchAvatarURL(uid int64) string { - client := &http.Client{Timeout: 15 * time.Second} addr := fmt.Sprintf("https://www.acfun.cn/rest/pc-direct/user/userInfo?userId=%d", uid) req, err := http.NewRequest("GET", addr, nil) if err != nil { return "" } req.Header.Set("User-Agent", "Mozilla/5.0") - resp, err := client.Do(req) + resp, err := avatarAPIClient.Do(req) if err != nil { return "" }