头像下载加上体积与像素上限,并不再每次新建连接
审查发现: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:
+26
-5
@@ -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 ""
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user