feat: add trusted proxy identity and overload stats

httpx/clientIP_test.go

@@ -0,0 +1,76 @@
+package httpx_test
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"git.yuntuo.hk/k2safe/go-omni/httpx"
+)
+
+func TestClientIPResolverUsesSocketPeerByDefault(t *testing.T) {
+	resolver, err := httpx.NewClientIPResolver()
+	if err != nil {
+		t.Fatal(err)
+	}
+	request := newClientIPRequest("192.0.2.10:443", "198.51.100.7")
+	if got := resolver.ClientIP(request); got != "192.0.2.10" {
+		t.Fatalf("untrusted forwarding header changed client IP: %q", got)
+	}
+}
+
+func TestClientIPResolverWalksTrustedForwardingChain(t *testing.T) {
+	resolver, err := httpx.NewClientIPResolver("10.0.0.0/8")
+	if err != nil {
+		t.Fatal(err)
+	}
+	request := newClientIPRequest("10.0.0.5:443", "198.51.100.7, 10.0.0.4")
+	if got := resolver.ClientIP(request); got != "198.51.100.7" {
+		t.Fatalf("expected original client IP, got %q", got)
+	}
+
+	request = newClientIPRequest("10.0.0.5:443", "198.51.100.7, 203.0.113.9, 10.0.0.4")
+	if got := resolver.ClientIP(request); got != "203.0.113.9" {
+		t.Fatalf("expected first untrusted hop, got %q", got)
+	}
+}
+
+func TestClientIPResolverRejectsMalformedForwardingData(t *testing.T) {
+	resolver, err := httpx.NewClientIPResolver("127.0.0.1")
+	if err != nil {
+		t.Fatal(err)
+	}
+	request := newClientIPRequest("127.0.0.1:443", "198.51.100.7, not-an-ip")
+	if got := resolver.ClientIP(request); got != "127.0.0.1" {
+		t.Fatalf("malformed forwarding header should fall back to peer, got %q", got)
+	}
+}
+
+func TestNewClientIPResolverRejectsHostnames(t *testing.T) {
+	if _, err := httpx.NewClientIPResolver("proxy.example.com"); err == nil || !strings.Contains(err.Error(), "invalid trusted proxy") {
+		t.Fatalf("expected hostname validation error, got %v", err)
+	}
+}
+
+func TestContextClientIPUsesResolverFromRequestContext(t *testing.T) {
+	resolver, err := httpx.NewClientIPResolver("127.0.0.1")
+	if err != nil {
+		t.Fatal(err)
+	}
+	request := newClientIPRequest("127.0.0.1:443", "198.51.100.7")
+	request = request.WithContext(httpx.WithClientIPResolver(request.Context(), resolver))
+	ctx := httpx.NewContext(request, httptest.NewRecorder(), nil)
+	if got := ctx.ClientIP(); got != "198.51.100.7" {
+		t.Fatalf("expected context client IP, got %q", got)
+	}
+}
+
+func newClientIPRequest(remote, forwarded string) *http.Request {
+	request := httptest.NewRequest(http.MethodGet, "http://example.test/", nil)
+	request.RemoteAddr = remote
+	if forwarded != "" {
+		request.Header.Set("X-Forwarded-For", forwarded)
+	}
+	return request
+}