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 "" }