Download this file
| 1 | package gate |
| 2 | |
| 3 | import ( |
| 4 | "sync" |
| 5 | "time" |
| 6 | ) |
| 7 | |
| 8 | const ( |
| 9 | limitBurst = 10 |
| 10 | limitPeriod = time.Minute |
| 11 | limitPrune = 10_000 |
| 12 | ) |
| 13 | |
| 14 | // Limiter is a per-key token bucket: limitBurst attempts, refilling at |
| 15 | // limitBurst per limitPeriod. bcrypt at cost 10 already costs tens of |
| 16 | // milliseconds a comparison; this exists so the CPU is spent on people and |
| 17 | // not on a script. |
| 18 | type Limiter struct { |
| 19 | mu sync.Mutex |
| 20 | now func() time.Time |
| 21 | buckets map[string]*bucket |
| 22 | } |
| 23 | |
| 24 | type bucket struct { |
| 25 | tokens float64 |
| 26 | last time.Time |
| 27 | } |
| 28 | |
| 29 | func NewLimiter(now func() time.Time) *Limiter { |
| 30 | return &Limiter{now: now, buckets: map[string]*bucket{}} |
| 31 | } |
| 32 | |
| 33 | // Allow spends one token for key, reporting whether one was available. |
| 34 | func (l *Limiter) Allow(key string) bool { |
| 35 | l.mu.Lock() |
| 36 | defer l.mu.Unlock() |
| 37 | now := l.now() |
| 38 | b, ok := l.buckets[key] |
| 39 | if !ok { |
| 40 | if len(l.buckets) >= limitPrune { |
| 41 | l.prune(now) |
| 42 | } |
| 43 | b = &bucket{tokens: limitBurst, last: now} |
| 44 | l.buckets[key] = b |
| 45 | } |
| 46 | refill := now.Sub(b.last).Seconds() * (limitBurst / limitPeriod.Seconds()) |
| 47 | b.tokens = min(limitBurst, b.tokens+refill) |
| 48 | b.last = now |
| 49 | if b.tokens < 1 { |
| 50 | return false |
| 51 | } |
| 52 | b.tokens-- |
| 53 | return true |
| 54 | } |
| 55 | |
| 56 | // prune drops every bucket that has been idle long enough to be full again; |
| 57 | // keeping it would change no answer. |
| 58 | func (l *Limiter) prune(now time.Time) { |
| 59 | for k, b := range l.buckets { |
| 60 | if now.Sub(b.last) >= limitPeriod { |
| 61 | delete(l.buckets, k) |
| 62 | } |
| 63 | } |
| 64 | } |
| 65 | |
amadan