rill/internal/utils/ip_test.go

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