package idear

import (
	"fmt"
	"net/http/httptest"
	"sync"
	"testing"
	"time"
)

// The limiter's own tests live inside the package because the thing
// worth proving about it — that the table is BOUNDED — is not visible
// from outside: a caller can only observe 429s, and a limiter that
// leaked a bucket per client would answer exactly the same 429s while
// growing until the process died.

func TestLimiterSpendsAndRefills(t *testing.T) {
	now := time.Now()
	l := newLimiter(RateLimit{Burst: 3, Every: time.Second})
	l.now = func() time.Time { return now }

	for i := range 3 {
		if !l.allow("a") {
			t.Fatalf("request %d of a burst of 3 was refused", i+1)
		}
	}
	if l.allow("a") {
		t.Fatal("the fourth request of a burst of 3 was allowed")
	}
	// Another client has their own budget: the limit is per client and
	// not global, or one prober would lock out every invitee.
	if !l.allow("b") {
		t.Fatal("a second client was refused on their first request")
	}

	// One token comes back per Every, and no more than Burst ever
	// accumulates.
	now = now.Add(time.Second)
	if !l.allow("a") {
		t.Fatal("no token came back after one refill interval")
	}
	if l.allow("a") {
		t.Fatal("more than one token came back in one interval")
	}
	now = now.Add(time.Hour)
	for i := range 3 {
		if !l.allow("a") {
			t.Fatalf("request %d after a long idle was refused", i+1)
		}
	}
	if l.allow("a") {
		t.Fatal("an idle client accumulated more than Burst tokens")
	}
}

func TestLimiterTableIsBounded(t *testing.T) {
	now := time.Now()
	const max = 8
	l := newLimiter(RateLimit{Burst: 2, Every: time.Second, Max: max})
	l.now = func() time.Time { return now }

	// Far more clients than the table may hold, each spending their
	// whole burst so none of them is sweepable.
	for i := range max * 10 {
		key := fmt.Sprintf("client-%d", i)
		l.allow(key)
		l.allow(key)
		if got := l.size(); got > max {
			t.Fatalf("the table holds %d buckets, above the bound of %d", got, max)
		}
	}
	if got := l.size(); got != max {
		t.Fatalf("the table holds %d buckets, want it filled to %d", got, max)
	}

	// A full table FAILS CLOSED for an unseen client — see RateLimit.Max.
	if l.allow("someone-new") {
		t.Error("a full table admitted an unseen client; it must fail closed")
	}
	// And a client already in the table is unaffected by the crowd.
	if l.allow("client-0") {
		t.Error("client-0 had spent its burst and was allowed anyway")
	}

	// Once the crowd's buckets refill they are swept, and the table
	// takes new clients again — the bound is on live clients, not a
	// permanent cap on how many the process may ever see.
	now = now.Add(time.Hour)
	if !l.allow("someone-new") {
		t.Error("a swept table still refused a new client")
	}
	if got := l.size(); got > max {
		t.Fatalf("the table holds %d buckets after a sweep", got)
	}
}

func TestLimiterIsConcurrencySafe(t *testing.T) {
	l := newLimiter(RateLimit{Burst: 1000, Every: time.Hour, Max: 16})
	var wg sync.WaitGroup
	for i := range 8 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := range 50 {
				l.allow(fmt.Sprintf("client-%d", (i+j)%16))
			}
		}()
	}
	wg.Wait()
	if got := l.size(); got > 16 {
		t.Fatalf("the table holds %d buckets, above the bound of 16", got)
	}
}

func TestClientIPKeysByNetworkNotAddress(t *testing.T) {
	key := func(remote string) string {
		r := httptest.NewRequest("GET", "/invitations/x", nil)
		r.RemoteAddr = remote
		return clientIP(r)
	}

	// IPv4: the address, without the port. A browser making a second
	// request from a new source port must share the first one's
	// budget, or the limit is no limit at all.
	if got := key("203.0.113.9:51234"); got != "203.0.113.9" {
		t.Errorf("clientIP = %q, want the address without the port", got)
	}
	if got := key("203.0.113.9:51235"); got != "203.0.113.9" {
		t.Errorf("clientIP = %q for a second port, want the same key", got)
	}
	// ...and a DIFFERENT IPv4 host is a different client. Folding v4
	// further would put a whole CGNAT behind one bucket.
	if key("203.0.113.9:1") == key("203.0.113.10:1") {
		t.Error("two IPv4 hosts share a key; the fold is too coarse")
	}

	// IPv6: the /64, because every host is handed one and can rotate
	// addresses inside it for free. Three addresses in one /64 are
	// one client.
	sixtyFour := key("[2001:db8:1:2::5]:443")
	if sixtyFour != "2001:db8:1:2::/64" {
		t.Errorf("clientIP = %q, want the /64 prefix", sixtyFour)
	}
	for _, other := range []string{"[2001:db8:1:2:ffff::9]:443", "[2001:db8:1:2:dead:beef:cafe:1]:80"} {
		if got := key(other); got != sixtyFour {
			t.Errorf("clientIP(%s) = %q, want the same /64 key %q", other, got, sixtyFour)
		}
	}
	// A different /64 is a different client, so a shared limit does
	// not fall out of the fold either.
	if got := key("[2001:db8:1:3::5]:443"); got == sixtyFour {
		t.Errorf("a neighbouring /64 shares the key %q", got)
	}
	// An IPv4-mapped address is unmapped first: one client must not
	// hold two budgets by switching representation.
	if got := key("[::ffff:203.0.113.9]:80"); got != "203.0.113.9" {
		t.Errorf("clientIP = %q for an IPv4-mapped address, want %q", got, "203.0.113.9")
	}
	// The zone is the local interface, not the client.
	if key("[fe80::1%eth0]:80") != key("[fe80::1%eth1]:80") {
		t.Error("the interface zone splits one client's budget in two")
	}
	// An address with no port at all (a unix socket, a test) is used
	// whole rather than dropped, so it still keys to something.
	if got := key("@"); got != "@" {
		t.Errorf("clientIP = %q, want the raw RemoteAddr when it has no port", got)
	}
}

// TestClientIPIgnoresForwardingHeaders pins the claim clientIP's doc
// comment makes outright: a header idear cannot verify is a header an
// attacker can spoof to mint unlimited budgets, so idear never reads
// one. Without this test, teaching clientIP to prefer X-Forwarded-For
// left the whole suite green while handing every prober an unlimited
// supply of fresh buckets.
func TestClientIPIgnoresForwardingHeaders(t *testing.T) {
	r := httptest.NewRequest("GET", "/invitations/x", nil)
	r.RemoteAddr = "203.0.113.9:51234"
	for _, header := range []string{"X-Forwarded-For", "X-Real-IP", "Forwarded", "CF-Connecting-IP", "True-Client-IP"} {
		r.Header.Set(header, "1.2.3.4")
	}
	if got := clientIP(r); got != "203.0.113.9" {
		t.Errorf("clientIP = %q with forwarding headers set, want the RemoteAddr host %q", got, "203.0.113.9")
	}
}

// TestIPv6RotationCannotLockOutOtherClients is the scenario the /64
// fold exists for, run end to end against the limiter.
//
// One machine rotates source addresses inside its own /64 — free, on
// any IPv6 host — and spends far more than the table can hold. Keyed
// on the bare address that fills the table, and allow then fails
// closed for every UNSEEN client: invitees, and orphans, whose only
// healing path is POST /invitations/{token}. Failing closed is right;
// the key was the bug.
func TestIPv6RotationCannotLockOutOtherClients(t *testing.T) {
	const max = 8
	l := newLimiter(RateLimit{Burst: 2, Every: time.Hour, Max: max})

	key := func(remote string) string {
		r := httptest.NewRequest("POST", "/invitations/x", nil)
		r.RemoteAddr = remote
		return clientIP(r)
	}

	for i := range max * 20 {
		l.allow(key(fmt.Sprintf("[2001:db8:1:2::%x]:443", i)))
	}
	if got := l.size(); got != 1 {
		t.Fatalf("one attacker's /64 filled %d buckets, want 1", got)
	}
	// The attacker is throttled on the budget their whole /64 shares.
	if l.allow(key("[2001:db8:1:2::ffff]:443")) {
		t.Error("an address rotation inside one /64 bought a fresh burst")
	}
	// And everybody else is still served — including the orphan whose
	// only way back into the instance is the route being defended.
	for _, invitee := range []string{"198.51.100.7:1", "[2001:db8:9:9::1]:443", "[2001:db8:aa::5]:443"} {
		if !l.allow(key(invitee)) {
			t.Errorf("a real client (%s) was locked out by one machine's address rotation", invitee)
		}
	}
}