sl_map_web/rate_limit.rs
1//! Token-bucket rate limiter for authenticated-user create endpoints,
2//! persisted in the `rate_buckets` table.
3//!
4//! Each `(category, user)` pair has its own bucket. A bucket has a
5//! per-category capacity and a per-category refill interval (one token
6//! per `seconds_per_token` seconds). On each acquisition the bucket is
7//! refilled based on elapsed wall-clock time (clamped to capacity) and
8//! one token is deducted; if fewer than one token is available the
9//! acquisition returns [`Error::TooManyRequests`] with a conservative
10//! `Retry-After` value equal to one refill interval.
11//!
12//! The refill and deduction happen in a single SQLite UPSERT so two
13//! concurrent acquisitions for the same key cannot both pass — SQLite
14//! serialises the write and the `ON CONFLICT DO UPDATE … WHERE`
15//! clause re-checks the post-refill token count on the second arrival.
16
17use chrono::Utc;
18use sqlx::SqlitePool;
19use uuid::Uuid;
20
21use crate::error::Error;
22
23/// Categories of rate-limited create endpoints. The string form maps
24/// directly to the `rate_buckets.category` column; the CHECK constraint
25/// in migration `0004_rate_buckets.sql` lists the same three values.
26#[derive(Debug, Clone, Copy)]
27pub enum RateCategory {
28 /// `POST /api/groups` — creating a new group.
29 GroupCreate,
30 /// `POST /api/notecards` — saving a notecard (personal or group).
31 NotecardCreate,
32 /// `POST /api/groups/{id}/invitations` — sending an invitation.
33 InvitationCreate,
34}
35
36impl RateCategory {
37 /// Persisted form used in the `category` column.
38 const fn as_str(self) -> &'static str {
39 match self {
40 Self::GroupCreate => "group_create",
41 Self::NotecardCreate => "notecard_create",
42 Self::InvitationCreate => "invitation_create",
43 }
44 }
45
46 /// `(capacity, seconds_per_token)` for this category. The refill
47 /// rate per hour is `3600 / seconds_per_token`.
48 ///
49 /// - `GroupCreate`: 12 burst, 1 token / 20 min (3 / hour)
50 /// - `NotecardCreate`: 60 burst, 1 token / 6 min (10 / hour)
51 /// - `InvitationCreate`: 250 burst, 1 token / 6 min (10 / hour)
52 const fn parameters(self) -> (u32, u32) {
53 match self {
54 Self::GroupCreate => (12, 1200),
55 Self::NotecardCreate => (60, 360),
56 Self::InvitationCreate => (250, 360),
57 }
58 }
59}
60
61/// Try to acquire one token from `category`'s bucket for `user_id`.
62/// Returns `Ok(())` on success and [`Error::TooManyRequests`] with a
63/// `retry_after_secs` equal to the per-category refill interval on
64/// failure.
65///
66/// # Errors
67///
68/// - [`Error::TooManyRequests`] when the bucket has fewer than 1 token
69/// after refill.
70/// - [`Error::Database`] if the UPSERT fails for any reason other than
71/// the conditional-update no-op (the lookup query is logged via
72/// `tracing::error!` before being collapsed).
73pub async fn try_acquire(
74 db: &SqlitePool,
75 category: RateCategory,
76 user_id: Uuid,
77) -> Result<(), Error> {
78 let (capacity, seconds_per_token) = category.parameters();
79 let now = Utc::now();
80 // A single atomic UPSERT. On first request for this (category,
81 // user) the INSERT path runs and the new row is created with
82 // `capacity - 1` tokens. On subsequent requests the ON CONFLICT
83 // path refills based on the elapsed time since `last_refill_at`,
84 // clamps at capacity, and deducts 1 — but only if the post-refill
85 // count is at least 1, otherwise the UPDATE is filtered out by the
86 // WHERE clause and `RETURNING` yields no row.
87 let row: Option<(f64,)> = sqlx::query_as(
88 "INSERT INTO rate_buckets (category, user_id, tokens, last_refill_at) \
89 VALUES (?1, ?2, ?3, ?4) \
90 ON CONFLICT(category, user_id) DO UPDATE SET \
91 tokens = MIN(?5, rate_buckets.tokens + \
92 (julianday(excluded.last_refill_at) \
93 - julianday(rate_buckets.last_refill_at)) * 86400.0 / ?6) - 1, \
94 last_refill_at = excluded.last_refill_at \
95 WHERE MIN(?5, rate_buckets.tokens + \
96 (julianday(excluded.last_refill_at) \
97 - julianday(rate_buckets.last_refill_at)) * 86400.0 / ?6) >= 1 \
98 RETURNING tokens",
99 )
100 .bind(category.as_str())
101 .bind(user_id.as_bytes().to_vec())
102 .bind(f64::from(capacity) - 1.0)
103 .bind(now)
104 .bind(f64::from(capacity))
105 .bind(f64::from(seconds_per_token))
106 .fetch_optional(db)
107 .await
108 .map_err(|err| {
109 tracing::error!("rate_limit acquire failed: {err}");
110 Error::Database
111 })?;
112 if row.is_some() {
113 Ok(())
114 } else {
115 Err(Error::TooManyRequests {
116 retry_after_secs: u64::from(seconds_per_token),
117 })
118 }
119}