feat(06-02): add fixed-window limiter, Store, and ClientIP
- MemoryStore mirrors Laravel tooManyAttempts-before-hit first-hit-wins - FixedWindowLimiter + throttle factory; success and 429 rate-limit headers - Trusted-proxy ClientIP; remove noOpLimit; keep Limiter interface seam
This commit is contained in:
Jakub Zych
260
surf/limiter_test.go
Normal file
260
surf/limiter_test.go
Normal file
@@ -0,0 +1,260 @@
|
||||
package surf
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.golem15.com/golem15/summercms/backpack"
|
||||
"git.golem15.com/golem15/summercms/bouncer"
|
||||
"git.golem15.com/golem15/summercms/pact"
|
||||
"git.golem15.com/golem15/summercms/party"
|
||||
)
|
||||
|
||||
func TestFixedWindowLimiterMemoryStoreWindow(t *testing.T) {
|
||||
s := NewMemoryStore(0)
|
||||
decay := 80 * time.Millisecond
|
||||
|
||||
if c := s.Hit("k", decay); c != 1 {
|
||||
t.Fatalf("hit 1 = %d", c)
|
||||
}
|
||||
if c := s.Hit("k", decay); c != 2 {
|
||||
t.Fatalf("hit 2 = %d", c)
|
||||
}
|
||||
if s.TooManyAttempts("k", 3) {
|
||||
t.Fatal("TooManyAttempts at count==max-1")
|
||||
}
|
||||
if c := s.Hit("k", decay); c != 3 {
|
||||
t.Fatalf("hit 3 = %d", c)
|
||||
}
|
||||
if !s.TooManyAttempts("k", 3) {
|
||||
t.Fatal("TooManyAttempts at count==max")
|
||||
}
|
||||
|
||||
time.Sleep(decay + 20*time.Millisecond)
|
||||
if s.TooManyAttempts("k", 3) {
|
||||
t.Fatal("TooManyAttempts after window elapsed")
|
||||
}
|
||||
if c := s.Hit("k", decay); c != 1 {
|
||||
t.Fatalf("fresh window hit = %d", c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedWindowLimiterMemoryStoreFirstHitWins(t *testing.T) {
|
||||
s := NewMemoryStore(0)
|
||||
decay := 200 * time.Millisecond
|
||||
s.Hit("k", decay)
|
||||
time.Sleep(120 * time.Millisecond)
|
||||
s.Hit("k", decay) // must not extend the original window
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
if s.TooManyAttempts("k", 1) {
|
||||
t.Fatal("window was extended by a later hit")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedWindowLimiterSuccessHeaders(t *testing.T) {
|
||||
lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
|
||||
h := lim.Middleware("3,1")(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte(`{"ok":true}`))
|
||||
}))
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d", rec.Code)
|
||||
}
|
||||
if rec.Header().Get("X-RateLimit-Limit") != "3" {
|
||||
t.Fatalf("limit = %q", rec.Header().Get("X-RateLimit-Limit"))
|
||||
}
|
||||
if rec.Header().Get("X-RateLimit-Remaining") != "2" {
|
||||
t.Fatalf("remaining = %q", rec.Header().Get("X-RateLimit-Remaining"))
|
||||
}
|
||||
if rec.Header().Get("Retry-After") != "" || rec.Header().Get("X-RateLimit-Reset") != "" {
|
||||
t.Fatal("Retry-After / X-RateLimit-Reset on success")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedWindowLimiterTooManyAttemptsHeaders(t *testing.T) {
|
||||
lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
|
||||
h := lim.Middleware("1,1")(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
req := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
|
||||
ok := httptest.NewRecorder()
|
||||
h.ServeHTTP(ok, req)
|
||||
if ok.Code != http.StatusOK {
|
||||
t.Fatalf("first status = %d", ok.Code)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d", rec.Code)
|
||||
}
|
||||
if rec.Header().Get("Retry-After") == "" {
|
||||
t.Fatal("missing Retry-After")
|
||||
}
|
||||
if rec.Header().Get("X-RateLimit-Reset") == "" {
|
||||
t.Fatal("missing X-RateLimit-Reset")
|
||||
}
|
||||
if rec.Header().Get("X-RateLimit-Limit") != "1" {
|
||||
t.Fatalf("limit = %q", rec.Header().Get("X-RateLimit-Limit"))
|
||||
}
|
||||
if rec.Header().Get("X-RateLimit-Remaining") != "0" {
|
||||
t.Fatalf("remaining = %q", rec.Header().Get("X-RateLimit-Remaining"))
|
||||
}
|
||||
if rec.Body.String() != `{"message":"Too Many Attempts."}` {
|
||||
t.Fatalf("body = %q", rec.Body.String())
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if payload["message"] != "Too Many Attempts." {
|
||||
t.Fatalf("payload = %v", payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedWindowLimiterStackedBuckets(t *testing.T) {
|
||||
okHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
|
||||
t.Run("exhaust A", func(t *testing.T) {
|
||||
lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
|
||||
if err := lim.RegisterBucket("demo", "bucket-a", Bucket{
|
||||
Max: 1, Decay: time.Minute,
|
||||
Key: func(*http.Request) string { return "a" },
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := lim.RegisterBucket("demo", "bucket-b", Bucket{
|
||||
Max: 100, Decay: time.Minute,
|
||||
Key: func(*http.Request) string { return "b" },
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := stackThrottle(lim, okHandler, "bucket-a", "bucket-b")
|
||||
req := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||||
first := httptest.NewRecorder()
|
||||
h.ServeHTTP(first, req)
|
||||
if first.Code != http.StatusOK {
|
||||
t.Fatalf("first = %d", first.Code)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("exhausted A status = %d", rec.Code)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("exhaust B", func(t *testing.T) {
|
||||
lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
|
||||
if err := lim.RegisterBucket("demo", "bucket-a", Bucket{
|
||||
Max: 100, Decay: time.Minute,
|
||||
Key: func(*http.Request) string { return "a-fresh" },
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := lim.RegisterBucket("demo", "bucket-b", Bucket{
|
||||
Max: 1, Decay: time.Minute,
|
||||
Key: func(*http.Request) string { return "b-only" },
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
h := stackThrottle(lim, okHandler, "bucket-a", "bucket-b")
|
||||
req := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||||
first := httptest.NewRecorder()
|
||||
h.ServeHTTP(first, req)
|
||||
if first.Code != http.StatusOK {
|
||||
t.Fatalf("first = %d", first.Code)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("exhausted B status = %d", rec.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestFixedWindowLimiterInlineThrottleKeys(t *testing.T) {
|
||||
okHandler := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
})
|
||||
|
||||
t.Run("principals differ", func(t *testing.T) {
|
||||
lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
|
||||
h := lim.Middleware("1,1")(okHandler)
|
||||
reqA := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
|
||||
reqA.RemoteAddr = "192.0.2.1:1"
|
||||
reqA = reqA.WithContext(bouncer.WithUser(reqA.Context(), &bouncer.Principal{ID: 1}))
|
||||
reqB := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
|
||||
reqB.RemoteAddr = "192.0.2.1:1"
|
||||
reqB = reqB.WithContext(bouncer.WithUser(reqB.Context(), &bouncer.Principal{ID: 2}))
|
||||
|
||||
a1 := httptest.NewRecorder()
|
||||
h.ServeHTTP(a1, reqA)
|
||||
a2 := httptest.NewRecorder()
|
||||
h.ServeHTTP(a2, reqA)
|
||||
if a2.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("user 1 second status = %d", a2.Code)
|
||||
}
|
||||
b1 := httptest.NewRecorder()
|
||||
h.ServeHTTP(b1, reqB)
|
||||
if b1.Code != http.StatusOK {
|
||||
t.Fatalf("user 2 should have a distinct key, status = %d", b1.Code)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("anonymous same IP+Host", func(t *testing.T) {
|
||||
lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
|
||||
h := lim.Middleware("1,1")(okHandler)
|
||||
req1 := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
|
||||
req1.RemoteAddr = "192.0.2.1:1"
|
||||
req2 := httptest.NewRequest(http.MethodGet, "http://example.test/x", nil)
|
||||
req2.RemoteAddr = "192.0.2.1:9"
|
||||
first := httptest.NewRecorder()
|
||||
h.ServeHTTP(first, req1)
|
||||
if first.Code != http.StatusOK {
|
||||
t.Fatalf("first = %d", first.Code)
|
||||
}
|
||||
second := httptest.NewRecorder()
|
||||
h.ServeHTTP(second, req2)
|
||||
if second.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("anonymous same IP+Host should share key, status = %d", second.Code)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestFixedWindowLimiterUnknownBucketFailsAssemble(t *testing.T) {
|
||||
p := assemblePlugin{id: "golem15.demo", use: []string{"throttle:missing"}}
|
||||
_, err := Assemble(backpack.New(nil), []party.Plugin{p})
|
||||
if err == nil || !strings.Contains(err.Error(), "missing") {
|
||||
t.Fatalf("want unknown throttle in error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedWindowLimiterDuplicateBucket(t *testing.T) {
|
||||
lim := NewFixedWindowLimiter(NewMemoryStore(0), nil)
|
||||
b := Bucket{Max: 1, Decay: time.Minute, Key: func(*http.Request) string { return "k" }}
|
||||
if err := lim.RegisterBucket("one", "shared", b); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
err := lim.RegisterBucket("two", "shared", b)
|
||||
if err == nil || !strings.Contains(err.Error(), "shared") || !strings.Contains(err.Error(), "one") {
|
||||
t.Fatalf("got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func stackThrottle(lim *FixedWindowLimiter, next http.Handler, names ...string) http.Handler {
|
||||
h := next
|
||||
for i := len(names) - 1; i >= 0; i-- {
|
||||
h = lim.Middleware(names[i])(h)
|
||||
}
|
||||
return h
|
||||
}
|
||||
|
||||
var _ pact.Router = (*Router)(nil)
|
||||
Reference in New Issue
Block a user