135 lines
3.7 KiB
Go
135 lines
3.7 KiB
Go
package mock
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestRateLimitWindows(t *testing.T) {
|
|
now := time.Date(2026, 10, 6, 10, 3, 20, 0, time.UTC) // first request of the hour
|
|
rl := NewRateLimiter(12, 5*time.Minute)
|
|
rl.now = func() time.Time { return now }
|
|
|
|
for i := 0; i < 12; i++ {
|
|
if ok, _ := rl.allow("a"); !ok {
|
|
t.Fatalf("request %d rejected", i+1)
|
|
}
|
|
}
|
|
now = now.Add(100 * time.Second) // 10:05:00, window ends 10:08:20
|
|
if ok, secs := rl.allow("a"); ok || secs != 200 {
|
|
t.Fatalf("want reject with 200s, got ok=%v secs=%d", ok, secs)
|
|
}
|
|
if ok, _ := rl.allow("b"); !ok {
|
|
t.Fatal("other client should be unaffected")
|
|
}
|
|
now = time.Date(2026, 10, 6, 10, 8, 20, 0, time.UTC) // next window opens
|
|
if ok, _ := rl.allow("a"); !ok {
|
|
t.Fatal("new window should allow")
|
|
}
|
|
|
|
// Window 10:58:20-11:03:20 is cut short by the new hour at 11:00:00.
|
|
now = time.Date(2026, 10, 6, 10, 58, 30, 0, time.UTC)
|
|
for i := 0; i < 12; i++ {
|
|
rl.allow("a")
|
|
}
|
|
if ok, secs := rl.allow("a"); ok || secs != 90 {
|
|
t.Fatalf("want reject with 90s, got ok=%v secs=%d", ok, secs)
|
|
}
|
|
now = time.Date(2026, 10, 6, 11, 0, 5, 0, time.UTC) // re-anchors here
|
|
if ok, _ := rl.allow("a"); !ok {
|
|
t.Fatal("new hour should allow")
|
|
}
|
|
}
|
|
|
|
func TestClientKey(t *testing.T) {
|
|
cases := map[string]string{
|
|
"192.0.2.1:1234": "192.0.2.1",
|
|
"[::1]:8080": "::1",
|
|
"no-port": "no-port", // not host:port, used as-is
|
|
}
|
|
for addr, want := range cases {
|
|
req := httptest.NewRequest("GET", "/", nil)
|
|
req.RemoteAddr = addr
|
|
if got := clientKey(req); got != want {
|
|
t.Errorf("clientKey(%q) = %q, want %q", addr, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRateLimitMiddleware(t *testing.T) {
|
|
now := time.Date(2026, 10, 6, 10, 0, 0, 0, time.UTC)
|
|
rl := NewRateLimiter(2, 5*time.Minute)
|
|
rl.now = func() time.Time { return now }
|
|
|
|
calls := 0
|
|
h := rl.Middleware(func(w http.ResponseWriter, req *http.Request) {
|
|
calls++
|
|
w.WriteHeader(http.StatusOK)
|
|
})
|
|
do := func(addr string) *httptest.ResponseRecorder {
|
|
req := httptest.NewRequest("GET", "/", nil)
|
|
req.RemoteAddr = addr
|
|
rec := httptest.NewRecorder()
|
|
h(rec, req)
|
|
return rec
|
|
}
|
|
|
|
for i := 0; i < 2; i++ {
|
|
if rec := do("192.0.2.1:1000"); rec.Code != http.StatusOK {
|
|
t.Fatalf("request %d: status %d", i+1, rec.Code)
|
|
}
|
|
}
|
|
// Same host, different port: still the same client.
|
|
rec := do("192.0.2.1:2000")
|
|
if rec.Code != http.StatusTooManyRequests {
|
|
t.Fatalf("status %d, want 429", rec.Code)
|
|
}
|
|
if got := rec.Header().Get("Retry-After"); got != "300" {
|
|
t.Errorf("Retry-After = %q, want 300", got)
|
|
}
|
|
if !strings.Contains(rec.Body.String(), "Rate limit of 2 exceeded") ||
|
|
!strings.Contains(rec.Body.String(), "300 seconds") {
|
|
t.Errorf("body = %q", rec.Body.String())
|
|
}
|
|
if calls != 2 {
|
|
t.Errorf("next called %d times, want 2", calls)
|
|
}
|
|
if rec := do("192.0.2.2:1000"); rec.Code != http.StatusOK {
|
|
t.Errorf("other client: status %d", rec.Code)
|
|
}
|
|
}
|
|
|
|
// Once more than 10000 clients are tracked, a new client in a new hour drops
|
|
// every client left over from earlier hours.
|
|
func TestRateLimitPrunesStaleClients(t *testing.T) {
|
|
now := time.Date(2026, 10, 6, 10, 30, 0, 0, time.UTC)
|
|
rl := NewRateLimiter(1, time.Minute)
|
|
rl.now = func() time.Time { return now }
|
|
|
|
for i := 0; i <= 10000; i++ {
|
|
rl.allow(strconv.Itoa(i))
|
|
}
|
|
if len(rl.clients) != 10001 {
|
|
t.Fatalf("tracking %d clients", len(rl.clients))
|
|
}
|
|
|
|
// Same hour: nothing is stale, so nothing is dropped.
|
|
rl.allow("same-hour")
|
|
if len(rl.clients) != 10002 {
|
|
t.Fatalf("same hour: tracking %d clients, want 10002", len(rl.clients))
|
|
}
|
|
|
|
now = now.Add(time.Hour)
|
|
rl.allow("new")
|
|
if len(rl.clients) != 1 {
|
|
t.Fatalf("after prune: tracking %d clients, want 1", len(rl.clients))
|
|
}
|
|
if _, ok := rl.clients["new"]; !ok {
|
|
t.Fatal("new client was pruned")
|
|
}
|
|
}
|