tamabox/internal/shared/ip_test.go

package shared

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

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

// setTrustedProxies 在测试内临时改 trusted_proxies 并重新初始化解析结果，
// 测试结束自动还原（trustedProxyNets 是包级状态，必须还原以免污染其它用例）。
func setTrustedProxies(t *testing.T, spec string) {
	t.Helper()
	old := conf.Security.TrustedProxies
	conf.Security.TrustedProxies = spec
	if err := InitTrustedProxies(); err != nil {
		t.Fatalf("InitTrustedProxies(%q) error = %v", spec, err)
	}
	t.Cleanup(func() {
		conf.Security.TrustedProxies = old
		if err := InitTrustedProxies(); err != nil {
			t.Fatalf("restore InitTrustedProxies() error = %v", err)
		}
	})
}

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

	// 回环对端 = cloudflared / 本机反代回源，这是生产实际形态。
	r := httptest.NewRequest("GET", "https://example.com/", nil)
	r.RemoteAddr = "127.0.0.1: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)
	}
}

// TestResolveClientIPIPv6LoopbackHeader 验证 ::1 对端同样被视为本机回源。
func TestResolveClientIPIPv6LoopbackHeader(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 = "[::1]:443"
	r.Header.Set("CF-Connecting-IP", "198.51.100.20")
	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 = "127.0.0.1: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 != "127.0.0.1" {
		t.Fatalf("ResolveClientIP() = %q, want remote address", got)
	}
}

// TestResolveClientIPRejectsForgedHeaderFromPublicPeer 是 P1-3 的回归测试。
//
// 攻击场景：双 CDN 并存时经 EdgeOne 打进来，EdgeOne 不剥离客户端自带的
// CF-Connecting-IP，旧实现无条件优先读它 → 客户端 IP 完全可控（限流 / 黑名单失效，
// 且可借 autoban 栽赃任意 IP）。修复后公网对端一律不采信权威头。
func TestResolveClientIPRejectsForgedHeaderFromPublicPeer(t *testing.T) {
	oldTrustProxy := conf.Security.TrustProxy
	defer func() { conf.Security.TrustProxy = oldTrustProxy }()
	conf.Security.TrustProxy = true
	setTrustedProxies(t, "") // 未登记任何代理 / 回源网段

	cases := []struct {
		name   string
		header map[string]string
	}{
		{"forged CF header via EdgeOne", map[string]string{
			"CF-Connecting-IP": "198.51.100.20",
			"EO-Client-IP":     "198.51.100.99",
		}},
		{"forged CF header only", map[string]string{
			"CF-Connecting-IP": "198.51.100.20",
		}},
		{"forged EO header only", map[string]string{
			"EO-Client-IP": "198.51.100.99",
		}},
		{"forged XFF and X-Real-IP", map[string]string{
			"X-Forwarded-For": "198.51.100.20",
			"X-Real-IP":       "198.51.100.21",
		}},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			r := httptest.NewRequest("GET", "https://example.com/", nil)
			r.RemoteAddr = "203.0.113.44:443" // 公网对端，不在 trusted_proxies 内
			for k, v := range tc.header {
				r.Header.Set(k, v)
			}
			if got := ResolveClientIP(r); got != "203.0.113.44" {
				t.Fatalf("ResolveClientIP() = %q, want TCP remote address（不可伪造的对端地址）", got)
			}
		})
	}
}

// TestResolveClientIPTrustedProxyHonorsCDNHeader 验证主动登记的 CDN 回源网段
// （或自己的 LB）仍可采信权威头 —— 这是「CDN 直连回源」部署形态的唯一正确配置方式。
func TestResolveClientIPTrustedProxyHonorsCDNHeader(t *testing.T) {
	oldTrustProxy := conf.Security.TrustProxy
	defer func() { conf.Security.TrustProxy = oldTrustProxy }()
	conf.Security.TrustProxy = true
	setTrustedProxies(t, "203.0.113.0/24")

	r := httptest.NewRequest("GET", "https://example.com/", nil)
	r.RemoteAddr = "203.0.113.44:443" // 落在已登记的回源网段内
	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 = "127.0.0.1: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 = "127.0.0.1:443"
	r.Header.Set("EO-Client-IP", "not-an-ip")
	if got := ResolveClientIP(r); got != "127.0.0.1" {
		t.Fatalf("ResolveClientIP() = %q, want remote address fallback", got)
	}
}

// TestResolveClientIPLoopbackDoesNotTrustXFF 锁定一条容易踩的边界：
// 回环对端只隐含信任「单一权威头」，不隐含信任 XFF / X-Real-IP（多跳链语义更弱）。
// 否则本机 cloudflared 场景下客户端只要多发一个 X-Real-IP 就能顶替权威头。
func TestResolveClientIPLoopbackDoesNotTrustXFF(t *testing.T) {
	oldTrustProxy := conf.Security.TrustProxy
	defer func() { conf.Security.TrustProxy = oldTrustProxy }()
	conf.Security.TrustProxy = true
	setTrustedProxies(t, "")

	r := httptest.NewRequest("GET", "https://example.com/", nil)
	r.RemoteAddr = "127.0.0.1:443"
	r.Header.Set("X-Forwarded-For", "198.51.100.20")
	r.Header.Set("X-Real-IP", "198.51.100.21")
	if got := ResolveClientIP(r); got != "127.0.0.1" {
		t.Fatalf("ResolveClientIP() = %q, want remote address（回环不应隐含信任 XFF）", 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)
	}
}