tamabox/internal/shared/ip_test.go

package shared

import (
	"net/http/httptest"
	"testing"

	"github.com/tamakyi/TamaBox/internal/conf"
)

func TestResolveClientIPCloudflareHeader(t *testing.T) {
	oldTrustProxy := conf.Security.TrustProxy
	defer func() { conf.Security.TrustProxy = oldTrustProxy }()
	conf.Security.TrustProxy = true

	r := httptest.NewRequest("GET", "https://example.com/", nil)
	r.RemoteAddr = "192.0.2.10:443"
	r.Header.Set("CF-Connecting-IP", "198.51.100.20")
	r.Header.Set("X-Real-IP", "198.51.100.30")
	if got := ResolveClientIP(r); got != "198.51.100.20" {
		t.Fatalf("ResolveClientIP() = %q, want Cloudflare client IP", got)
	}
}

func TestResolveClientIPIgnoresHeadersWithoutTrust(t *testing.T) {
	oldTrustProxy := conf.Security.TrustProxy
	defer func() { conf.Security.TrustProxy = oldTrustProxy }()
	conf.Security.TrustProxy = false

	r := httptest.NewRequest("GET", "https://example.com/", nil)
	r.RemoteAddr = "192.0.2.10:443"
	r.Header.Set("CF-Connecting-IP", "198.51.100.20")
	r.Header.Set("EO-Client-IP", "198.51.100.99")
	if got := ResolveClientIP(r); got != "192.0.2.10" {
		t.Fatalf("ResolveClientIP() = %q, want remote address", got)
	}
}

// TestResolveClientIPEdgeOneHeader 验证腾讯 EdgeOne 的权威客户端头 EO-Client-IP
// 在 TrustProxy=true 时被信任（双 CDN 部署：境内 EdgeOne 回源不带 CF-Connecting-IP）。
func TestResolveClientIPEdgeOneHeader(t *testing.T) {
	oldTrustProxy := conf.Security.TrustProxy
	defer func() { conf.Security.TrustProxy = oldTrustProxy }()
	conf.Security.TrustProxy = true

	r := httptest.NewRequest("GET", "https://example.com/", nil)
	r.RemoteAddr = "203.0.113.44:443" // 模拟 EdgeOne 边缘回源
	r.Header.Set("EO-Client-IP", "198.51.100.99")
	r.Header.Set("X-Forwarded-For", "203.0.113.99")
	if got := ResolveClientIP(r); got != "198.51.100.99" {
		t.Fatalf("ResolveClientIP() = %q, want EdgeOne client IP", got)
	}
}

// TestResolveClientIPPrefersCloudflareOverEdgeOne 锁定权威头的优先级：
// 两个头并存（配置遗漏剥离规则时的异常态）时以 CF-Connecting-IP 为准，行为可预期。
func TestResolveClientIPPrefersCloudflareOverEdgeOne(t *testing.T) {
	oldTrustProxy := conf.Security.TrustProxy
	defer func() { conf.Security.TrustProxy = oldTrustProxy }()
	conf.Security.TrustProxy = true

	r := httptest.NewRequest("GET", "https://example.com/", nil)
	r.RemoteAddr = "203.0.113.44:443"
	r.Header.Set("CF-Connecting-IP", "198.51.100.20")
	r.Header.Set("EO-Client-IP", "198.51.100.99")
	if got := ResolveClientIP(r); got != "198.51.100.20" {
		t.Fatalf("ResolveClientIP() = %q, want Cloudflare client IP", got)
	}
}

// TestResolveClientIPEdgeOneInvalidFallsBack 验证 EO-Client-IP 为非法值时
// 不 panic 且回退到 RemoteAddr（与 CF-Connecting-IP 的健壮性一致）。
func TestResolveClientIPEdgeOneInvalidFallsBack(t *testing.T) {
	oldTrustProxy := conf.Security.TrustProxy
	defer func() { conf.Security.TrustProxy = oldTrustProxy }()
	conf.Security.TrustProxy = true

	r := httptest.NewRequest("GET", "https://example.com/", nil)
	r.RemoteAddr = "203.0.113.44:443"
	r.Header.Set("EO-Client-IP", "not-an-ip")
	if got := ResolveClientIP(r); got != "203.0.113.44" {
		t.Fatalf("ResolveClientIP() = %q, want remote address fallback", got)
	}
}

// TestLimitKey 是 P1-4 的回归测试：限流 / 自动封禁用的聚合键必须
//   - IPv4：原样返回（每地址独立计数）；
//   - IPv6：掩码到 /64（防止单 /64 内轮换海量地址绕过阈值、并避免键无限膨胀）；
//   - 非法输入：原样返回，保持调用方行为可预期。
func TestLimitKey(t *testing.T) {
	cases := []struct {
		name string
		in   string
		want string
	}{
		{"ipv4 passthrough", "203.0.113.10", "203.0.113.10"},
		{"ipv4-mapped collapses to v4", "::ffff:203.0.113.10", "203.0.113.10"},
		{"ipv6 aggregated to /64", "2001:db8:1:2:3:4:5:6", "2001:db8:1:2::/64"},
		// LimitKey 假定输入已由 ResolveClientIP 归一化（已剥离端口）。若误传入带端口的
		// 字符串，ParseIP 失败会原样透传 —— 这是防御性约定，需锁定以免行为漂移。
		{"ipv4 with port passes through", "203.0.113.10:8080", "203.0.113.10:8080"},
		{"ipv6 with port passes through", "[2001:db8:1:2:3:4:5:6]:443", "[2001:db8:1:2:3:4:5:6]:443"},
		{"invalid input passthrough", "not-an-ip", "not-an-ip"},
		{"empty passthrough", "", ""},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if got := LimitKey(tc.in); got != tc.want {
				t.Fatalf("LimitKey(%q) = %q, want %q", tc.in, got, tc.want)
			}
		})
	}
}

// TestLimitKeyIPv6SamePrefixCollapses 验证：同一 /64 内任意两台主机应聚合到同一键，
// 不同 /64 必须落到不同键 —— 这是限速 / 封禁抗 IPv6 轮换绕过的核心不变量。
func TestLimitKeyIPv6SamePrefixCollapses(t *testing.T) {
	a := LimitKey("2001:db8:1:2:3:4:5:6")
	b := LimitKey("2001:db8:1:2:ff:fe::1") // 同一 /64，仅主机位不同
	if a != b {
		t.Fatalf("同一 /64 内地址应聚合到同一键，got %q vs %q", a, b)
	}

	c := LimitKey("2001:db8:1:3:3:4:5:6") // 不同 /64
	if a == c {
		t.Fatalf("不同 /64 应落到不同键，两者均为 %q", a)
	}
}