package middleware

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

// serveRealIP runs a request from peer with the given headers through RealIP
// and returns the ClientInfo the next handler sees.
func serveRealIP(
	t *testing.T, opts RealIPOptions, peer string, h map[string]string,
) ClientInfo {
	t.Helper()
	var got ClientInfo
	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = Client(r)
	})

	req := httptest.NewRequest("GET", "/", nil)
	req.RemoteAddr = peer
	for k, v := range h {
		req.Header.Set(k, v)
	}

	RealIP(opts)(next).ServeHTTP(httptest.NewRecorder(), req)

	return got
}

var defaultHeaders = []string{"X-Forwarded-For", "Forwarded", "X-Real-IP"}

func TestRealIP(t *testing.T) {
	opts := RealIPOptions{
		Trusted: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")},
		Headers: defaultHeaders,
	}
	// Spoofed left-most entry is ignored; right-most untrusted wins.
	headers := map[string]string{
		"X-Forwarded-For":   "6.6.6.6, 203.0.113.9, 127.0.0.1",
		"X-Forwarded-Proto": "https",
	}
	c := serveRealIP(t, opts, "127.0.0.1:1234", headers)
	if c.IP.String() != "203.0.113.9" || c.Scheme != "https" || !c.ViaProxy {
		t.Fatalf("xff: %+v", c)
	}

	headers = map[string]string{
		"Forwarded": `for="[2001:db8::7]:4711";proto=https`,
	}
	c = serveRealIP(t, opts, "127.0.0.1:1234", headers)
	if c.IP.String() != "2001:db8::7" {
		t.Fatalf("forwarded: %+v", c)
	}
}

func TestUntrustedPeerHeadersIgnored(t *testing.T) {
	opts := RealIPOptions{
		Trusted: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
		Headers: defaultHeaders,
	}
	c := serveRealIP(
		t,
		opts,
		"127.0.0.1:1234",
		map[string]string{"X-Forwarded-For": "203.0.113.9"},
	)
	if c.IP.String() != "127.0.0.1" || c.ViaProxy {
		t.Fatalf("untrusted XFF honoured: %+v", c)
	}
}