diff --git a/internal/api/ratelimit_test.go b/internal/api/ratelimit_test.go new file mode 100644 index 0000000..5f8dfd0 --- /dev/null +++ b/internal/api/ratelimit_test.go @@ -0,0 +1,433 @@ +package api + +import ( + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" +) + +func TestNewRateLimiter(t *testing.T) { + rl := NewRateLimiter(5, time.Minute) + defer rl.Stop() + + if rl == nil { + t.Fatal("NewRateLimiter returned nil") + } + if rl.limit != 5 { + t.Errorf("limit = %d, want 5", rl.limit) + } + if rl.window != time.Minute { + t.Errorf("window = %v, want %v", rl.window, time.Minute) + } + if rl.attempts == nil { + t.Error("attempts map is nil") + } +} + +func TestRateLimiter_Allow_Basic(t *testing.T) { + rl := NewRateLimiter(3, time.Minute) + defer rl.Stop() + + ip := "192.168.1.1" + + // First 3 attempts should be allowed + for i := 1; i <= 3; i++ { + if !rl.Allow(ip) { + t.Errorf("attempt %d should be allowed", i) + } + } + + // 4th attempt should be denied + if rl.Allow(ip) { + t.Error("4th attempt should be denied") + } + + // 5th attempt should also be denied + if rl.Allow(ip) { + t.Error("5th attempt should be denied") + } +} + +func TestRateLimiter_Allow_DifferentIPs(t *testing.T) { + rl := NewRateLimiter(2, time.Minute) + defer rl.Stop() + + ip1 := "192.168.1.1" + ip2 := "192.168.1.2" + + // Both IPs should have independent limits + if !rl.Allow(ip1) { + t.Error("ip1 attempt 1 should be allowed") + } + if !rl.Allow(ip2) { + t.Error("ip2 attempt 1 should be allowed") + } + if !rl.Allow(ip1) { + t.Error("ip1 attempt 2 should be allowed") + } + if !rl.Allow(ip2) { + t.Error("ip2 attempt 2 should be allowed") + } + + // Both should now be at their limit + if rl.Allow(ip1) { + t.Error("ip1 attempt 3 should be denied") + } + if rl.Allow(ip2) { + t.Error("ip2 attempt 3 should be denied") + } +} + +func TestRateLimiter_Allow_WindowExpiry(t *testing.T) { + // Use a very short window for testing + rl := NewRateLimiter(2, 50*time.Millisecond) + defer rl.Stop() + + ip := "192.168.1.1" + + // Use up the limit + if !rl.Allow(ip) { + t.Error("attempt 1 should be allowed") + } + if !rl.Allow(ip) { + t.Error("attempt 2 should be allowed") + } + if rl.Allow(ip) { + t.Error("attempt 3 should be denied") + } + + // Wait for window to expire + time.Sleep(60 * time.Millisecond) + + // Should be allowed again + if !rl.Allow(ip) { + t.Error("attempt after window expiry should be allowed") + } +} + +func TestRateLimiter_Allow_SlidingWindow(t *testing.T) { + // Test that the window is truly sliding (not fixed intervals) + rl := NewRateLimiter(2, 100*time.Millisecond) + defer rl.Stop() + + ip := "192.168.1.1" + + // First attempt + if !rl.Allow(ip) { + t.Error("attempt 1 should be allowed") + } + + // Wait half the window + time.Sleep(60 * time.Millisecond) + + // Second attempt + if !rl.Allow(ip) { + t.Error("attempt 2 should be allowed") + } + + // Third attempt should be denied (both still in window) + if rl.Allow(ip) { + t.Error("attempt 3 should be denied") + } + + // Wait for first attempt to expire (another 50ms should be enough) + time.Sleep(50 * time.Millisecond) + + // Now should be allowed (first attempt expired, second still valid) + if !rl.Allow(ip) { + t.Error("attempt 4 should be allowed after first expires") + } +} + +func TestRateLimiter_Allow_ZeroLimit(t *testing.T) { + rl := NewRateLimiter(0, time.Minute) + defer rl.Stop() + + // With limit 0, nothing should be allowed + if rl.Allow("192.168.1.1") { + t.Error("should deny with limit 0") + } +} + +func TestRateLimiter_Allow_LargeLimit(t *testing.T) { + rl := NewRateLimiter(1000, time.Minute) + defer rl.Stop() + + ip := "192.168.1.1" + + // All 1000 attempts should be allowed + for i := 1; i <= 1000; i++ { + if !rl.Allow(ip) { + t.Errorf("attempt %d should be allowed", i) + } + } + + // 1001st should be denied + if rl.Allow(ip) { + t.Error("attempt 1001 should be denied") + } +} + +func TestRateLimiter_Allow_EmptyIP(t *testing.T) { + rl := NewRateLimiter(2, time.Minute) + defer rl.Stop() + + // Empty string IP should still work as a valid key + if !rl.Allow("") { + t.Error("empty IP attempt 1 should be allowed") + } + if !rl.Allow("") { + t.Error("empty IP attempt 2 should be allowed") + } + if rl.Allow("") { + t.Error("empty IP attempt 3 should be denied") + } +} + +func TestRateLimiter_Cleanup(t *testing.T) { + rl := NewRateLimiter(5, 50*time.Millisecond) + defer rl.Stop() + + // Add some attempts + rl.Allow("192.168.1.1") + rl.Allow("192.168.1.2") + rl.Allow("192.168.1.3") + + // Verify attempts are tracked + rl.mu.RLock() + beforeCount := len(rl.attempts) + rl.mu.RUnlock() + + if beforeCount != 3 { + t.Errorf("expected 3 IPs tracked, got %d", beforeCount) + } + + // Wait for window to expire + time.Sleep(60 * time.Millisecond) + + // Manually trigger cleanup + rl.cleanup() + + // All should be cleaned up + rl.mu.RLock() + afterCount := len(rl.attempts) + rl.mu.RUnlock() + + if afterCount != 0 { + t.Errorf("expected 0 IPs after cleanup, got %d", afterCount) + } +} + +func TestRateLimiter_Cleanup_PartialExpiry(t *testing.T) { + rl := NewRateLimiter(5, 50*time.Millisecond) + defer rl.Stop() + + // Add attempt for IP1 + rl.Allow("192.168.1.1") + + // Wait for it to expire + time.Sleep(60 * time.Millisecond) + + // Add attempt for IP2 (this one is fresh) + rl.Allow("192.168.1.2") + + // Cleanup should remove IP1 but keep IP2 + rl.cleanup() + + rl.mu.RLock() + _, hasIP1 := rl.attempts["192.168.1.1"] + _, hasIP2 := rl.attempts["192.168.1.2"] + rl.mu.RUnlock() + + if hasIP1 { + t.Error("IP1 should have been cleaned up") + } + if !hasIP2 { + t.Error("IP2 should still be present") + } +} + +func TestRateLimiter_Concurrent(t *testing.T) { + rl := NewRateLimiter(100, time.Minute) + defer rl.Stop() + + var wg sync.WaitGroup + numGoroutines := 10 + attemptsPerGoroutine := 20 + + // Track results per goroutine + results := make([][]bool, numGoroutines) + for i := 0; i < numGoroutines; i++ { + results[i] = make([]bool, attemptsPerGoroutine) + } + + // Concurrent access from same IP + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(idx int) { + defer wg.Done() + for j := 0; j < attemptsPerGoroutine; j++ { + results[idx][j] = rl.Allow("shared-ip") + } + }(i) + } + + wg.Wait() + + // Count allowed attempts + allowed := 0 + for i := 0; i < numGoroutines; i++ { + for j := 0; j < attemptsPerGoroutine; j++ { + if results[i][j] { + allowed++ + } + } + } + + // Should have exactly 100 allowed (the limit) + if allowed != 100 { + t.Errorf("expected exactly 100 allowed attempts, got %d", allowed) + } +} + +func TestRateLimiter_Stop(t *testing.T) { + rl := NewRateLimiter(5, time.Minute) + + // Should be able to stop without issues + rl.Stop() + + // After stop, Allow should still work (stop only affects cleanup goroutine) + if !rl.Allow("192.168.1.1") { + t.Error("Allow should still work after Stop") + } +} + +func TestRateLimiter_Middleware_Allowed(t *testing.T) { + rl := NewRateLimiter(5, time.Minute) + defer rl.Stop() + + handlerCalled := false + handler := func(w http.ResponseWriter, r *http.Request) { + handlerCalled = true + w.WriteHeader(http.StatusOK) + } + + wrapped := rl.Middleware(handler) + + req := httptest.NewRequest(http.MethodGet, "/test", nil) + req.RemoteAddr = "192.168.1.1:12345" + w := httptest.NewRecorder() + + wrapped(w, req) + + if !handlerCalled { + t.Error("handler should have been called") + } + if w.Code != http.StatusOK { + t.Errorf("status = %d, want %d", w.Code, http.StatusOK) + } +} + +func TestRateLimiter_Middleware_Denied(t *testing.T) { + rl := NewRateLimiter(1, time.Minute) + defer rl.Stop() + + handlerCalls := 0 + handler := func(w http.ResponseWriter, r *http.Request) { + handlerCalls++ + w.WriteHeader(http.StatusOK) + } + + wrapped := rl.Middleware(handler) + + // First request should succeed + req1 := httptest.NewRequest(http.MethodGet, "/test", nil) + req1.RemoteAddr = "192.168.1.1:12345" + w1 := httptest.NewRecorder() + wrapped(w1, req1) + + if w1.Code != http.StatusOK { + t.Errorf("first request status = %d, want %d", w1.Code, http.StatusOK) + } + + // Second request should be rate limited + req2 := httptest.NewRequest(http.MethodGet, "/test", nil) + req2.RemoteAddr = "192.168.1.1:12345" + w2 := httptest.NewRecorder() + wrapped(w2, req2) + + if w2.Code != http.StatusTooManyRequests { + t.Errorf("second request status = %d, want %d", w2.Code, http.StatusTooManyRequests) + } + + if handlerCalls != 1 { + t.Errorf("handler called %d times, want 1", handlerCalls) + } +} + +func TestRateLimiter_Middleware_XForwardedFor(t *testing.T) { + rl := NewRateLimiter(1, time.Minute) + defer rl.Stop() + + handlerCalls := 0 + handler := func(w http.ResponseWriter, r *http.Request) { + handlerCalls++ + } + + wrapped := rl.Middleware(handler) + + // Request with X-Forwarded-For header + req1 := httptest.NewRequest(http.MethodGet, "/test", nil) + req1.RemoteAddr = "127.0.0.1:12345" + req1.Header.Set("X-Forwarded-For", "203.0.113.1") + w1 := httptest.NewRecorder() + wrapped(w1, req1) + + // Second request with same X-Forwarded-For should be rate limited + req2 := httptest.NewRequest(http.MethodGet, "/test", nil) + req2.RemoteAddr = "127.0.0.1:12345" + req2.Header.Set("X-Forwarded-For", "203.0.113.1") + w2 := httptest.NewRecorder() + wrapped(w2, req2) + + if w2.Code != http.StatusTooManyRequests { + t.Errorf("second request with same X-Forwarded-For should be rate limited, got %d", w2.Code) + } + + // Request with different X-Forwarded-For should be allowed + req3 := httptest.NewRequest(http.MethodGet, "/test", nil) + req3.RemoteAddr = "127.0.0.1:12345" + req3.Header.Set("X-Forwarded-For", "203.0.113.2") + w3 := httptest.NewRecorder() + wrapped(w3, req3) + + if w3.Code == http.StatusTooManyRequests { + t.Error("request with different X-Forwarded-For should be allowed") + } + + if handlerCalls != 2 { + t.Errorf("handler called %d times, want 2", handlerCalls) + } +} + +func TestRateLimiter_Middleware_ResponseBody(t *testing.T) { + rl := NewRateLimiter(0, time.Minute) // Immediately rate limit + defer rl.Stop() + + handler := func(w http.ResponseWriter, r *http.Request) {} + wrapped := rl.Middleware(handler) + + req := httptest.NewRequest(http.MethodGet, "/test", nil) + req.RemoteAddr = "192.168.1.1:12345" + w := httptest.NewRecorder() + + wrapped(w, req) + + body := w.Body.String() + expected := "Rate limit exceeded. Please try again later.\n" + if body != expected { + t.Errorf("body = %q, want %q", body, expected) + } +}