头像下载加上体积与像素上限，并不再每次新建连接
审查发现：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
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
+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