lightmux-contrib/realip/realip_test.go

package realip

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

const defaultTestRemoteAddr = "192.0.2.1:1234"

func captureRemoteAddr(got *string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		*got = r.RemoteAddr
	})
}

func TestNew(t *testing.T) {
	cases := []struct {
		name    string
		headers map[string]string
		want    string
	}{
		{
			name: "no headers",
			want: defaultTestRemoteAddr,
		},
		{
			name:    "CF-Connecting-IP",
			headers: map[string]string{"CF-Connecting-IP": "203.0.113.5"},
			want:    "203.0.113.5",
		},
		{
			name:    "True-Client-IP",
			headers: map[string]string{"True-Client-IP": "203.0.113.6"},
			want:    "203.0.113.6",
		},
		{
			name:    "X-Real-IP",
			headers: map[string]string{"X-Real-IP": "203.0.113.7"},
			want:    "203.0.113.7",
		},
		{
			name:    "X-Forwarded-For single",
			headers: map[string]string{"X-Forwarded-For": "203.0.113.8"},
			want:    "203.0.113.8",
		},
		{
			name:    "X-Forwarded-For list",
			headers: map[string]string{"X-Forwarded-For": "203.0.113.9, 10.0.0.1, 10.0.0.2"},
			want:    "203.0.113.9",
		},
		{
			name:    "X-Forwarded-For with spaces",
			headers: map[string]string{"X-Forwarded-For": "  203.0.113.10 , 10.0.0.1"},
			want:    "203.0.113.10",
		},
		{
			name: "precedence CF over XFF",
			headers: map[string]string{
				"CF-Connecting-IP": "203.0.113.11",
				"X-Forwarded-For":  "198.51.100.1",
			},
			want: "203.0.113.11",
		},
		{
			name: "invalid then valid",
			headers: map[string]string{
				"CF-Connecting-IP": "not-an-ip",
				"X-Real-IP":        "203.0.113.12",
			},
			want: "203.0.113.12",
		},
		{
			name:    "IPv6",
			headers: map[string]string{"X-Real-IP": "2001:db8::1"},
			want:    "2001:db8::1",
		},
		{
			name:    "all invalid falls through",
			headers: map[string]string{"CF-Connecting-IP": "garbage"},
			want:    defaultTestRemoteAddr,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			var got string
			h := New()(captureRemoteAddr(&got))

			req := httptest.NewRequest(http.MethodGet, "/", nil)
			for k, v := range tc.headers {
				req.Header.Set(k, v)
			}
			h.ServeHTTP(httptest.NewRecorder(), req)

			if got != tc.want {
				t.Errorf("r.RemoteAddr = %q, want %q", got, tc.want)
			}
		})
	}
}

func TestNewDoesNotMutateCallerRequest(t *testing.T) {
	var seen string
	h := New()(captureRemoteAddr(&seen))

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("X-Real-IP", "203.0.113.13")
	original := req.RemoteAddr

	h.ServeHTTP(httptest.NewRecorder(), req)

	if seen != "203.0.113.13" {
		t.Errorf("handler saw r.RemoteAddr = %q, want %q", seen, "203.0.113.13")
	}
	if req.RemoteAddr != original {
		t.Errorf("caller's request was mutated: RemoteAddr = %q, want %q", req.RemoteAddr, original)
	}
}

func TestNewWithTrustedPeer(t *testing.T) {
	var got string
	h := New(netip.MustParsePrefix("10.0.0.0/8"))(captureRemoteAddr(&got))

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = "10.1.2.3:55555"
	req.Header.Set("X-Real-IP", "203.0.113.20")

	h.ServeHTTP(httptest.NewRecorder(), req)

	if got != "203.0.113.20" {
		t.Errorf("trusted peer header not honored: r.RemoteAddr = %q, want %q", got, "203.0.113.20")
	}
}

func TestNewWithUntrustedPeer(t *testing.T) {
	var got string
	h := New(netip.MustParsePrefix("10.0.0.0/8"))(captureRemoteAddr(&got))

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = "203.0.113.99:55555"
	req.Header.Set("X-Real-IP", "198.51.100.1")

	h.ServeHTTP(httptest.NewRecorder(), req)

	if got != "203.0.113.99:55555" {
		t.Errorf("untrusted peer should not have header honored: r.RemoteAddr = %q, want %q", got, "203.0.113.99:55555")
	}
}

func TestNewWithMultiplePrefixesIPv6(t *testing.T) {
	var got string
	h := New(
		netip.MustParsePrefix("10.0.0.0/8"),
		netip.MustParsePrefix("2001:db8::/32"),
	)(captureRemoteAddr(&got))

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = "[2001:db8::abcd]:55555"
	req.Header.Set("X-Real-IP", "203.0.113.30")

	h.ServeHTTP(httptest.NewRecorder(), req)

	if got != "203.0.113.30" {
		t.Errorf("IPv6 peer match should honor header: r.RemoteAddr = %q, want %q", got, "203.0.113.30")
	}
}

func TestNewWithUnparseableRemoteAddr(t *testing.T) {
	var got string
	h := New(netip.MustParsePrefix("10.0.0.0/8"))(captureRemoteAddr(&got))

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.RemoteAddr = "not-an-addr"
	req.Header.Set("CF-Connecting-IP", "203.0.113.50")

	h.ServeHTTP(httptest.NewRecorder(), req)

	if got != "not-an-addr" {
		t.Errorf("unparseable RemoteAddr should pass through unchanged: got %q", got)
	}
}