| 1 | package calls |
| 2 | |
| 3 | import ( |
| 4 | "sync" |
| 5 | "time" |
| 6 | ) |
| 7 | |
| 8 | // maxSignalBuckets bounds the signalling limiter's map (review Important |
| 9 | // 1): a generous multiple of what any real deployment holds at once |
| 10 | // (sixteen participations to a room), so legitimate traffic never |
| 11 | // evicts anything, while an attacker opening session after session |
| 12 | // cannot grow the map without bound. |
| 13 | const maxSignalBuckets = 4096 |
| 14 | |
| 15 | // rateLimiter is a per-key token bucket (spec §4.4): rate tokens are |
| 16 | // added to a key's bucket per second, up to burst held, and one message |
| 17 | // spends one token. POST /m/{sqid}/signal keys it on session and room |
| 18 | // (session+"|"+sqid), not on anything the request body claims: the |
| 19 | // session half comes from the session cookie the request actually |
| 20 | // presented, so nobody can spend another session's budget or dodge |
| 21 | // their own by rotating a claimed peer id (review Important 1). |
| 22 | // |
| 23 | // Entries are process-local, like the hub itself, and bounded by |
| 24 | // maxKeys: once the map holds that many keys, the least recently |
| 25 | // touched bucket is evicted to make room, so an attacker opening |
| 26 | // unbounded sessions cannot grow the map without bound. A bucket is |
| 27 | // also dropped once idle for a minute; either way it is simply |
| 28 | // recreated at burst on its next visit, which is correct — a key that |
| 29 | // has been quiet owes nothing. The idle sweep runs inline, every |
| 30 | // thousandth call, rather than on a timer: nothing here needs a |
| 31 | // goroutine, and the call rate this guards is exactly what would drive |
| 32 | // the sweep anyway. |
| 33 | type rateLimiter struct { |
| 34 | rate float64 // tokens added per second |
| 35 | burst float64 // bucket capacity |
| 36 | maxKeys int // 0 means unbounded — used only by tests; named to |
| 37 | // avoid shadowing the builtin cap (review Task 5 item 3) |
| 38 | |
| 39 | mu sync.Mutex |
| 40 | buckets map[string]*tokenBucket |
| 41 | calls uint64 |
| 42 | } |
| 43 | |
| 44 | type tokenBucket struct { |
| 45 | tokens float64 |
| 46 | seenAt time.Time |
| 47 | } |
| 48 | |
| 49 | func newRateLimiter(rate, burst, maxKeys int) *rateLimiter { |
| 50 | return &rateLimiter{rate: float64(rate), burst: float64(burst), maxKeys: maxKeys, buckets: map[string]*tokenBucket{}} |
| 51 | } |
| 52 | |
| 53 | // newRateLimiterEvery is a limiter that adds one token per interval: for |
| 54 | // the report endpoints, whose rates are slower than one a second. |
| 55 | func newRateLimiterEvery(every time.Duration, burst, maxKeys int) *rateLimiter { |
| 56 | return &rateLimiter{rate: 1 / every.Seconds(), burst: float64(burst), maxKeys: maxKeys, buckets: map[string]*tokenBucket{}} |
| 57 | } |
| 58 | |
| 59 | // allow spends one token for key, refilling by the time elapsed since |
| 60 | // its last visit. ok is false when the bucket is empty, in which case |
| 61 | // retryAfter is how long until one token is available. |
| 62 | func (l *rateLimiter) allow(key string, now time.Time) (ok bool, retryAfter time.Duration) { |
| 63 | l.mu.Lock() |
| 64 | defer l.mu.Unlock() |
| 65 | |
| 66 | l.calls++ |
| 67 | if l.calls%1000 == 0 { |
| 68 | l.sweepLocked(now) |
| 69 | } |
| 70 | |
| 71 | b, found := l.buckets[key] |
| 72 | if !found { |
| 73 | if l.maxKeys > 0 && len(l.buckets) >= l.maxKeys { |
| 74 | l.evictOldestLocked() |
| 75 | } |
| 76 | b = &tokenBucket{tokens: l.burst, seenAt: now} |
| 77 | l.buckets[key] = b |
| 78 | } else if elapsed := now.Sub(b.seenAt).Seconds(); elapsed > 0 { |
| 79 | b.tokens = min(l.burst, b.tokens+elapsed*l.rate) |
| 80 | b.seenAt = now |
| 81 | } |
| 82 | |
| 83 | if b.tokens < 1 { |
| 84 | need := (1 - b.tokens) / l.rate |
| 85 | return false, time.Duration(need * float64(time.Second)) |
| 86 | } |
| 87 | b.tokens-- |
| 88 | return true, 0 |
| 89 | } |
| 90 | |
| 91 | // sweepLocked drops any bucket idle for a full minute. |
| 92 | func (l *rateLimiter) sweepLocked(now time.Time) { |
| 93 | for k, b := range l.buckets { |
| 94 | if now.Sub(b.seenAt) > time.Minute { |
| 95 | delete(l.buckets, k) |
| 96 | } |
| 97 | } |
| 98 | } |
| 99 | |
| 100 | // evictOldestLocked drops the bucket least recently touched, so a |
| 101 | // bounded limiter never grows past maxKeys: a burst of distinct keys |
| 102 | // displaces the quietest one rather than growing the map without bound |
| 103 | // (review Important 1). O(n) in the bucket count, which only matters |
| 104 | // right at the limit — everyday traffic never reaches it. |
| 105 | func (l *rateLimiter) evictOldestLocked() { |
| 106 | var oldestKey string |
| 107 | var oldestAt time.Time |
| 108 | first := true |
| 109 | for k, b := range l.buckets { |
| 110 | if first || b.seenAt.Before(oldestAt) { |
| 111 | oldestKey, oldestAt, first = k, b.seenAt, false |
| 112 | } |
| 113 | } |
| 114 | if !first { |
| 115 | delete(l.buckets, oldestKey) |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | // size is how many keys the limiter currently holds — a test-only |
| 120 | // window into whether a request minted a bucket at all (review Task 5 |
| 121 | // item 1's rotation test). |
| 122 | func (l *rateLimiter) size() int { |
| 123 | l.mu.Lock() |
| 124 | defer l.mu.Unlock() |
| 125 | return len(l.buckets) |
| 126 | } |
| 127 | |
amadan