feat: add trusted proxy identity and overload stats

httpx/clientIP.go

@@ -0,0 +1,105 @@
+package httpx
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"net/http"
+	"net/netip"
+	"strings"
+)
+
+// ClientIPResolver trusts X-Forwarded-For only from explicitly configured
+// proxy addresses. It is immutable and safe to share between requests.
+// The zero value (and a nil resolver) trusts no proxy.
+type ClientIPResolver struct {
+	trusted []netip.Prefix
+}
+
+// NewClientIPResolver accepts IP addresses or CIDR networks, never hostnames.
+func NewClientIPResolver(proxies ...string) (*ClientIPResolver, error) {
+	r := &ClientIPResolver{}
+	for _, value := range proxies {
+		value = strings.TrimSpace(value)
+		if address, err := netip.ParseAddr(value); err == nil && address.Zone() == "" {
+			address = address.Unmap()
+			r.trusted = append(r.trusted, netip.PrefixFrom(address, address.BitLen()))
+			continue
+		}
+		prefix, err := netip.ParsePrefix(value)
+		if err != nil {
+			return nil, fmt.Errorf("invalid trusted proxy %q: expected an IP address or CIDR", value)
+		}
+		if prefix.Addr().Is4In6() {
+			if prefix.Bits() < 96 {
+				return nil, fmt.Errorf("invalid trusted proxy %q: mapped IPv4 prefix must be at least /96", value)
+			}
+			prefix = netip.PrefixFrom(prefix.Addr().Unmap(), prefix.Bits()-96)
+		}
+		r.trusted = append(r.trusted, prefix.Masked())
+	}
+	return r, nil
+}
+
+func (r *ClientIPResolver) trusts(address netip.Addr) bool {
+	if r != nil {
+		for _, prefix := range r.trusted {
+			if prefix.Contains(address) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// ClientIP walks X-Forwarded-For from right to left, stopping at the first
+// untrusted address. Missing or malformed forwarding data falls back to the
+// socket peer. X-Real-IP, Forwarded and provider-specific headers are ignored.
+func (r *ClientIPResolver) ClientIP(request *http.Request) string {
+	if request == nil {
+		return ""
+	}
+	host := request.RemoteAddr
+	if value, _, err := net.SplitHostPort(host); err == nil {
+		host = value
+	}
+	peer, err := netip.ParseAddr(host)
+	if err != nil {
+		return ""
+	}
+	peer = peer.WithZone("").Unmap()
+	if !r.trusts(peer) {
+		return peer.String()
+	}
+	values := request.Header.Values("X-Forwarded-For")
+	var chain []netip.Addr
+	for _, value := range values {
+		for _, part := range strings.Split(value, ",") {
+			address, err := netip.ParseAddr(strings.TrimSpace(part))
+			if err != nil || address.Zone() != "" {
+				return peer.String()
+			}
+			chain = append(chain, address.Unmap())
+		}
+	}
+	current := peer
+	for i := len(chain) - 1; i >= 0 && r.trusts(current); i-- {
+		current = chain[i]
+	}
+	return current.String()
+}
+
+type clientIPResolverKey struct{}
+
+// WithClientIPResolver configures ClientIP for Context values derived from ctx.
+// Adapters should set this before running application middleware.
+func WithClientIPResolver(ctx context.Context, resolver *ClientIPResolver) context.Context {
+	return context.WithValue(ctx, clientIPResolverKey{}, resolver)
+}
+
+// ClientIP returns the verified client address without a port. By default it
+// returns the socket peer; App's trusted proxy configuration enables XFF.
+func (c *Context) ClientIP() string {
+	resolver, _ := c.Context().Value(clientIPResolverKey{}).(*ClientIPResolver)
+	return resolver.ClientIP(c.request)
+}

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