ModelRouter/internal/gateway/login_guard_test.go

package gateway

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestLoginBruteForceLockout(t *testing.T) {
	g := newTestGateway(t)
	body := `{"key":"wrong-key-XXXX"}`

	// 5 failures allowed (each answered 401), 6th+ locked out
	var gotLocked bool
	for i := 1; i <= 8; i++ {
		req := httptest.NewRequest("POST", "/api/login", strings.NewReader(body))
		rr := httptest.NewRecorder()
		g.handleLoginAPI(rr, req)
		if rr.Code != http.StatusUnauthorized {
			t.Fatalf("attempt %d: status=%d want 401", i, rr.Code)
		}
		if i >= 6 {
			// after lockout the response must be identical (no info leak)
			if !strings.Contains(rr.Body.String(), "invalid gateway api key") {
				t.Fatalf("lockout reply changed body: %s", rr.Body.String())
			}
			gotLocked = true
		}
	}
	if !gotLocked {
		t.Fatal("expected lockout by attempt 6")
	}

	// correct key from the SAME ip must also be rejected while locked
	okBody := `{"key":"sk-test"}`
	req := httptest.NewRequest("POST", "/api/login", strings.NewReader(okBody))
	rr := httptest.NewRecorder()
	g.handleLoginAPI(rr, req)
	if rr.Code != http.StatusUnauthorized {
		t.Fatalf("locked-out valid key: status=%d want 401", rr.Code)
	}

	// a different IP is unaffected
	req = httptest.NewRequest("POST", "/api/login", strings.NewReader(okBody))
	req.RemoteAddr = "10.9.9.9:5555"
	rr = httptest.NewRecorder()
	g.handleLoginAPI(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("other IP valid key: status=%d want 200 body=%s", rr.Code, rr.Body.String())
	}
	var resp map[string]interface{}
	_ = json.Unmarshal(rr.Body.Bytes(), &resp)
	if resp["ok"] != true {
		t.Fatalf("login ok flag missing: %s", rr.Body.String())
	}
}

func TestLoginSuccessResetsFailures(t *testing.T) {
	g := newTestGateway(t)
	mkBad := func() *http.Request { return httptest.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"nope"}`)) }
	for i := 0; i < loginMaxFails-1; i++ { // one below the lockout threshold
		rr := httptest.NewRecorder()
		g.handleLoginAPI(rr, mkBad())
		if rr.Code != http.StatusUnauthorized {
			t.Fatal("setup failure")
		}
	}
	// success clears the counter
	rr := httptest.NewRecorder()
	g.handleLoginAPI(rr, httptest.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"sk-test"}`)))
	if rr.Code != http.StatusOK {
		t.Fatalf("valid login failed: %d", rr.Code)
	}
	// now up to loginMaxFails-1 more failures are still accepted before lock
	for i := 0; i < loginMaxFails-1; i++ {
		last := httptest.NewRecorder()
		g.handleLoginAPI(last, mkBad())
	}
	last := httptest.NewRecorder()
	g.handleLoginAPI(last, mkBad())
	// count is at loginMaxFails-1 failures since reset; this next one reaches
	// the threshold but the attempt itself is still evaluated normally (401),
	// and the FOLLOWING one must be locked.
	if last.Code != http.StatusUnauthorized {
		t.Fatalf("threshold attempt should still answer 401: %d", last.Code)
	}
	next := httptest.NewRecorder()
	g.handleLoginAPI(next, mkBad())
	if next.Code != http.StatusUnauthorized {
		t.Fatal("want 401")
	}
	// verify lock actually engaged via the guard state
	g.loginMu.Lock()
	f := g.loginFails["192.0.2.1"]
	engaged := f != nil && f.count >= loginMaxFails && f.until > 0
	g.loginMu.Unlock()
	if !engaged {
		t.Fatal("expected engaged lockout state")
	}
}

func TestClientIPExtraction(t *testing.T) {
	r := httptest.NewRequest("POST", "/api/login", nil)
	r.RemoteAddr = "203.0.113.7:12345"
	if got := clientIP(r); got != "203.0.113.7" {
		t.Errorf("clientIP=%q want 203.0.113.7", got)
	}
	r.Header.Set("X-Real-IP", "198.51.100.9")
	if got := clientIP(r); got != "198.51.100.9" {
		t.Errorf("clientIP with X-Real-IP=%q want proxy value", got)
	}
}