147 lines
3.6 KiB
Go
147 lines
3.6 KiB
Go
package httpx
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRateLimiter_WithinLimit(t *testing.T) {
|
|
rl := NewRateLimiter(time.Second, 10)
|
|
key := "test-key"
|
|
|
|
for i := 0; i < 10; i++ {
|
|
if !rl.Allow(key) {
|
|
t.Errorf("request %d should be allowed (within limit)", i+1)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_ExceedLimit(t *testing.T) {
|
|
rl := NewRateLimiter(time.Second, 10)
|
|
key := "test-key"
|
|
|
|
// First 10 requests allowed
|
|
for i := 0; i < 10; i++ {
|
|
rl.Allow(key)
|
|
}
|
|
|
|
// 11th request should be rejected
|
|
if rl.Allow(key) {
|
|
t.Error("11th request should be rejected (exceed limit)")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_DifferentKeys(t *testing.T) {
|
|
rl := NewRateLimiter(time.Second, 10)
|
|
|
|
// Use up all quota for key1
|
|
for i := 0; i < 10; i++ {
|
|
rl.Allow("key1")
|
|
}
|
|
|
|
// key1 should be rejected now
|
|
if rl.Allow("key1") {
|
|
t.Error("key1 should be rejected after exhausting quota")
|
|
}
|
|
|
|
// key2 should still be allowed (different key, independent quota)
|
|
if !rl.Allow("key2") {
|
|
t.Error("key2 should be allowed (different key does not share quota)")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_CleanupOldEntries(t *testing.T) {
|
|
rl := NewRateLimiter(50*time.Millisecond, 5)
|
|
key := "cleanup-key"
|
|
|
|
// Use up all quota
|
|
for i := 0; i < 5; i++ {
|
|
rl.Allow(key)
|
|
}
|
|
|
|
// Verify limit is reached
|
|
if rl.Allow(key) {
|
|
t.Error("should be at limit before cleanup")
|
|
}
|
|
|
|
// Wait for window to expire
|
|
time.Sleep(60 * time.Millisecond)
|
|
|
|
// After window expires, should be allowed again
|
|
if !rl.Allow(key) {
|
|
t.Error("request should be allowed after old entries are cleaned up")
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_WithRateLimit(t *testing.T) {
|
|
rl := NewRateLimiter(time.Second, 2)
|
|
|
|
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusOK)
|
|
})
|
|
|
|
wrapped := rl.WithRateLimit(handler)
|
|
|
|
// First two requests should succeed
|
|
for i := 0; i < 2; i++ {
|
|
req := httptest.NewRequest("GET", "/", nil)
|
|
req.RemoteAddr = "192.168.1.1:1234"
|
|
rec := httptest.NewRecorder()
|
|
wrapped.ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusOK {
|
|
t.Errorf("request %d: expected 200, got %d", i+1, rec.Code)
|
|
}
|
|
}
|
|
|
|
// Third request should be rate limited (429)
|
|
req := httptest.NewRequest("GET", "/", nil)
|
|
req.RemoteAddr = "192.168.1.1:1234"
|
|
rec := httptest.NewRecorder()
|
|
wrapped.ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusTooManyRequests {
|
|
t.Errorf("expected 429, got %d", rec.Code)
|
|
}
|
|
}
|
|
|
|
func TestRateLimiter_WithRateLimit_XForwardedFor(t *testing.T) {
|
|
rl := NewRateLimiter(time.Second, 1)
|
|
|
|
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusOK)
|
|
})
|
|
|
|
wrapped := rl.WithRateLimit(handler)
|
|
|
|
// First request with X-Forwarded-For should succeed
|
|
req := httptest.NewRequest("GET", "/", nil)
|
|
req.RemoteAddr = "192.168.1.1:1234"
|
|
req.Header.Set("X-Forwarded-For", "10.0.0.1")
|
|
rec := httptest.NewRecorder()
|
|
wrapped.ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusOK {
|
|
t.Errorf("first request: expected 200, got %d", rec.Code)
|
|
}
|
|
|
|
// Second request with same IP in X-Forwarded-For should be rejected
|
|
req = httptest.NewRequest("GET", "/", nil)
|
|
req.RemoteAddr = "192.168.1.1:1234"
|
|
req.Header.Set("X-Forwarded-For", "10.0.0.1")
|
|
rec = httptest.NewRecorder()
|
|
wrapped.ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusTooManyRequests {
|
|
t.Errorf("second request: expected 429, got %d", rec.Code)
|
|
}
|
|
|
|
// Different X-Forwarded-For IP should succeed
|
|
req = httptest.NewRequest("GET", "/", nil)
|
|
req.RemoteAddr = "192.168.1.1:1234"
|
|
req.Header.Set("X-Forwarded-For", "10.0.0.2")
|
|
rec = httptest.NewRecorder()
|
|
wrapped.ServeHTTP(rec, req)
|
|
if rec.Code != http.StatusOK {
|
|
t.Errorf("different IP: expected 200, got %d", rec.Code)
|
|
}
|
|
}
|