common/xff/fastlyxff/xff_test.go

package fastlyxff

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

func TestFastlyIPRanges(t *testing.T) {
	fastlyxff, err := New("fastly.json")
	if err != nil {
		t.Fatalf("could not load test data: %s", err)
	}

	data, err := fastlyxff.EchoTrustOption()
	if err != nil {
		t.Fatalf("could not parse test data: %s", err)
	}

	if len(data) < 10 {
		t.Logf("only got %d prefixes, expected more", len(data))
		t.Fail()
	}
}

func TestHTTPMiddleware(t *testing.T) {
	// Create a test FastlyXFF instance with known IP ranges
	xff := &FastlyXFF{
		IPv4: []string{"192.0.2.0/24", "203.0.113.0/24"},
		IPv6: []string{"2001:db8::/32"},
	}

	middleware := xff.HTTPMiddleware()

	tests := []struct {
		name           string
		remoteAddr     string
		xForwardedFor  string
		expectedRealIP string
	}{
		{
			name:           "direct connection",
			remoteAddr:     "198.51.100.1:12345",
			xForwardedFor:  "",
			expectedRealIP: "198.51.100.1",
		},
		{
			name:           "trusted proxy with XFF",
			remoteAddr:     "192.0.2.1:80",
			xForwardedFor:  "198.51.100.1",
			expectedRealIP: "198.51.100.1",
		},
		{
			name:           "trusted proxy with multiple XFF",
			remoteAddr:     "192.0.2.1:80",
			xForwardedFor:  "198.51.100.1, 203.0.113.1",
			expectedRealIP: "198.51.100.1",
		},
		{
			name:           "untrusted proxy ignored",
			remoteAddr:     "198.51.100.2:80",
			xForwardedFor:  "10.0.0.1",
			expectedRealIP: "198.51.100.2",
		},
		{
			name:           "IPv6 trusted proxy",
			remoteAddr:     "[2001:db8::1]:80",
			xForwardedFor:  "198.51.100.1",
			expectedRealIP: "198.51.100.1",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create test handler that captures both GetRealIP and r.RemoteAddr
			var capturedRealIP, capturedRemoteAddr string
			handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				capturedRealIP = GetRealIP(r)
				capturedRemoteAddr = r.RemoteAddr
				w.WriteHeader(http.StatusOK)
			})

			// Create request with middleware
			req := httptest.NewRequest("GET", "/", nil)
			req.RemoteAddr = tt.remoteAddr
			if tt.xForwardedFor != "" {
				req.Header.Set("X-Forwarded-For", tt.xForwardedFor)
			}

			rr := httptest.NewRecorder()
			middleware(handler).ServeHTTP(rr, req)

			// Test GetRealIP function
			if capturedRealIP != tt.expectedRealIP {
				t.Errorf("GetRealIP: expected %s, got %s", tt.expectedRealIP, capturedRealIP)
			}

			// Test that r.RemoteAddr is updated with real IP and port 0
			// (since the original port is from the proxy, not the real client)
			expectedRemoteAddr := net.JoinHostPort(tt.expectedRealIP, "0")
			if capturedRemoteAddr != expectedRemoteAddr {
				t.Errorf("RemoteAddr: expected %s, got %s", expectedRemoteAddr, capturedRemoteAddr)
			}
		})
	}
}

func TestIsTrustedProxy(t *testing.T) {
	xff := &FastlyXFF{
		IPv4: []string{"192.0.2.0/24", "203.0.113.0/24"},
		IPv6: []string{"2001:db8::/32"},
	}

	tests := []struct {
		ip       string
		expected bool
	}{
		{"192.0.2.1", true},
		{"192.0.2.255", true},
		{"203.0.113.1", true},
		{"192.0.3.1", false},
		{"198.51.100.1", false},
		{"2001:db8::1", true},
		{"2001:db8:ffff::1", true},
		{"2001:db9::1", false},
		{"invalid-ip", false},
	}

	for _, tt := range tests {
		t.Run(tt.ip, func(t *testing.T) {
			result := xff.isTrustedProxy(tt.ip)
			if result != tt.expected {
				t.Errorf("isTrustedProxy(%s) = %v, expected %v", tt.ip, result, tt.expected)
			}
		})
	}
}

func TestExtractRealIP(t *testing.T) {
	xff := &FastlyXFF{
		IPv4: []string{"192.0.2.0/24"},
		IPv6: []string{"2001:db8::/32"},
	}

	tests := []struct {
		name          string
		remoteAddr    string
		xForwardedFor string
		expected      string
	}{
		{
			name:          "no XFF header",
			remoteAddr:    "198.51.100.1:12345",
			xForwardedFor: "",
			expected:      "198.51.100.1",
		},
		{
			name:          "trusted proxy with single IP",
			remoteAddr:    "192.0.2.1:80",
			xForwardedFor: "198.51.100.1",
			expected:      "198.51.100.1",
		},
		{
			name:          "trusted proxy with multiple IPs",
			remoteAddr:    "192.0.2.1:80",
			xForwardedFor: "198.51.100.1, 203.0.113.5",
			expected:      "198.51.100.1",
		},
		{
			name:          "untrusted proxy",
			remoteAddr:    "198.51.100.1:80",
			xForwardedFor: "10.0.0.1",
			expected:      "198.51.100.1",
		},
		{
			name:          "empty XFF",
			remoteAddr:    "192.0.2.1:80",
			xForwardedFor: "",
			expected:      "192.0.2.1",
		},
		{
			name:          "malformed remote addr",
			remoteAddr:    "192.0.2.1",
			xForwardedFor: "198.51.100.1",
			expected:      "198.51.100.1",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest("GET", "/", nil)
			req.RemoteAddr = tt.remoteAddr
			if tt.xForwardedFor != "" {
				req.Header.Set("X-Forwarded-For", tt.xForwardedFor)
			}

			result := xff.extractRealIP(req)
			if result != tt.expected {
				t.Errorf("extractRealIP() = %s, expected %s", result, tt.expected)
			}
		})
	}
}

func TestGetRealIPWithoutMiddleware(t *testing.T) {
	req := httptest.NewRequest("GET", "/", nil)
	req.RemoteAddr = "198.51.100.1:12345"

	realIP := GetRealIP(req)
	expected := "198.51.100.1"
	if realIP != expected {
		t.Errorf("GetRealIP() = %s, expected %s", realIP, expected)
	}
}