diff --git a/per_source_ratelimit.go b/per_source_ratelimit.go index 6c41422..d9e51ea 100644 --- a/per_source_ratelimit.go +++ b/per_source_ratelimit.go @@ -62,10 +62,11 @@ func (l *perSourceAuthLimiter) allow(req *http.Request) bool { if now.Sub(l.lastSweep) >= time.Minute { l.sweep(now) } - e := l.entries[ip] + key := limiterKey(ip) + e := l.entries[key] if e == nil { e = &perSourceAuthEntry{limiter: rate.NewLimiter(l.limit, l.burst)} - l.entries[ip] = e + l.entries[key] = e } e.lastSeen = now return e.limiter.Allow() @@ -115,3 +116,19 @@ func isInternalSource(ip string) bool { return p.IsLoopback() || p.IsPrivate() || p.IsLinkLocalUnicast() || p.IsLinkLocalMulticast() || p.IsUnspecified() } + +// ipv6PrefixBits is the IPv6 prefix that shares one bucket: a host holds a +// whole /64, so per-address keys would give it 2^64 buckets. +const ipv6PrefixBits = 64 + +// limiterKey maps a source IP to its bucket key (IPv6 collapses to its /64). +func limiterKey(ip string) string { + p := net.ParseIP(ip) + if p == nil { + return ip + } + if p.To4() != nil { + return p.String() + } + return (&net.IPNet{IP: p.Mask(net.CIDRMask(ipv6PrefixBits, 128)), Mask: net.CIDRMask(ipv6PrefixBits, 128)}).String() +} diff --git a/per_source_ratelimit_ipv6_test.go b/per_source_ratelimit_ipv6_test.go new file mode 100644 index 0000000..61bd1a8 --- /dev/null +++ b/per_source_ratelimit_ipv6_test.go @@ -0,0 +1,49 @@ +package traefikoidc + +import ( + "fmt" + "net/http/httptest" + "testing" +) + +// An IPv6 /64 holder must share one bucket; per-address keys let one host +// mint 2^64 buckets and bypass the throttle. +func TestPerSourceLimiter_IPv6SamePrefixSharesBucket(t *testing.T) { + const burst = 3 + l := newPerSourceAuthLimiter(burst) + allowed := 0 + for i := 0; i < 50; i++ { + req := httptest.NewRequest("GET", "/", nil) + req.RemoteAddr = fmt.Sprintf("[2001:db8:1:2::%x]:1234", i+1) + if l.allow(req) { + allowed++ + } + } + if allowed > burst { + t.Fatalf("one /64 got %d allows, want <= %d", allowed, burst) + } +} + +func TestPerSourceLimiter_IPv6DifferentPrefixIndependent(t *testing.T) { + l := newPerSourceAuthLimiter(1) + a := httptest.NewRequest("GET", "/", nil) + a.RemoteAddr = "[2001:db8:1:2::1]:1" + b := httptest.NewRequest("GET", "/", nil) + b.RemoteAddr = "[2001:db8:1:3::1]:1" + if !l.allow(a) || !l.allow(b) { + t.Fatal("distinct /64 prefixes must have independent buckets") + } +} + +func TestLimiterKey(t *testing.T) { + cases := map[string]string{ + "192.0.2.1": "192.0.2.1", + "2001:db8:1:2::9": "2001:db8:1:2::/64", + "::ffff:192.0.2.1": "192.0.2.1", + } + for in, want := range cases { + if got := limiterKey(in); got != want { + t.Errorf("limiterKey(%q)=%q want %q", in, got, want) + } + } +}