106 lines
3.6 KiB
Go
106 lines
3.6 KiB
Go
package test
|
|
|
|
import (
|
|
"RenewCMS/api/middleware"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func newConnectionKeyLimiter(trustedProxies []string, maxAttempts int, window time.Duration) http.Handler {
|
|
limited := middleware.NewRateLimiter(maxAttempts, window)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
return middleware.TrustedProxyRealIP(trustedProxies)(limited)
|
|
}
|
|
|
|
func sendFromConnection(handler http.Handler, remoteAddr string, realIP string) int {
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/auth/login", nil)
|
|
req.RemoteAddr = remoteAddr
|
|
if realIP != "" {
|
|
req.Header.Set("X-Real-IP", realIP)
|
|
}
|
|
rec := httptest.NewRecorder()
|
|
handler.ServeHTTP(rec, req)
|
|
return rec.Code
|
|
}
|
|
|
|
func TestRateLimitKeysOnTheConnectionHost(t *testing.T) {
|
|
t.Run("New connections from the same address share one counter whatever their source port", func(t *testing.T) {
|
|
handler := newConnectionKeyLimiter(nil, 5, time.Hour)
|
|
|
|
blocked := 0
|
|
for port := 40000; port < 40100; port++ {
|
|
if sendFromConnection(handler, fmt.Sprintf("203.0.113.7:%d", port), "") == http.StatusTooManyRequests {
|
|
blocked++
|
|
}
|
|
}
|
|
|
|
if blocked != 95 {
|
|
t.Errorf("expected 95 of 100 attempts blocked, got %d", blocked)
|
|
}
|
|
})
|
|
|
|
t.Run("IPv6 connections from the same address share one counter", func(t *testing.T) {
|
|
handler := newConnectionKeyLimiter(nil, 1, time.Hour)
|
|
|
|
first := sendFromConnection(handler, "[2001:db8::1]:50000", "")
|
|
second := sendFromConnection(handler, "[2001:db8::1]:50001", "")
|
|
|
|
if first != http.StatusOK || second != http.StatusTooManyRequests {
|
|
t.Errorf("expected 200 then 429, got %d then %d", first, second)
|
|
}
|
|
})
|
|
|
|
t.Run("Different addresses keep independent counters", func(t *testing.T) {
|
|
handler := newConnectionKeyLimiter(nil, 1, time.Hour)
|
|
|
|
first := sendFromConnection(handler, "203.0.113.7:1", "")
|
|
second := sendFromConnection(handler, "203.0.113.8:1", "")
|
|
|
|
if first != http.StatusOK || second != http.StatusOK {
|
|
t.Errorf("expected both 200, got %d and %d", first, second)
|
|
}
|
|
})
|
|
|
|
t.Run("A forwarded client keeps one counter across the proxy's connections", func(t *testing.T) {
|
|
handler := newConnectionKeyLimiter([]string{"10.0.0.1"}, 1, time.Hour)
|
|
|
|
first := sendFromConnection(handler, "10.0.0.1:51000", "198.51.100.9")
|
|
second := sendFromConnection(handler, "10.0.0.1:51001", "198.51.100.9")
|
|
|
|
if first != http.StatusOK || second != http.StatusTooManyRequests {
|
|
t.Errorf("expected 200 then 429, got %d then %d", first, second)
|
|
}
|
|
})
|
|
|
|
t.Run("A counter starts over once its window has passed", func(t *testing.T) {
|
|
handler := newConnectionKeyLimiter(nil, 1, 20*time.Millisecond)
|
|
|
|
first := sendFromConnection(handler, "203.0.113.7:1", "")
|
|
second := sendFromConnection(handler, "203.0.113.7:2", "")
|
|
time.Sleep(30 * time.Millisecond)
|
|
third := sendFromConnection(handler, "203.0.113.7:3", "")
|
|
|
|
if first != http.StatusOK || second != http.StatusTooManyRequests || third != http.StatusOK {
|
|
t.Errorf("expected 200, 429, 200, got %d, %d, %d", first, second, third)
|
|
}
|
|
})
|
|
|
|
t.Run("Sweeping expired counters never resets a counter whose window is still open", func(t *testing.T) {
|
|
handler := newConnectionKeyLimiter(nil, 1, 60*time.Millisecond)
|
|
|
|
_ = sendFromConnection(handler, "198.51.100.1:1", "")
|
|
time.Sleep(40 * time.Millisecond)
|
|
first := sendFromConnection(handler, "203.0.113.7:1", "")
|
|
time.Sleep(30 * time.Millisecond)
|
|
second := sendFromConnection(handler, "203.0.113.7:2", "")
|
|
|
|
if first != http.StatusOK || second != http.StatusTooManyRequests {
|
|
t.Errorf("expected 200 then 429 across a sweep, got %d then %d", first, second)
|
|
}
|
|
})
|
|
}
|