82 lines
2.0 KiB
Go
82 lines
2.0 KiB
Go
package server
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// RatePolicy is a fixed-window policy. Buckets expire naturally and cleanup
|
|
// prevents unauthenticated traffic from consuming unbounded memory.
|
|
type RatePolicy struct {
|
|
Limit int
|
|
Window time.Duration
|
|
}
|
|
|
|
type rateBucket struct {
|
|
Count int
|
|
WindowStart time.Time
|
|
LastSeen time.Time
|
|
}
|
|
|
|
type rateLimiter struct {
|
|
mu sync.Mutex
|
|
buckets map[string]rateBucket
|
|
now func() time.Time
|
|
}
|
|
|
|
const maxRateLimitBuckets = 10000
|
|
|
|
func newRateLimiter(now func() time.Time) *rateLimiter {
|
|
if now == nil {
|
|
now = time.Now
|
|
}
|
|
return &rateLimiter{buckets: make(map[string]rateBucket), now: now}
|
|
}
|
|
|
|
// Allow consumes one attempt and reports the remaining wait when limited.
|
|
func (l *rateLimiter) Allow(key string, policy RatePolicy) (bool, time.Duration) {
|
|
if policy.Limit <= 0 || policy.Window <= 0 {
|
|
return true, 0
|
|
}
|
|
now := l.now().UTC()
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
bucket := l.buckets[key]
|
|
if bucket.WindowStart.IsZero() && len(l.buckets) >= maxRateLimitBuckets {
|
|
// An attacker can vary IP/account keys faster than age cleanup runs.
|
|
// Evict the least recently seen bucket to keep memory strictly bounded.
|
|
var oldestKey string
|
|
var oldest time.Time
|
|
for candidate, existing := range l.buckets {
|
|
if oldestKey == "" || existing.LastSeen.Before(oldest) {
|
|
oldestKey, oldest = candidate, existing.LastSeen
|
|
}
|
|
}
|
|
if oldestKey != "" {
|
|
delete(l.buckets, oldestKey)
|
|
}
|
|
}
|
|
if bucket.WindowStart.IsZero() || !now.Before(bucket.WindowStart.Add(policy.Window)) {
|
|
bucket = rateBucket{WindowStart: now}
|
|
}
|
|
bucket.LastSeen = now
|
|
if bucket.Count >= policy.Limit {
|
|
l.buckets[key] = bucket
|
|
return false, bucket.WindowStart.Add(policy.Window).Sub(now)
|
|
}
|
|
bucket.Count++
|
|
l.buckets[key] = bucket
|
|
return true, 0
|
|
}
|
|
|
|
func (l *rateLimiter) Cleanup(maxAge time.Duration) {
|
|
now := l.now().UTC()
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
for key, bucket := range l.buckets {
|
|
if now.Sub(bucket.LastSeen) > maxAge {
|
|
delete(l.buckets, key)
|
|
}
|
|
}
|
|
}
|