package utils import ( "context" "net" "net/http" "net/http/httptest" "os" "testing" "github.com/gin-gonic/gin" ) func TestMain(m *testing.M) { gin.SetMode(gin.TestMode) os.Exit(m.Run()) } func newContext(t *testing.T, remoteAddr string, headers map[string]string) *gin.Context { t.Helper() req := httptest.NewRequest(http.MethodGet, "/", nil) req.RemoteAddr = remoteAddr for key, value := range headers { req.Header.Set(key, value) } w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = req return c } func setTrusted(t *testing.T, cidrs ...string) { t.Helper() if err := SetTrustedProxies(cidrs); err != nil { t.Fatalf("SetTrustedProxies(%v) 失败: %v", cidrs, err) } t.Cleanup(func() { _ = SetTrustedProxies(nil) }) } func TestClientIPUntrustedPeer(t *testing.T) { setTrusted(t) c := newContext(t, "192.0.2.10:5555", map[string]string{ "CF-Connecting-IP": "203.0.113.9", "X-Real-IP": "203.0.113.8", }) if got := ClientIP(c); got != "192.0.2.10" { t.Errorf("ClientIP = %q, 期望直连地址 192.0.2.10", got) } } func TestClientIPHeaderPriority(t *testing.T) { setTrusted(t, "192.0.2.0/24") cases := []struct { name string headers map[string]string want string }{ {"cloudflare", map[string]string{"CF-Connecting-IP": "203.0.113.1", "X-Real-IP": "203.0.113.2"}, "203.0.113.1"}, {"true-client", map[string]string{"True-Client-IP": "203.0.113.3", "X-Real-IP": "203.0.113.2"}, "203.0.113.3"}, {"ali-cdn", map[string]string{"Ali-CDN-Real-IP": "203.0.113.4"}, "203.0.113.4"}, {"real-ip", map[string]string{"X-Real-IP": "203.0.113.5"}, "203.0.113.5"}, {"client-ip", map[string]string{"X-Client-IP": "203.0.113.6"}, "203.0.113.6"}, {"fastly", map[string]string{"Fastly-Client-IP": "203.0.113.7"}, "203.0.113.7"}, } for _, tc := range cases { c := newContext(t, "192.0.2.10:5555", tc.headers) if got := ClientIP(c); got != tc.want { t.Errorf("%s: ClientIP = %q, 期望 %q", tc.name, got, tc.want) } } } func TestClientIPInvalidValueFallsThrough(t *testing.T) { setTrusted(t, "192.0.2.10") c := newContext(t, "192.0.2.10:5555", map[string]string{ "CF-Connecting-IP": "not-an-ip", "X-Real-IP": "203.0.113.5", }) if got := ClientIP(c); got != "203.0.113.5" { t.Errorf("ClientIP = %q, 期望 203.0.113.5", got) } } func TestClientIPXForwardedFor(t *testing.T) { setTrusted(t, "192.0.2.0/24", "10.0.0.0/8") c := newContext(t, "192.0.2.10:5555", map[string]string{ "X-Forwarded-For": "198.51.100.1, 10.0.0.5", }) if got := ClientIP(c); got != "198.51.100.1" { t.Errorf("应取最右侧非可信 IP, ClientIP = %q, 期望 198.51.100.1", got) } c = newContext(t, "192.0.2.10:5555", map[string]string{ "X-Forwarded-For": "10.0.0.1, 10.0.0.2", }) if got := ClientIP(c); got != "10.0.0.1" { t.Errorf("全部可信时应取最左, ClientIP = %q, 期望 10.0.0.1", got) } } func TestClientIPForwarded(t *testing.T) { setTrusted(t, "192.0.2.10") c := newContext(t, "192.0.2.10:5555", map[string]string{ "Forwarded": `for=203.0.113.9;proto=https, for="[2001:db8::1]:4711"`, }) if got := ClientIP(c); got != "2001:db8::1" { t.Errorf("ClientIP = %q, 期望 2001:db8::1", got) } } func TestClientIPNormalization(t *testing.T) { setTrusted(t, "192.0.2.10") cases := []struct { header string value string want string }{ {"CF-Connecting-IP", "192.0.2.99:443", "192.0.2.99"}, {"X-Real-IP", `"198.51.100.7"`, "198.51.100.7"}, {"X-Real-IP", " 198.51.100.8 ", "198.51.100.8"}, {"X-Real-IP", "[2001:db8::2]:8080", "2001:db8::2"}, {"X-Real-IP", "2001:db8::3", "2001:db8::3"}, } for _, tc := range cases { c := newContext(t, "192.0.2.10:5555", map[string]string{tc.header: tc.value}) if got := ClientIP(c); got != tc.want { t.Errorf("%s=%q: ClientIP = %q, 期望 %q", tc.header, tc.value, got, tc.want) } } } func TestClientIPUnixSocketTrusted(t *testing.T) { setTrusted(t) req := httptest.NewRequest(http.MethodGet, "/", nil) req.RemoteAddr = "@" req = req.WithContext(context.WithValue(req.Context(), http.LocalAddrContextKey, &net.UnixAddr{Name: "web.sock", Net: "unix"})) req.Header.Set("CF-Connecting-IP", "203.0.113.9") w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = req if got := ClientIP(c); got != "203.0.113.9" { t.Errorf("unix socket 应视为可信, ClientIP = %q, 期望 203.0.113.9", got) } } func TestSetTrustedProxies(t *testing.T) { if err := SetTrustedProxies([]string{"192.0.2.0/24", "198.51.100.7", "2001:db8::/32", " "}); err != nil { t.Fatalf("合法输入不应报错: %v", err) } t.Cleanup(func() { _ = SetTrustedProxies(nil) }) if err := SetTrustedProxies([]string{"not-a-cidr"}); err == nil { t.Error("非法 CIDR 应报错") } } func TestRemoteIP(t *testing.T) { if got := RemoteIP(newContext(t, "192.0.2.1:1234", nil)); got != "192.0.2.1" { t.Errorf("RemoteIP = %q, 期望 192.0.2.1", got) } if got := RemoteIP(newContext(t, "@", nil)); got != "" { t.Errorf("unix socket RemoteIP = %q, 期望空串", got) } }