67 lines
1.5 KiB
Go
67 lines
1.5 KiB
Go
|
package ratelimit
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestLimiterAllowsUpToMax(t *testing.T) {
|
||
|
|
l := New(3, time.Minute)
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
|
req.RemoteAddr = "1.2.3.4:1234"
|
||
|
|
|
||
|
|
for i := 0; i < 3; i++ {
|
||
|
|
if !l.Allow(req) {
|
||
|
|
t.Fatalf("request %d rejected, want allowed", i+1)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if l.Allow(req) {
|
||
|
|
t.Fatal("4th request allowed, want rejected")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestLimiterSeparatesIPs(t *testing.T) {
|
||
|
|
l := New(2, time.Minute)
|
||
|
|
reqA := httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
|
reqA.RemoteAddr = "1.1.1.1:1"
|
||
|
|
reqB := httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
|
reqB.RemoteAddr = "2.2.2.2:2"
|
||
|
|
|
||
|
|
l.Allow(reqA)
|
||
|
|
l.Allow(reqA)
|
||
|
|
if l.Allow(reqA) {
|
||
|
|
t.Fatal("IP A over limit, want rejected")
|
||
|
|
}
|
||
|
|
if !l.Allow(reqB) {
|
||
|
|
t.Fatal("IP B should have its own budget")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestLimiterWindowResets(t *testing.T) {
|
||
|
|
l := New(1, 50*time.Millisecond)
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
|
req.RemoteAddr = "9.9.9.9:9"
|
||
|
|
|
||
|
|
if !l.Allow(req) {
|
||
|
|
t.Fatal("first request rejected")
|
||
|
|
}
|
||
|
|
if l.Allow(req) {
|
||
|
|
t.Fatal("second request within window allowed, want rejected")
|
||
|
|
}
|
||
|
|
time.Sleep(60 * time.Millisecond)
|
||
|
|
if !l.Allow(req) {
|
||
|
|
t.Fatal("request after window reset rejected")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestClientIPPrefersForwarded(t *testing.T) {
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||
|
|
req.RemoteAddr = "10.0.0.5:1234"
|
||
|
|
req.Header.Set("X-Forwarded-For", "203.0.113.9, 10.0.0.1")
|
||
|
|
if got := ClientIP(req); got != "203.0.113.9" {
|
||
|
|
t.Fatalf("ClientIP = %q, want 203.0.113.9", got)
|
||
|
|
}
|
||
|
|
}
|