83 lines
2.1 KiB
Go
83 lines
2.1 KiB
Go
|
package rfc2136
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Per-key token bucket. Hamilton M8: a compromised TSIG key — or a
|
||
|
|
// misconfigured client retrying forever — must not be able to drive
|
||
|
|
// unbounded UPDATE traffic. Each UPDATE costs disk IOPS, a git commit,
|
||
|
|
// and a slot in the SOA serial counter (9999/day per zone). 100
|
||
|
|
// UPDATEs/minute per key is well above any legitimate ACME workflow
|
||
|
|
// (a full renewal storm across our ~84 zones might emit ~200 UPDATEs
|
||
|
|
// total over several minutes); anything beyond is suspicious.
|
||
|
|
const (
|
||
|
|
defaultRateBurst = 100 // max tokens
|
||
|
|
defaultRatePeriod = time.Minute // refill window
|
||
|
|
)
|
||
|
|
|
||
|
|
// rateLimiter is a goroutine-safe per-key token bucket. The zero value
|
||
|
|
// is unusable; construct via newRateLimiter.
|
||
|
|
type rateLimiter struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
buckets map[string]*bucket
|
||
|
|
burst float64 // max tokens
|
||
|
|
period time.Duration // time to fully refill
|
||
|
|
}
|
||
|
|
|
||
|
|
type bucket struct {
|
||
|
|
tokens float64
|
||
|
|
lastRefill time.Time
|
||
|
|
}
|
||
|
|
|
||
|
|
func newRateLimiter(burst int, period time.Duration) *rateLimiter {
|
||
|
|
if burst <= 0 {
|
||
|
|
burst = defaultRateBurst
|
||
|
|
}
|
||
|
|
if period <= 0 {
|
||
|
|
period = defaultRatePeriod
|
||
|
|
}
|
||
|
|
return &rateLimiter{
|
||
|
|
buckets: make(map[string]*bucket),
|
||
|
|
burst: float64(burst),
|
||
|
|
period: period,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// allow attempts to take one token for `key`. Returns true if a token
|
||
|
|
// was available, false otherwise. New keys start full (burst tokens).
|
||
|
|
//
|
||
|
|
// Refill is continuous: tokens accumulate at burst/period per second.
|
||
|
|
// The bucket caps at burst tokens.
|
||
|
|
func (rl *rateLimiter) allow(key string, now time.Time) bool {
|
||
|
|
rl.mu.Lock()
|
||
|
|
defer rl.mu.Unlock()
|
||
|
|
|
||
|
|
b, ok := rl.buckets[key]
|
||
|
|
if !ok {
|
||
|
|
// First time we see this key — start the bucket full so
|
||
|
|
// legitimate clients don't see refill delays at boot.
|
||
|
|
rl.buckets[key] = &bucket{
|
||
|
|
tokens: rl.burst - 1,
|
||
|
|
lastRefill: now,
|
||
|
|
}
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
|
||
|
|
// Refill: tokens earned since last access.
|
||
|
|
elapsed := now.Sub(b.lastRefill).Seconds()
|
||
|
|
earned := elapsed * (rl.burst / rl.period.Seconds())
|
||
|
|
b.tokens += earned
|
||
|
|
if b.tokens > rl.burst {
|
||
|
|
b.tokens = rl.burst
|
||
|
|
}
|
||
|
|
b.lastRefill = now
|
||
|
|
|
||
|
|
if b.tokens >= 1.0 {
|
||
|
|
b.tokens -= 1.0
|
||
|
|
return true
|
||
|
|
}
|
||
|
|
return false
|
||
|
|
}
|