acfunlive-remake/avatar_test.go

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)
	if err := os.WriteFile(old, []byte("cached-png"), 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 || string(data) != "cached-png" {
		t.Fatalf("搬过来的内容不对：%v %q", err, data)
	}
	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)
	}
}

// ==================== 头像下载的容量上限 ====================
// 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)
	}
}