package main

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

func TestThrottleBlocksRepeatedGuesses(t *testing.T) {
	th := newThrottle()

	for i := 0; i < guessBurst; i++ {
		if !th.allow("198.51.100.7") {
			t.Fatalf("guess %d refused inside the burst", i+1)
		}
	}
	if th.allow("198.51.100.7") {
		t.Error("guess allowed past the burst")
	}
	// A different address has its own allowance.
	if !th.allow("198.51.100.8") {
		t.Error("a second address was blocked by the first one's guesses")
	}
}

func TestClientIPIgnoresForwardedHeaderFromDirectClients(t *testing.T) {
	// Connecting straight from the internet: X-Forwarded-For is attacker
	// input, so a forged value must not create a fresh rate-limit bucket.
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.RemoteAddr = "203.0.113.9:44321"
	r.Header.Set("X-Forwarded-For", "1.2.3.4")

	if got := clientIP(r); got != "203.0.113.9" {
		t.Errorf("clientIP = %q, want the real peer 203.0.113.9", got)
	}
}

func TestClientIPTakesLastForwardedEntryBehindProxy(t *testing.T) {
	// Arriving through Traefik on the container network. The proxy appends
	// the address it saw, so the last entry is the trustworthy one and the
	// forged entry in front of it must be ignored.
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.RemoteAddr = "172.18.0.4:53000"
	r.Header.Set("X-Forwarded-For", "1.2.3.4, 198.51.100.22")

	if got := clientIP(r); got != "198.51.100.22" {
		t.Errorf("clientIP = %q, want 198.51.100.22", got)
	}
}

func TestClientIPFallsBackWhenNoForwardedHeader(t *testing.T) {
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.RemoteAddr = "172.18.0.4:53000"

	if got := clientIP(r); got != "172.18.0.4" {
		t.Errorf("clientIP = %q, want 172.18.0.4", got)
	}
}

func TestAuthRateLimitsWrongPasswords(t *testing.T) {
	handler := auth("hunter2", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusTeapot)
	}))

	send := func(pass string) int {
		r := httptest.NewRequest(http.MethodGet, "/", nil)
		r.RemoteAddr = "203.0.113.5:40000"
		r.SetBasicAuth("", pass)
		w := httptest.NewRecorder()
		handler.ServeHTTP(w, r)
		return w.Code
	}

	for i := 0; i < guessBurst; i++ {
		if code := send("wrong"); code != http.StatusUnauthorized {
			t.Fatalf("guess %d returned %d, want 401", i+1, code)
		}
	}
	if code := send("wrong"); code != http.StatusTooManyRequests {
		t.Errorf("guess past the burst returned %d, want 429", code)
	}
}

func TestAuthDoesNotSpendAllowanceOnTheBrowserHandshake(t *testing.T) {
	handler := auth("hunter2", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusTeapot)
	}))

	// Every session opens with a credential-less request. Charging those
	// would lock a family out by simply opening the app a few times.
	for i := 0; i < guessBurst*4; i++ {
		r := httptest.NewRequest(http.MethodGet, "/", nil)
		r.RemoteAddr = "203.0.113.6:40000"
		w := httptest.NewRecorder()
		handler.ServeHTTP(w, r)
		if w.Code != http.StatusUnauthorized {
			t.Fatalf("handshake %d returned %d, want 401", i+1, w.Code)
		}
	}

	// The correct password still works afterwards.
	r := httptest.NewRequest(http.MethodGet, "/", nil)
	r.RemoteAddr = "203.0.113.6:40000"
	r.SetBasicAuth("", "hunter2")
	w := httptest.NewRecorder()
	handler.ServeHTTP(w, r)
	if w.Code != http.StatusTeapot {
		t.Errorf("correct password returned %d, want the wrapped handler", w.Code)
	}
}