头像下载加上体积与像素上限,并不再每次新建连接

审查发现: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 全过。
This commit is contained in:
tamakyi
2026-09-26 03:06:54 +08:00
parent c55fe3539d
commit 4e8471deea
2 changed files with 125 additions and 5 deletions
+26 -5
View File
@@ -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 ""
}