Skip to main content

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}