package api

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

	"git.k8s-git-k6n.net/mats/fancy-dns/internal/config"
)

// rateLimited is the config knob every test in this file turns: the per-minute
// request budget and the per-IP failure lockout are the same struct.
func rateLimited(c *config.Config, perMinute, failures int, lockout time.Duration) {
	modeAPIKey(c)
	c.API.RateLimit = config.AuthRate{PerMinute: perMinute, Failures: failures, Lockout: config.Duration(lockout)}
}

// behindProxy is the X-Forwarded-For header that makes a request look like it
// came from ip. The harness's own connections come from 127.0.0.1, so a test
// that wants two distinct clients has to trust the loopback proxy.
func behindProxy(ip string) string { return "X-Forwarded-For: " + ip }

// badLogin posts a login with a key that is not configured. It is what a
// brute-force attempt looks like from the server's side.
func badLogin(h *harness, header string) response {
	return h.do(http.MethodPost, "/api/auth/login", `{"api_key":"nope"}`, header)
}

// goodLogin posts a login with the key the harness was configured with.
func goodLogin(h *harness, header string) response {
	return h.do(http.MethodPost, "/api/auth/login", `{"api_key":"rw-key"}`, header)
}

// TestLoginBruteForceEventuallyHitsTheRateLimiter is the brute-force case, and
// the reason api.rate_limit exists at all: hammering the login endpoint with
// guessed keys must eventually be refused with 429 rather than answered with an
// unlimited stream of 401s.
//
// The per_minute limit is documented as "requests per minute per IP for the
// login endpoint", and the login route is registered through the identity
// wrapper, so nothing used to ever count these requests.
func TestLoginBruteForceEventuallyHitsTheRateLimiter(t *testing.T) {
	h := newHarness(t, func(c *config.Config) {
		// A generous per-minute budget so the lockout, not the request
		// count, is what stops the attack.
		rateLimited(c, 1000, 3, 5*time.Minute)
	})

	var locked response
	for i := 0; i < 20; i++ {
		res := badLogin(h, "")
		if res.status == http.StatusUnauthorized {
			continue
		}
		locked = res
		break
	}
	if locked.status != http.StatusTooManyRequests {
		t.Fatalf("20 bad logins from one IP never hit the rate limit: last status %d, want 429: %s",
			locked.status, locked.body)
	}
	if locked.header.Get("Retry-After") == "" {
		t.Error("429 from the login limiter carries no Retry-After header")
	}

	// A *correct* key during the lockout must be refused too. Otherwise the
	// lockout is only an obstacle to the attacker and the operator whose key
	// was typo'd gets in anyway.
	res := goodLogin(h, "")
	if res.status != http.StatusTooManyRequests {
		t.Fatalf("a correct key during the lockout = %d, want 429: %s", res.status, res.body)
	}
	if _, ok := res.json(t)["token"]; ok {
		t.Error("the lockout still issued a session token")
	}

	// The lockout is not a login-only affair in this design: the same per-IP
	// state gates the rest of the API, so a brute-force attempt from a host
	// that also runs the console cannot keep hammering other routes either.
	if res := h.get("/api/whoami"); res.status != http.StatusTooManyRequests {
		t.Errorf("an authed route during the lockout = %d, want 429: %s", res.status, res.body)
	}
}

// TestLoginLockoutExpiresWithTheClock checks the lockout is temporary: once the
// configured lockout has passed, the same IP is served again. The clock is the
// server's, so the test advances it instead of sleeping.
func TestLoginLockoutExpiresWithTheClock(t *testing.T) {
	h := newHarness(t, func(c *config.Config) {
		rateLimited(c, 1000, 3, 5*time.Minute)
	})
	base := time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)
	now := base
	h.srv.now = func() time.Time { return now }

	for i := 0; i < 3; i++ {
		if res := badLogin(h, ""); res.status != http.StatusUnauthorized {
			t.Fatalf("bad login %d = %d, want 401: %s", i+1, res.status, res.body)
		}
	}

	// Still inside the lockout window.
	now = base.Add(5*time.Minute - time.Second)
	if res := goodLogin(h, ""); res.status != http.StatusTooManyRequests {
		t.Fatalf("one second before the lockout expires = %d, want 429: %s", res.status, res.body)
	}

	// Past it.
	now = base.Add(5*time.Minute + time.Second)
	res := goodLogin(h, "")
	if res.status != http.StatusOK {
		t.Fatalf("login after the lockout expired = %d, want 200: %s", res.status, res.body)
	}
	if _, ok := res.json(t)["token"]; !ok {
		t.Error("login after the lockout expired returned no token")
	}
}

// TestLoginLockoutIsPerIP checks one attacker cannot lock out everyone else
// behind the same proxy: the state is keyed on the forwarded client address,
// which is only honoured from a trusted proxy.
func TestLoginLockoutIsPerIP(t *testing.T) {
	h := newHarness(t, func(c *config.Config) {
		rateLimited(c, 1000, 3, 5*time.Minute)
		c.API.TrustedProxies = []string{"127.0.0.0/8"}
	})
	const attacker = "198.51.100.7"
	for i := 0; i < 10; i++ {
		if res := badLogin(h, behindProxy(attacker)); res.status == http.StatusOK {
			t.Fatalf("bad login %d was accepted: %s", i+1, res.body)
		}
	}
	if res := badLogin(h, behindProxy(attacker)); res.status != http.StatusTooManyRequests {
		t.Fatalf("the attacker was never locked out: %d: %s", res.status, res.body)
	}
	// A different client behind the same proxy must be untouched.
	res := goodLogin(h, behindProxy("203.0.113.9"))
	if res.status != http.StatusOK {
		t.Fatalf("a second client behind the same proxy = %d, want 200: %s", res.status, res.body)
	}
	// And so must the proxy's own address (no X-Forwarded-For at all).
	if res := goodLogin(h, ""); res.status != http.StatusOK {
		t.Fatalf("the proxy address itself = %d, want 200: %s", res.status, res.body)
	}
}

// TestLoginLockoutDisabledWithoutFailures pins the config side: failures <= 0
// turns the lockout off, so an operator who has not opted in cannot lock
// themselves out of the console.
func TestLoginLockoutDisabledWithoutFailures(t *testing.T) {
	h := newHarness(t, func(c *config.Config) {
		rateLimited(c, 0, 0, 5*time.Minute)
	})
	for i := 0; i < 25; i++ {
		if res := badLogin(h, ""); res.status != http.StatusUnauthorized {
			t.Fatalf("bad login %d = %d, want 401: no lockout is configured: %s", i+1, res.status, res.body)
		}
	}
	if res := goodLogin(h, ""); res.status != http.StatusOK {
		t.Fatalf("login after 25 failures = %d, want 200: %s", res.status, res.body)
	}
}

// TestSuccessfulLoginClearsTheFailureCounter checks a success resets the
// failure counter: an operator who fat-fingers the password three times, waits
// out the lockout and then logs in correctly must not be one typo away from
// another lockout.
func TestSuccessfulLoginClearsTheFailureCounter(t *testing.T) {
	h := newHarness(t, func(c *config.Config) {
		rateLimited(c, 1000, 3, time.Minute)
	})
	base := time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)
	now := base
	h.srv.now = func() time.Time { return now }

	for i := 0; i < 3; i++ {
		badLogin(h, "")
	}
	now = base.Add(2 * time.Minute) // past the one-minute lockout
	if res := goodLogin(h, ""); res.status != http.StatusOK {
		t.Fatalf("login after the lockout = %d, want 200: %s", res.status, res.body)
	}
	// The counter started again: without the reset the first post-success
	// failure would already be locked out. Three more failures are still
	// answered 401 and the fourth re-arms the lockout, which is exactly what
	// "the counter restarted" means.
	for i := 0; i < 3; i++ {
		if res := badLogin(h, ""); res.status != http.StatusUnauthorized {
			t.Fatalf("post-success failure %d = %d, want 401: the counter was not reset", i+1, res.status)
		}
	}
	if res := badLogin(h, ""); res.status != http.StatusTooManyRequests {
		t.Errorf("the fourth post-success failure = %d, want 429: %s", res.status, res.body)
	}
}

// TestSuccessfulLoginDoesNotResetTheRequestWindow pins a deliberate separation:
// a successful login clears the *failure* state, but it must not hand the
// client a fresh per-minute request budget. clearLoginFailures used to delete
// the whole per-IP state, so one successful login in the middle of a burst
// wiped the rate-limit window too — an attacker holding one valid key could use
// it to make brute force free, and the fix is to only reset what the login
// knows about.
func TestSuccessfulLoginDoesNotResetTheRequestWindow(t *testing.T) {
	h := newHarness(t, func(c *config.Config) {
		// Four requests a minute and no lockout at all, so the window is the
		// only thing that can refuse: three failures, then a correct key,
		// then the budget is spent.
		rateLimited(c, 4, 0, 5*time.Minute)
	})
	for i := 0; i < 3; i++ {
		if res := badLogin(h, ""); res.status != http.StatusUnauthorized {
			t.Fatalf("bad login %d = %d, want 401: %s", i+1, res.status, res.body)
		}
	}
	if res := goodLogin(h, ""); res.status != http.StatusOK {
		t.Fatalf("the correct key = %d, want 200: %s", res.status, res.body)
	}
	res := badLogin(h, "")
	if res.status != http.StatusTooManyRequests {
		t.Fatalf("after four requests in the window = %d, want 429: a successful login must not reset the window: %s",
			res.status, res.body)
	}
}

// TestAuthedRouteRateLimit checks the per-minute budget on the authenticated
// routes, which is where allowRequest has always been wired: N+1 requests in a
// window get 429 with a Retry-After, and the Nth is still served.
func TestAuthedRouteRateLimit(t *testing.T) {
	h := newHarness(t, func(c *config.Config) { rateLimited(c, 5, 0, 5*time.Minute) })
	for i := 0; i < 5; i++ {
		if res := h.get("/api/whoami"); res.status != http.StatusOK {
			t.Fatalf("request %d = %d, want 200: %s", i+1, res.status, res.body)
		}
	}
	res := h.get("/api/whoami")
	if res.status != http.StatusTooManyRequests {
		t.Fatalf("the 6th request in a minute = %d, want 429: %s", res.status, res.body)
	}
	if got := res.header.Get("Retry-After"); got == "" {
		t.Error("429 from the per-minute limit carries no Retry-After header")
	}
	// The refused request must not have reached the handler.
	if h.fake.count("Status") != 0 {
		t.Error("a rate-limited request still reached a hook")
	}
}

// TestAuthedRouteRateLimitWindowResets checks the window is a minute wide and
// is driven by the server's clock, not by a timer that has to expire in real
// time.
func TestAuthedRouteRateLimitWindowResets(t *testing.T) {
	h := newHarness(t, func(c *config.Config) { rateLimited(c, 3, 0, 5*time.Minute) })
	base := time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)
	now := base
	h.srv.now = func() time.Time { return now }

	for i := 0; i < 3; i++ {
		if res := h.get("/api/whoami"); res.status != http.StatusOK {
			t.Fatalf("request %d = %d, want 200: %s", i+1, res.status, res.body)
		}
	}
	if res := h.get("/api/whoami"); res.status != http.StatusTooManyRequests {
		t.Fatalf("the 4th request = %d, want 429: %s", res.status, res.body)
	}
	// One second before the window rolls over is still inside it.
	now = base.Add(time.Minute - time.Second)
	if res := h.get("/api/whoami"); res.status != http.StatusTooManyRequests {
		t.Fatalf("one second before the window rolls = %d, want 429: %s", res.status, res.body)
	}
	now = base.Add(time.Minute)
	if res := h.get("/api/whoami"); res.status != http.StatusOK {
		t.Fatalf("first request in the next window = %d, want 200: %s", res.status, res.body)
	}
}

// TestAuthedRouteRateLimitIsPerIP checks the budget is per client address, so
// one noisy operator cannot lock out the rest of the console's users.
func TestAuthedRouteRateLimitIsPerIP(t *testing.T) {
	h := newHarness(t, func(c *config.Config) {
		rateLimited(c, 2, 0, 5*time.Minute)
		c.API.TrustedProxies = []string{"127.0.0.0/8"}
	})
	for i := 0; i < 3; i++ {
		res := h.do(http.MethodGet, "/api/whoami", "", "X-API-Key: rw-key", behindProxy("198.51.100.1"))
		want := http.StatusOK
		if i == 2 {
			want = http.StatusTooManyRequests
		}
		if res.status != want {
			t.Fatalf("request %d from the noisy client = %d, want %d: %s", i+1, res.status, want, res.body)
		}
	}
	if res := h.do(http.MethodGet, "/api/whoami", "", "X-API-Key: rw-key", behindProxy("203.0.113.9")); res.status != http.StatusOK {
		t.Fatalf("a different client = %d, want 200: %s", res.status, res.body)
	}
}

// TestAuthedRouteRateLimitDisabledWhenNotConfigured pins the default: with
// rate_limit unset (which is what a config that never mentions it gets) the
// authenticated routes are not limited at all.
func TestAuthedRouteRateLimitDisabledWhenNotConfigured(t *testing.T) {
	h := newHarness(t, nil) // testConfig leaves RateLimit zero
	for i := 0; i < 20; i++ {
		if res := h.get("/api/whoami"); res.status != http.StatusOK {
			t.Fatalf("request %d with no rate limit configured = %d, want 200: %s", i+1, res.status, res.body)
		}
	}
}

// TestRateLimitCleanupDropsStaleEntries exercises the opportunistic cleanup in
// allowRequest. There is no sweeper goroutine: the per-IP map only shrinks
// when a request comes in and finds it over the 4096-entry cap, so an attacker
// rotating source addresses is the case that matters, and an entry that has
// not been touched for ten minutes is the one that goes.
func TestRateLimitCleanupDropsStaleEntries(t *testing.T) {
	h := newHarness(t, func(c *config.Config) {
		// A budget no test can hit: this is about the map, not the counter.
		rateLimited(c, 1_000_000, 0, time.Minute)
	})
	base := time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)
	now := base
	h.srv.now = func() time.Time { return now }

	// Fill the map past the cap with one request per synthetic client.
	allow := func(ip string) {
		r := httptest.NewRequest(http.MethodGet, "/api/whoami", nil)
		r.RemoteAddr = ip + ":40000"
		if !h.srv.allowRequest(httptest.NewRecorder(), r) {
			t.Fatalf("client %s was rate limited", ip)
		}
	}
	for i := 0; i < 4097; i++ {
		allow(fmt.Sprintf("10.%d.%d.%d", i/65536, (i/256)%256, i%256))
	}
	if got := len(h.srv.limits); got <= 4096 {
		t.Fatalf("the map holds %d entries, want more than the 4096 cap", got)
	}

	// Nothing is stale yet, so a request at the same instant must not shrink
	// the map: the cap is a trigger, not a deadline.
	allow("10.9.9.9")
	if got := len(h.srv.limits); got <= 4096 {
		t.Errorf("the cleanup dropped %d entries that were not stale", 4097+1-got)
	}

	// Past the ten-minute idle horizon every entry is stale, and the next
	// request sweeps them.
	now = base.Add(11 * time.Minute)
	allow("10.9.9.8")
	if got := len(h.srv.limits); got > 2 {
		t.Errorf("after the idle horizon the map holds %d entries, want only the two live clients", got)
	}
}