package main import ( "bytes" "crypto/sha1" "encoding/base64" "fmt" "image" "image/color" "image/gif" "image/jpeg" "image/png" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "sync/atomic" "testing" "time" ) // 2x2 的静态 webp(纯红),用来确认 webp 解码器确实注册上了 const tinyWebPBase64 = "UklGRjwAAABXRUJQVlA4IDAAAADQAQCdASoCAAIAAgA0JaACdLoB+AADsAD+8Oj3/yC5YXXI1/8gP+QH/ID/+PIAAAA=" // 多帧 gif 解出来应该是第一帧(AcFun 上有动图头像,压成静态就用第一帧) func TestDecodeAvatarGifTakesFirstFrame(t *testing.T) { pal := color.Palette{color.RGBA{R: 255, A: 255}, color.RGBA{B: 255, A: 255}} first := image.NewPaletted(image.Rect(0, 0, 8, 8), pal) second := image.NewPaletted(image.Rect(0, 0, 8, 8), pal) for i := range second.Pix { second.Pix[i] = 1 } var buf bytes.Buffer if err := gif.EncodeAll(&buf, &gif.GIF{ Image: []*image.Paletted{first, second}, Delay: []int{10, 10}, }); err != nil { t.Fatal(err) } src, err := decodeAvatar(buf.Bytes(), "image/gif") if err != nil { t.Fatalf("gif 解码失败:%v", err) } r, g, b, _ := src.At(4, 4).RGBA() if r>>8 != 255 || g>>8 != 0 || b>>8 != 0 { t.Errorf("应该取第一帧(红),实际拿到 r=%d g=%d b=%d", r>>8, g>>8, b>>8) } } // jpg / png / 静态 webp 都能解,解不出来的要给中文原因 func TestDecodeAvatarFormats(t *testing.T) { var pngBuf, jpgBuf bytes.Buffer img := image.NewRGBA(image.Rect(0, 0, 8, 8)) if err := png.Encode(&pngBuf, img); err != nil { t.Fatal(err) } if err := jpeg.Encode(&jpgBuf, img, nil); err != nil { t.Fatal(err) } webpData, err := base64.StdEncoding.DecodeString(tinyWebPBase64) if err != nil { t.Fatal(err) } cases := []struct { name string data []byte contentType string wantErr bool }{ {"png", pngBuf.Bytes(), "image/png", false}, {"jpeg", jpgBuf.Bytes(), "image/jpeg", false}, {"webp", webpData, "image/webp", false}, {"不是图片", []byte("这不是图片"), "text/html", true}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { src, err := decodeAvatar(c.data, c.contentType) if c.wantErr { if err == nil { t.Fatal("应该报错") } if !bytes.Contains([]byte(err.Error()), []byte(c.contentType)) { t.Errorf("错误信息里应该带上 Content-Type 方便排查:%v", err) } return } if err != nil { t.Fatalf("解码失败:%v", err) } if src == nil { t.Fatal("解出来是空的") } }) } } // 解出来的头像要能走通圆形裁切 + 缓存用的 PNG 编码(gif 是 *image.Paletted,和 jpg/png 的底层类型都不同) func TestDecodeAvatarGifGoesThroughCircleAvatar(t *testing.T) { pal := color.Palette{color.RGBA{G: 255, A: 255}, color.RGBA{R: 255, B: 255, A: 255}} frame := image.NewPaletted(image.Rect(0, 0, 16, 16), pal) var buf bytes.Buffer if err := gif.EncodeAll(&buf, &gif.GIF{Image: []*image.Paletted{frame}, Delay: []int{10}}); err != nil { t.Fatal(err) } src, err := decodeAvatar(buf.Bytes(), "image/gif") if err != nil { t.Fatal(err) } var out bytes.Buffer if err := png.Encode(&out, circleAvatar(src, avatarCacheSize)); err != nil { t.Fatalf("圆形裁切后的 PNG 编码失败:%v", err) } if out.Len() == 0 { t.Fatal("编码出来是空的") } } // 头像缓存在程序所在目录,不再跟着保存目录跑 func TestAvatarCachePathUnderExeDir(t *testing.T) { old := exeDir defer func() { exeDir = old }() exeDir = filepath.Join("C:", "app", "acfunlive") want := filepath.Join("C:", "app", "acfunlive", avatarCacheDir) if got := avatarCachePath(); got != want { t.Errorf("应缓存到 %q,实际 %q", want, got) } // 取不到程序目录时退回相对路径,别拼出 /avatar_cache 这种绝对路径 exeDir = "" if got := avatarCachePath(); got != avatarCacheDir { t.Errorf("没有程序目录时应退回 %q,实际 %q", avatarCacheDir, got) } } // 旧版本把头像存在录播目录里,换保存目录后要能直接搬过来用,不重新下载 func TestEnsureAvatarMigratesLegacyCache(t *testing.T) { root := t.TempDir() cacheDir := filepath.Join(root, "program", avatarCacheDir) legacyDir := filepath.Join(root, "record", avatarCacheDir) if err := os.MkdirAll(legacyDir, 0755); err != nil { t.Fatal(err) } // 这个地址连不通:命中旧缓存就不该去下载(下载会重试 3 次,很慢) url := "http://127.0.0.1:9/never.png" sum := sha1.Sum([]byte(url)) name := fmt.Sprintf("%d_%x.png", int64(7), sum[:6]) old := filepath.Join(legacyDir, name) // 旧位置放的是真 PNG(旧版本自己也是用 png.Encode 写的): // 现在搬过来之前会校验完整性,随便一段字节的话会被当成半截文件丢掉 var pngBuf bytes.Buffer if err := png.Encode(&pngBuf, image.NewRGBA(image.Rect(0, 0, 8, 8))); err != nil { t.Fatal(err) } if err := os.WriteFile(old, pngBuf.Bytes(), 0644); err != nil { t.Fatal(err) } got := ensureAvatar(7, url, cacheDir, legacyDir) want := filepath.Join(cacheDir, name) if got != want { t.Fatalf("应搬到 %q,实际 %q", want, got) } data, err := os.ReadFile(want) if err != nil || !bytes.Equal(data, pngBuf.Bytes()) { t.Fatalf("搬过来的内容不对:%v(%d 字节 vs %d 字节)", err, len(data), pngBuf.Len()) } if _, err := os.Stat(old); !os.IsNotExist(err) { t.Errorf("旧位置的文件应被搬走,实际还在:%v", err) } // 已经在程序目录里的直接命中,不用管旧目录 if got := ensureAvatar(7, url, cacheDir, legacyDir); got != want { t.Errorf("命中缓存应返回 %q,实际 %q", want, got) } } // 旧位置那份也是半截文件时,不能搬过来用:会一路传给 ffmpeg 当输入 func TestEnsureAvatarRejectsBrokenLegacyCache(t *testing.T) { withJobFlags(t, true, false, false) root := t.TempDir() cacheDir := filepath.Join(root, "program", avatarCacheDir) legacyDir := filepath.Join(root, "record", avatarCacheDir) if err := os.MkdirAll(legacyDir, 0755); err != nil { t.Fatal(err) } url := "http://127.0.0.1:9/never.png" old := filepath.Join(legacyDir, avatarCacheName(8, url)) if err := os.WriteFile(old, []byte("这根本不是 PNG"), 0644); err != nil { t.Fatal(err) } // 连不通的地址 + 坏缓存 → 拿不到头像,但绝不能把坏文件留在缓存目录里 if got := ensureAvatar(8, url, cacheDir); got != "" { t.Fatalf("坏缓存不该被采用,实际返回 %q", got) } if validCachedAvatar(filepath.Join(cacheDir, avatarCacheName(8, url))) { t.Fatal("坏文件不该出现在缓存目录里") } } // ==================== 头像下载的容量上限 ==================== // 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) } }