Skip to main content

max / makenotwork

Add livechat rate limiter Token bucket per (user, room). Chat needs its own budget: Multithreaded caps posting at 15/60s, a sensible forum rate and an unusable chat one, and the per-IP governor layers key on a connection that chat holds open across every message. A bucket rather than a fixed window because chat traffic is bursty by nature. Keyed by the pair so being talkative in one room does not silence someone in another, and a room's budget cannot be spent from outside it. Eviction is a memory control, not tidiness: one bucket per pair is unbounded state keyed by something an attacker influences, under a hard 512M cap. recommended_idle_timeout returns the shortest timeout that cannot turn eviction into a bypass, since a fully refilled bucket answers identically to a fresh one, and that equivalence is tested. `now` is injected rather than read internally so the behavior is testable without sleeping, and saturates backwards so a stale instant cannot panic. Design: wiki livechat-design Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Author: Max Johnson <me@maxj.phd> · 2026-07-19 00:31 UTC
Commit: abfb09482b8c663f1414e11ce82c3bfb4b38df4c
Parent: 7fc6d99
2 files changed, +332 insertions, -0 deletions
M shared/livechat/src/lib.rs +2
@@ -42,6 +42,7 @@ mod event;
42 42 mod hub;
43 43 mod ids;
44 44 mod message;
45 + mod rate_limit;
45 46 mod retention;
46 47 mod room;
47 48 #[cfg(feature = "axum")]
@@ -54,6 +55,7 @@ pub use event::ChatEvent;
54 55 pub use hub::{Hub, HubLimits, Subscription};
55 56 pub use ids::{MessageId, Nonce, RoomId, UserId};
56 57 pub use message::{MAX_MESSAGE_LEN, Message, MessageBody};
58 + pub use rate_limit::{RateDecision, RateLimiter, RateLimits};
57 59 pub use retention::{MAX_AGE_CEILING, MAX_MESSAGES_CEILING, Retention};
58 60 pub use room::{Room, RoomState};
59 61 #[cfg(feature = "axum")]
A shared/livechat/src/rate_limit.rs +330
@@ -0,0 +1,330 @@
1 + //! Per-(user, room) send budget.
2 + //!
3 + //! Chat needs its own limiter rather than borrowing either host's. Multithreaded
4 + //! caps posting at 15 per 60 seconds, which is a sensible forum-post rate and an
5 + //! unusable chat rate. The per-IP `tower_governor` layers do not fit either: they
6 + //! key on the connection, and chat's connection is long-lived and shared by every
7 + //! message on it.
8 + //!
9 + //! A token bucket rather than a fixed window, because the shape of chat traffic is
10 + //! bursty by nature. Someone answering a question types four short lines in five
11 + //! seconds and then goes quiet for a minute; a fixed window either refuses the
12 + //! fourth line or has to be set so loose it stops being a limit.
13 + //!
14 + //! # Eviction is not optional
15 + //!
16 + //! One bucket per (user, room) pair, held in a map, is unbounded state keyed by
17 + //! something an attacker influences. Under a hard 512M cgroup cap that is a memory
18 + //! exhaustion vector, not untidiness, so [`RateLimiter::evict_idle`] exists and
19 + //! the host is expected to call it on a timer. Multithreaded's per-IP governor
20 + //! already sweeps its own buckets every 300 seconds; this wants the same
21 + //! treatment.
22 +
23 + use std::time::{Duration, Instant};
24 +
25 + use dashmap::DashMap;
26 +
27 + use crate::ids::{RoomId, UserId};
28 +
29 + /// Send budget.
30 + #[derive(Debug, Clone, Copy, PartialEq, Eq)]
31 + pub struct RateLimits {
32 + /// Messages allowed back to back from a standing start.
33 + pub burst: u32,
34 + /// Long-run rate, in messages per minute.
35 + pub sustain_per_min: u32,
36 + }
37 +
38 + impl Default for RateLimits {
39 + fn default() -> Self {
40 + // Ten in a row, then one every two seconds. Fast enough that a person
41 + // talking normally never meets it, slow enough that a script is throttled
42 + // to something a room can absorb.
43 + Self {
44 + burst: 10,
45 + sustain_per_min: 30,
46 + }
47 + }
48 + }
49 +
50 + impl RateLimits {
51 + fn refill_per_sec(&self) -> f64 {
52 + f64::from(self.sustain_per_min) / 60.0
53 + }
54 + }
55 +
56 + #[derive(Debug, Clone, Copy)]
57 + struct Bucket {
58 + tokens: f64,
59 + last: Instant,
60 + }
61 +
62 + /// Outcome of a send check.
63 + #[derive(Debug, Clone, Copy, PartialEq, Eq)]
64 + pub enum RateDecision {
65 + Allow,
66 + /// Refused. `retry_after` is when a token next becomes available, so the
67 + /// caller can tell the client rather than leaving it to guess.
68 + Deny { retry_after: Duration },
69 + }
70 +
71 + impl RateDecision {
72 + pub fn is_allowed(&self) -> bool {
73 + matches!(self, Self::Allow)
74 + }
75 + }
76 +
77 + /// Token buckets keyed by (user, room).
78 + ///
79 + /// Keyed by the pair rather than by user alone so that being talkative in one
80 + /// room does not silence someone in another, and so a room's budget cannot be
81 + /// consumed from outside it.
82 + pub struct RateLimiter {
83 + buckets: DashMap<(UserId, RoomId), Bucket>,
84 + limits: RateLimits,
85 + }
86 +
87 + impl RateLimiter {
88 + pub fn new(limits: RateLimits) -> Self {
89 + Self {
90 + buckets: DashMap::new(),
91 + limits,
92 + }
93 + }
94 +
95 + /// Take a token for one message.
96 + ///
97 + /// `now` is passed in rather than read here so the behavior is testable
98 + /// without sleeping, and so a caller batching work can use one consistent
99 + /// instant.
100 + pub fn check(&self, user: UserId, room: RoomId, now: Instant) -> RateDecision {
101 + let capacity = f64::from(self.limits.burst);
102 + let refill = self.limits.refill_per_sec();
103 +
104 + let mut bucket = self.buckets.entry((user, room)).or_insert(Bucket {
105 + tokens: capacity,
106 + last: now,
107 + });
108 +
109 + // `saturating_duration_since` rather than subtraction: a caller passing a
110 + // `now` older than the last one (batching, or a clock read from a
111 + // different task) would otherwise panic.
112 + let elapsed = now.saturating_duration_since(bucket.last).as_secs_f64();
113 + bucket.tokens = (bucket.tokens + elapsed * refill).min(capacity);
114 + bucket.last = now;
115 +
116 + if bucket.tokens >= 1.0 {
117 + bucket.tokens -= 1.0;
118 + RateDecision::Allow
119 + } else {
120 + let deficit = 1.0 - bucket.tokens;
121 + RateDecision::Deny {
122 + retry_after: Duration::from_secs_f64(deficit / refill),
123 + }
124 + }
125 + }
126 +
127 + /// Drop buckets untouched for longer than `idle_for`.
128 + ///
129 + /// A bucket at full tokens carries no information: recreating it produces the
130 + /// same answer as keeping it. So eviction cannot make the limiter more
131 + /// permissive as long as `idle_for` is at least the time it takes a bucket to
132 + /// refill, which is what [`RateLimiter::recommended_idle_timeout`] returns.
133 + ///
134 + /// Returns how many were dropped.
135 + pub fn evict_idle(&self, now: Instant, idle_for: Duration) -> usize {
136 + let before = self.buckets.len();
137 + self.buckets
138 + .retain(|_, bucket| now.saturating_duration_since(bucket.last) < idle_for);
139 + before - self.buckets.len()
140 + }
141 +
142 + /// The shortest idle timeout that cannot make eviction a bypass: long enough
143 + /// for an empty bucket to have refilled completely.
144 + pub fn recommended_idle_timeout(&self) -> Duration {
145 + Duration::from_secs_f64(f64::from(self.limits.burst) / self.limits.refill_per_sec())
146 + }
147 +
148 + /// Live bucket count. Diagnostics, and the number to watch if eviction is
149 + /// ever suspected of not running.
150 + pub fn bucket_count(&self) -> usize {
151 + self.buckets.len()
152 + }
153 + }
154 +
155 + #[cfg(test)]
156 + mod tests {
157 + use super::*;
158 + use uuid::Uuid;
159 +
160 + fn user() -> UserId {
161 + UserId(Uuid::new_v4())
162 + }
163 +
164 + fn room() -> RoomId {
165 + RoomId(Uuid::new_v4())
166 + }
167 +
168 + fn limiter(burst: u32, sustain_per_min: u32) -> RateLimiter {
169 + RateLimiter::new(RateLimits {
170 + burst,
171 + sustain_per_min,
172 + })
173 + }
174 +
175 + #[test]
176 + fn burst_is_spendable_immediately_then_refused() {
177 + let rl = limiter(3, 60);
178 + let (u, r, t0) = (user(), room(), Instant::now());
179 +
180 + for i in 0..3 {
181 + assert!(rl.check(u, r, t0).is_allowed(), "burst message {i}");
182 + }
183 + assert!(!rl.check(u, r, t0).is_allowed(), "burst exhausted");
184 + }
185 +
186 + #[test]
187 + fn refuses_with_a_usable_retry_after() {
188 + // 60/min is one per second, so an empty bucket is one second from a token.
189 + let rl = limiter(1, 60);
190 + let (u, r, t0) = (user(), room(), Instant::now());
191 +
192 + assert!(rl.check(u, r, t0).is_allowed());
193 + match rl.check(u, r, t0) {
194 + RateDecision::Deny { retry_after } => {
195 + let ms = retry_after.as_millis();
196 + assert!((900..=1100).contains(&ms), "retry_after was {ms}ms");
197 + }
198 + RateDecision::Allow => panic!("should have been refused"),
199 + }
200 + }
201 +
202 + #[test]
203 + fn tokens_come_back_over_time() {
204 + let rl = limiter(2, 60);
205 + let (u, r, t0) = (user(), room(), Instant::now());
206 +
207 + assert!(rl.check(u, r, t0).is_allowed());
208 + assert!(rl.check(u, r, t0).is_allowed());
209 + assert!(!rl.check(u, r, t0).is_allowed());
210 +
211 + assert!(rl.check(u, r, t0 + Duration::from_secs(1)).is_allowed());
212 + }
213 +
214 + #[test]
215 + fn refill_does_not_exceed_the_burst_ceiling() {
216 + // Idling for an hour must not bank an hour of messages.
217 + let rl = limiter(3, 60);
218 + let (u, r, t0) = (user(), room(), Instant::now());
219 +
220 + assert!(rl.check(u, r, t0).is_allowed());
221 + let long_later = t0 + Duration::from_secs(3600);
222 +
223 + for i in 0..3 {
224 + assert!(rl.check(u, r, long_later).is_allowed(), "message {i}");
225 + }
226 + assert!(
227 + !rl.check(u, r, long_later).is_allowed(),
228 + "banked more than the burst"
229 + );
230 + }
231 +
232 + #[test]
233 + fn the_sustained_rate_is_actually_the_sustained_rate() {
234 + // Spend the burst, then send at exactly the refill interval forever.
235 + let rl = limiter(5, 60);
236 + let (u, r) = (user(), room());
237 + let mut t = Instant::now();
238 +
239 + for _ in 0..5 {
240 + assert!(rl.check(u, r, t).is_allowed());
241 + }
242 + for i in 0..100 {
243 + t += Duration::from_secs(1);
244 + assert!(rl.check(u, r, t).is_allowed(), "sustained message {i}");
245 + }
246 + // And faster than that is still refused.
247 + assert!(!rl.check(u, r, t).is_allowed());
248 + }
249 +
250 + #[test]
251 + fn budgets_are_per_user_and_per_room() {
252 + let rl = limiter(1, 60);
253 + let (u1, u2) = (user(), user());
254 + let (r1, r2) = (room(), room());
255 + let t0 = Instant::now();
256 +
257 + assert!(rl.check(u1, r1, t0).is_allowed());
258 + assert!(!rl.check(u1, r1, t0).is_allowed(), "same pair is spent");
259 + assert!(rl.check(u2, r1, t0).is_allowed(), "other user unaffected");
260 + assert!(rl.check(u1, r2, t0).is_allowed(), "other room unaffected");
261 + }
262 +
263 + #[test]
264 + fn a_backwards_now_does_not_panic() {
265 + let rl = limiter(2, 60);
266 + let (u, r) = (user(), room());
267 + let t = Instant::now() + Duration::from_secs(10);
268 +
269 + assert!(rl.check(u, r, t).is_allowed());
270 + // Earlier than the bucket's last touch: elapsed saturates to zero rather
271 + // than underflowing.
272 + assert!(rl.check(u, r, t - Duration::from_secs(5)).is_allowed());
273 + }
274 +
275 + #[test]
276 + fn eviction_drops_only_idle_buckets() {
277 + let rl = limiter(5, 60);
278 + let (busy, idle) = (user(), user());
279 + let r = room();
280 + let t0 = Instant::now();
281 +
282 + rl.check(busy, r, t0);
283 + rl.check(idle, r, t0);
284 + assert_eq!(rl.bucket_count(), 2);
285 +
286 + let later = t0 + Duration::from_secs(60);
287 + rl.check(busy, r, later);
288 +
289 + assert_eq!(rl.evict_idle(later, Duration::from_secs(30)), 1);
290 + assert_eq!(rl.bucket_count(), 1);
291 + }
292 +
293 + #[test]
294 + fn eviction_at_the_recommended_timeout_is_not_a_bypass() {
295 + // The concern: spend the burst, wait out eviction, and get a fresh bucket
296 + // sooner than refilling would have allowed. At the recommended timeout the
297 + // bucket has fully refilled anyway, so the two paths agree.
298 + let rl = limiter(5, 60);
299 + let (u, r, t0) = (user(), room(), Instant::now());
300 +
301 + for _ in 0..5 {
302 + assert!(rl.check(u, r, t0).is_allowed());
303 + }
304 + assert!(!rl.check(u, r, t0).is_allowed());
305 +
306 + let timeout = rl.recommended_idle_timeout();
307 + assert_eq!(timeout, Duration::from_secs(5), "5 tokens at 1/sec");
308 +
309 + let after = t0 + timeout;
310 + assert_eq!(rl.evict_idle(after, timeout), 1);
311 +
312 + // A fresh bucket gives exactly the burst, which is what waiting the same
313 + // period without eviction would also have given.
314 + for i in 0..5 {
315 + assert!(rl.check(u, r, after).is_allowed(), "post-eviction {i}");
316 + }
317 + assert!(!rl.check(u, r, after).is_allowed(), "no extra budget");
318 + }
319 +
320 + #[test]
321 + fn default_limits_are_not_the_forum_post_limit() {
322 + // The whole reason this module exists: 15/min would make chat unusable.
323 + let d = RateLimits::default();
324 + assert!(
325 + d.sustain_per_min > 15,
326 + "chat sustain must exceed the forum post rate"
327 + );
328 + assert!(d.burst >= 5, "chat must absorb a short burst of typing");
329 + }
330 + }