tamabox/internal/route/banned_test.go

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("")
}