// Package ratelimit is the gateway's in-memory anti-abuse limiter: a token // bucket per key (golang.org/x/time/rate). The connect edge keys the public // class per client IP, the authenticated class per user id, and a stricter // sub-limit guards the email-code path; the admin proxy keys per IP. Buckets are // swept lazily so an idle key does not leak memory. package ratelimit import ( "sync" "time" "golang.org/x/time/rate" ) // Policy is a token-bucket rate and burst. type Policy struct { Limit rate.Limit Burst int } // PerMinute builds a Policy allowing perMinute events per minute with the given // burst. func PerMinute(perMinute, burst int) Policy { return Policy{Limit: rate.Limit(float64(perMinute) / 60.0), Burst: burst} } // Per builds a Policy allowing events per window with the given burst. func Per(events int, window time.Duration, burst int) Policy { return Policy{Limit: rate.Limit(float64(events) / window.Seconds()), Burst: burst} } // staleAfter is how long an unused bucket is retained before the lazy sweep // discards it; sweepInterval bounds how often the sweep runs. const ( staleAfter = 10 * time.Minute sweepInterval = time.Minute ) // Limiter holds the per-key token buckets. type Limiter struct { now func() time.Time mu sync.Mutex buckets map[string]*bucket lastSweep time.Time } type bucket struct { lim *rate.Limiter seen time.Time } // New constructs an empty Limiter. func New() *Limiter { now := func() time.Time { return time.Now() } return &Limiter{now: now, buckets: make(map[string]*bucket), lastSweep: now()} } // Allow reports whether one event under key is permitted by policy, consuming a // token when it is. func (l *Limiter) Allow(key string, p Policy) bool { l.mu.Lock() defer l.mu.Unlock() now := l.now() l.sweepLocked(now) b, ok := l.buckets[key] if !ok { b = &bucket{lim: rate.NewLimiter(p.Limit, p.Burst)} l.buckets[key] = b } b.seen = now return b.lim.Allow() } // sweepLocked discards buckets unused for staleAfter, at most once per // sweepInterval. The caller holds l.mu. func (l *Limiter) sweepLocked(now time.Time) { if now.Sub(l.lastSweep) < sweepInterval { return } l.lastSweep = now for k, b := range l.buckets { if now.Sub(b.seen) > staleAfter { delete(l.buckets, k) } } }