Aleksey Shakhmatov
2d4a06e715
Key on RemoteAddr by default; honor X-Forwarded-For only when the peer is a configured trusted proxy (WithTrustedProxies), walking right-to-left to the first untrusted hop. This closes a trivial rate-limit bypass and the matching unbounded-bucket DoS via spoofed headers. Add WithMaxKeys with opportunistic eviction of idle (fully-refilled) buckets to bound memory. Drop the hand-rolled indexOf in favor of stdlib.
63 lines
1.8 KiB
Go
63 lines
1.8 KiB
Go
package server
|
|
|
|
import (
|
|
"net/http"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.codelab.vc/pkg/httpx/internal/clock"
|
|
)
|
|
|
|
func newTestRequest(remoteAddr, xff string) *http.Request {
|
|
r := &http.Request{RemoteAddr: remoteAddr, Header: http.Header{}}
|
|
if xff != "" {
|
|
r.Header.Set("X-Forwarded-For", xff)
|
|
}
|
|
return r
|
|
}
|
|
|
|
// TestLimiterSweepEvictsIdleBuckets verifies that sweep removes fully-refilled
|
|
// (idle) buckets while preserving buckets that still hold an active limit, so
|
|
// memory is bounded without resetting live clients' allowances.
|
|
func TestLimiterSweepEvictsIdleBuckets(t *testing.T) {
|
|
clk := clock.Mock(time.Now())
|
|
o := &rateLimitOptions{rate: 1, burst: 5, clock: clk, maxKeys: 1 << 30}
|
|
lim := &limiter{opts: o}
|
|
|
|
// "idle" makes a single request, then time passes so it refills to full.
|
|
lim.allow("idle")
|
|
clk.Advance(10 * time.Second)
|
|
|
|
// "active" drains its whole burst at the (advanced) current time.
|
|
for i := 0; i < 6; i++ {
|
|
lim.allow("active")
|
|
}
|
|
|
|
lim.sweep()
|
|
|
|
if _, ok := lim.buckets.Load("idle"); ok {
|
|
t.Error("fully-refilled idle bucket was not evicted")
|
|
}
|
|
if _, ok := lim.buckets.Load("active"); !ok {
|
|
t.Error("active bucket with a partial limit was wrongly evicted")
|
|
}
|
|
}
|
|
|
|
// TestClientKeyTrustedProxy exercises the X-Forwarded-For walk used behind a
|
|
// trusted proxy, independent of the HTTP layer.
|
|
func TestClientKeyTrustedProxy(t *testing.T) {
|
|
o := &rateLimitOptions{}
|
|
WithTrustedProxies("192.168.0.0/16")(o)
|
|
|
|
r := newTestRequest("192.168.1.10:443", "203.0.113.7, 192.168.1.10")
|
|
if got := o.clientKey(r); got != "203.0.113.7" {
|
|
t.Fatalf("clientKey = %q, want real client 203.0.113.7", got)
|
|
}
|
|
|
|
// Untrusted peer: X-Forwarded-For must be ignored entirely.
|
|
r = newTestRequest("203.0.113.7:443", "10.0.0.1")
|
|
if got := o.clientKey(r); got != "203.0.113.7" {
|
|
t.Fatalf("clientKey = %q, want peer 203.0.113.7 (XFF ignored)", got)
|
|
}
|
|
}
|