package route import ( "encoding/base64" "net/http" "net/http/httptest" "os" "path/filepath" "strconv" "strings" "testing" ) // tinyPNG 是 1×1 的合法 PNG,用来验证「配了图就回图」这条分支。 const tinyPNG = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" func writeTempImage(t *testing.T, data []byte) string { t.Helper() p := filepath.Join(t.TempDir(), "ban.png") if err := os.WriteFile(p, data, 0o600); err != nil { t.Fatal(err) } return p } func decodeTinyPNG(t *testing.T) []byte { t.Helper() b, err := base64.StdEncoding.DecodeString(tinyPNG) if err != nil { t.Fatal(err) } return b } // TestWriteBannedTextFallback 守住默认行为:没配图时就是改造前的纯文字 403。 func TestWriteBannedTextFallback(t *testing.T) { if loadBanImage("") { t.Fatal("空路径不应装载出图") } rec := httptest.NewRecorder() writeBanned(rec, http.MethodGet) if rec.Code != http.StatusForbidden { t.Fatalf("status = %d, want 403", rec.Code) } if ct := rec.Header().Get("Content-Type"); !strings.HasPrefix(ct, "text/plain") { t.Errorf("Content-Type = %q, want text/plain(没配图时回退纯文字)", ct) } if body := rec.Body.String(); body != banText { t.Errorf("body = %q, want %q", body, banText) } if cl := rec.Header().Get("Content-Length"); cl != strconv.Itoa(len(banText)) { t.Errorf("Content-Length = %q, want %d", cl, len(banText)) } } // TestWriteBannedServesConfiguredImage 守住「配了图就回图」。 func TestWriteBannedServesConfiguredImage(t *testing.T) { data := decodeTinyPNG(t) if !loadBanImage(writeTempImage(t, data)) { t.Fatal("合法的 PNG 应装载成功") } t.Cleanup(func() { loadBanImage("") }) rec := httptest.NewRecorder() writeBanned(rec, http.MethodGet) if rec.Code != http.StatusForbidden { t.Fatalf("status = %d, want 403(换 body 不改变状态码语义)", rec.Code) } if ct := rec.Header().Get("Content-Type"); ct != "image/png" { t.Errorf("Content-Type = %q, want image/png(按内容嗅探,不看扩展名)", ct) } if cl := rec.Header().Get("Content-Length"); cl != strconv.Itoa(len(data)) { t.Errorf("Content-Length = %q, want %d", cl, len(data)) } if body := rec.Body.Bytes(); string(body) != string(data) { t.Errorf("body 与文件内容不一致(%d vs %d 字节)", len(body), len(data)) } if cc := rec.Header().Get("Cache-Control"); !strings.Contains(cc, "no-store") { t.Errorf("Cache-Control = %q, want 含 no-store", cc) } } // TestLoadBanImageRejectsBadInput 覆盖各种配置错误:都必须降级为纯文字, // 而不是报错、更不是把文件内容原样吐出去。 func TestLoadBanImageRejectsBadInput(t *testing.T) { t.Cleanup(func() { loadBanImage("") }) cases := []struct { name string path string }{ {"不存在的文件", filepath.Join(t.TempDir(), "nope.png")}, {"是个目录", t.TempDir()}, {"不是图片", writeTempImage(t, []byte("[server]\nsalt = secret\n"))}, {"空文件", writeTempImage(t, nil)}, } for _, c := range cases { if loadBanImage(c.path) { t.Errorf("%s: 不应装载成功", c.name) } rec := httptest.NewRecorder() writeBanned(rec, http.MethodGet) if ct := rec.Header().Get("Content-Type"); !strings.HasPrefix(ct, "text/plain") { t.Errorf("%s: 应回退纯文字,实际 Content-Type = %q", c.name, ct) } if body := rec.Body.String(); strings.Contains(body, "salt") { t.Errorf("%s: 把文件内容当响应体吐了出去:%q", c.name, body) } } } // TestLoadBanImageTooLarge 守住体积上限:被封来源反复重试,大图等于放大出网带宽。 func TestLoadBanImageTooLarge(t *testing.T) { big := make([]byte, maxBanImageBytes+1) copy(big, []byte("\x89PNG\r\n\x1a\n")) if loadBanImage(writeTempImage(t, big)) { t.Error("超过 2 MB 的文件不应被装载") } loadBanImage("") } // TestWriteBannedHead 覆盖 HEAD:头照给(含 Content-Length),body 不发。 func TestWriteBannedHead(t *testing.T) { for _, withImage := range []bool{false, true} { if withImage { data := decodeTinyPNG(t) if !loadBanImage(writeTempImage(t, data)) { t.Fatal("装载失败") } } else { loadBanImage("") } rec := httptest.NewRecorder() writeBanned(rec, http.MethodHead) if rec.Code != http.StatusForbidden { t.Errorf("withImage=%v: status = %d, want 403", withImage, rec.Code) } if rec.Body.Len() != 0 { t.Errorf("withImage=%v: HEAD 不应有 body,实际 %d 字节", withImage, rec.Body.Len()) } if rec.Header().Get("Content-Length") == "" { t.Errorf("withImage=%v: HEAD 也应给出 Content-Length", withImage) } } loadBanImage("") }