go-core/middleware/rate_limit_test.go

package middleware

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

func TestIPRateLimiterRejectsExcessRequestsAcrossAuthPaths(t *testing.T) {
	limiter := NewIPRateLimiter(2, time.Minute, 5*time.Minute)
	h := limiter.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))

	for i := 0; i < 2; i++ {
		req := httptest.NewRequest(http.MethodPost, "/auth/login", nil)
		req.RemoteAddr = "203.0.113.10:12345"
		rr := httptest.NewRecorder()
		h.ServeHTTP(rr, req)
		if rr.Code != http.StatusNoContent {
			t.Fatalf("request %d expected 204, got %d", i+1, rr.Code)
		}
	}

	third := httptest.NewRequest(http.MethodPost, "/auth/login", nil)
	third.RemoteAddr = "203.0.113.10:12345"
	rr := httptest.NewRecorder()
	h.ServeHTTP(rr, third)
	if rr.Code != http.StatusTooManyRequests {
		t.Fatalf("expected 429 for third request, got %d", rr.Code)
	}
	if rr.Header().Get("Retry-After") == "" {
		t.Fatal("expected Retry-After header")
	}

	otherPath := httptest.NewRequest(http.MethodPost, "/auth/register", nil)
	otherPath.RemoteAddr = "203.0.113.10:12345"
	rr2 := httptest.NewRecorder()
	h.ServeHTTP(rr2, otherPath)
	if rr2.Code != http.StatusTooManyRequests {
		t.Fatalf("expected shared limiter bucket per IP across paths, got %d", rr2.Code)
	}
}

func TestIPRateLimiterCleanupRunsPeriodicallyWithoutEntryThreshold(t *testing.T) {
	limiter := NewIPRateLimiter(5, 20*time.Millisecond, 20*time.Millisecond)

	first := httptest.NewRequest(http.MethodPost, "/auth/login", nil)
	first.RemoteAddr = "203.0.113.10:12345"
	if !limiter.allow(first) {
		t.Fatal("expected first request to be allowed")
	}

	time.Sleep(50 * time.Millisecond)

	second := httptest.NewRequest(http.MethodPost, "/auth/login", nil)
	second.RemoteAddr = "203.0.113.11:12345"
	if !limiter.allow(second) {
		t.Fatal("expected second request to be allowed")
	}

	limiter.mu.Lock()
	defer limiter.mu.Unlock()
	if len(limiter.entries) != 1 {
		t.Fatalf("expected stale entries to be cleaned, got %d entries", len(limiter.entries))
	}
	if _, exists := limiter.entries["203.0.113.11"]; !exists {
		t.Fatal("expected current client key to remain after cleanup")
	}
}

func TestIPRateLimiterEvictsOldestEntryAtCapacity(t *testing.T) {
	limiter := NewIPRateLimiter(5, time.Minute, time.Hour)
	limiter.maxEntries = 2

	req1 := httptest.NewRequest(http.MethodPost, "/auth/login", nil)
	req1.RemoteAddr = "203.0.113.10:12345"
	if !limiter.allow(req1) {
		t.Fatal("expected first request to be allowed")
	}

	time.Sleep(2 * time.Millisecond)

	req2 := httptest.NewRequest(http.MethodPost, "/auth/login", nil)
	req2.RemoteAddr = "203.0.113.11:12345"
	if !limiter.allow(req2) {
		t.Fatal("expected second request to be allowed")
	}

	time.Sleep(2 * time.Millisecond)

	req3 := httptest.NewRequest(http.MethodPost, "/auth/login", nil)
	req3.RemoteAddr = "203.0.113.12:12345"
	if !limiter.allow(req3) {
		t.Fatal("expected third request to be allowed")
	}

	limiter.mu.Lock()
	defer limiter.mu.Unlock()
	if len(limiter.entries) != 2 {
		t.Fatalf("expected limiter size to remain capped at 2, got %d", len(limiter.entries))
	}
	if _, exists := limiter.entries["203.0.113.10"]; exists {
		t.Fatal("expected oldest entry to be evicted")
	}
	if _, exists := limiter.entries["203.0.113.11"]; !exists {
		t.Fatal("expected second entry to remain")
	}
	if _, exists := limiter.entries["203.0.113.12"]; !exists {
		t.Fatal("expected newest entry to remain")
	}
}

func TestIPRateLimiterWindowResetAllowsRequestAgain(t *testing.T) {
	limiter := NewIPRateLimiter(1, 20*time.Millisecond, time.Minute)

	req := httptest.NewRequest(http.MethodPost, "/auth/login", nil)
	req.RemoteAddr = "203.0.113.10:12345"
	if !limiter.allow(req) {
		t.Fatal("expected first request to be allowed")
	}
	if limiter.allow(req) {
		t.Fatal("expected second request in same window to be blocked")
	}

	time.Sleep(25 * time.Millisecond)
	if !limiter.allow(req) {
		t.Fatal("expected request to be allowed after window reset")
	}
}

func TestClientKeyVariants(t *testing.T) {
	t.Run("host port", func(t *testing.T) {
		req := httptest.NewRequest(http.MethodGet, "/", nil)
		req.RemoteAddr = "203.0.113.10:4321"
		if got := clientKey(req); got != "203.0.113.10" {
			t.Fatalf("expected parsed host, got %q", got)
		}
	})

	t.Run("raw host", func(t *testing.T) {
		req := httptest.NewRequest(http.MethodGet, "/", nil)
		req.RemoteAddr = "203.0.113.11"
		if got := clientKey(req); got != "203.0.113.11" {
			t.Fatalf("expected raw host, got %q", got)
		}
	})

	t.Run("empty uses unknown", func(t *testing.T) {
		req := httptest.NewRequest(http.MethodGet, "/", nil)
		req.RemoteAddr = "   "
		if got := clientKey(req); got != "unknown" {
			t.Fatalf("expected unknown fallback, got %q", got)
		}
	})
}

func TestNewIPRateLimiterAppliesDefaults(t *testing.T) {
	limiter := NewIPRateLimiter(0, 0, 0)

	if limiter.limit != 60 {
		t.Fatalf("expected default limit 60, got %d", limiter.limit)
	}
	if limiter.window != time.Minute {
		t.Fatalf("expected default window %v, got %v", time.Minute, limiter.window)
	}
	if limiter.ttl != 10*time.Minute {
		t.Fatalf("expected default ttl %v, got %v", 10*time.Minute, limiter.ttl)
	}
	if limiter.maxEntries != defaultRateLimiterMaxEntries {
		t.Fatalf("expected maxEntries %d, got %d", defaultRateLimiterMaxEntries, limiter.maxEntries)
	}
	if limiter.entries == nil {
		t.Fatal("expected entries map to be initialized")
	}
}