alibi_sqlx/
rate_limit.rs

1//! Atomic rolling rate limits persisted independently of an application's auth schema.
2
3use crate::pool::{Exec, SqlxPool};
4use crate::sql::Sql;
5use crate::store::migrator::{Migration, apply};
6use alibi_core::middleware::rate_limit::bucket::{self, LongestWindow, Step};
7use alibi_core::store::SchemaMigrator;
8use alibi_core::{AuthResult, EndpointRateLimit, RateLimitDecision, RateLimitStorage};
9use async_trait::async_trait;
10use std::sync::Arc;
11
12const RATE_LIMIT_SCHEMA: Migration = Migration {
13    name: "m20261002_000001_rate_limit",
14    sqlite: r#"CREATE TABLE IF NOT EXISTS "rate_limit" ( "key" varchar NOT NULL PRIMARY KEY, "count" double NOT NULL, "expires_at" integer, "last_request" integer NOT NULL );"#,
15    postgres: r#"CREATE TABLE IF NOT EXISTS "rate_limit" ( "key" varchar NOT NULL PRIMARY KEY, "count" double precision NOT NULL, "expires_at" bigint, "last_request" bigint NOT NULL );"#,
16};
17
18/// One stored quota bucket.
19#[derive(Clone, Debug, PartialEq, sqlx::FromRow)]
20pub struct RateLimitRow {
21    pub key: String,
22    pub count: f64,
23    pub last_request: i64,
24    pub expires_at: Option<i64>,
25}
26
27/// Shared database storage. Run [`SchemaMigrator::migrate`] before serving
28/// requests; ordinary auth migrations do not create or claim its table.
29/// Compare-and-swap predicates preserve quotas across independent application instances.
30#[derive(Clone, Debug)]
31pub struct SqlxRateLimitStorage {
32    pool: SqlxPool,
33    longest_window: Arc<LongestWindow>,
34}
35
36impl SqlxRateLimitStorage {
37    #[must_use]
38    pub fn new(pool: impl Into<SqlxPool>) -> Self {
39        Self {
40            pool: pool.into(),
41            longest_window: Arc::new(LongestWindow::new()),
42        }
43    }
44
45    const fn exec(&self) -> Exec<'_> {
46        Exec::pool(&self.pool)
47    }
48
49    async fn prune(&self, now: i64) {
50        let Some(cutoff) = self.longest_window.prune_cutoff(now) else {
51            return;
52        };
53        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM \"rate_limit\" WHERE ");
54        sql.compare("rate_limit", "last_request", " < ", cutoff);
55        sql.push(" AND ");
56        sql.compare("rate_limit", "expires_at", " <= ", now);
57        // Pruning failures do not undo a successful quota consumption. The
58        // next admitted new/reset bucket retries cleanup without logging keys.
59        if self.exec().execute(sql).await.is_err() {
60            tracing::warn!("Rate-limit cleanup failed");
61        }
62    }
63}
64
65/// Installs the opt-in rate-limit table using its own migration ledger.
66/// Migration errors are returned without serving requests on missing storage.
67#[async_trait]
68impl SchemaMigrator for SqlxRateLimitStorage {
69    async fn migrate(&self) -> AuthResult<()> {
70        apply(
71            &self.pool,
72            "better_auth_rate_limit_migrations",
73            &[RATE_LIMIT_SCHEMA],
74        )
75        .await
76    }
77}
78
79#[async_trait]
80impl RateLimitStorage for SqlxRateLimitStorage {
81    fn observe_window(&self, window: f64) {
82        self.longest_window.observe(window);
83    }
84
85    async fn consume(&self, key: &str, rule: &EndpointRateLimit) -> AuthResult<RateLimitDecision> {
86        self.observe_window(rule.window_seconds);
87        let backend = self.exec().engine();
88        loop {
89            let mut select = Sql::with(
90                backend,
91                "SELECT \"rate_limit\".\"key\", \"rate_limit\".\"count\", \"rate_limit\".\"last_request\", \"rate_limit\".\"expires_at\" FROM \"rate_limit\" WHERE ",
92            );
93            select.compare("rate_limit", "key", " = ", key);
94            select.push(" LIMIT 1");
95            let observed = self.exec().fetch_optional::<RateLimitRow>(select).await?;
96            let now = chrono::Utc::now().timestamp_millis();
97            let Some(observed) = observed else {
98                let mut insert = Sql::with(
99                    backend,
100                    "INSERT INTO \"rate_limit\" (\"key\", \"count\", \"last_request\", \"expires_at\") VALUES (",
101                );
102                insert.bind(key);
103                insert.push(", ");
104                insert.bind(1.0_f64);
105                insert.push(", ");
106                insert.bind(now);
107                insert.push(", ");
108                insert.bind(bucket::expires_at(now, rule.window_seconds));
109                insert.push(") ON CONFLICT (\"key\") DO NOTHING");
110                if self.exec().execute(insert).await? == 1 {
111                    self.prune(now).await;
112                    return Ok(RateLimitDecision::Allowed);
113                }
114                continue;
115            };
116            let (count, reset) =
117                match bucket::step(observed.count, observed.last_request, now, rule) {
118                    Step::Blocked { retry_after } => {
119                        return Ok(RateLimitDecision::Blocked { retry_after });
120                    }
121                    Step::Allowed { count, reset } => (count, reset),
122                };
123            let mut update = Sql::with(backend, "UPDATE \"rate_limit\" SET \"count\" = ");
124            update.bind(count);
125            update.push(", \"last_request\" = ");
126            update.bind(now);
127            update.push(", \"expires_at\" = ");
128            update.bind(bucket::expires_at(now, rule.window_seconds));
129            update.push(" WHERE ");
130            update.compare("rate_limit", "key", " = ", key);
131            update.push(" AND ");
132            update.compare("rate_limit", "last_request", " = ", observed.last_request);
133            update.push(" AND ");
134            update.compare("rate_limit", "count", " = ", observed.count);
135            if self.exec().execute(update).await? == 1 {
136                if reset {
137                    self.prune(now).await;
138                }
139                return Ok(RateLimitDecision::Allowed);
140            }
141        }
142    }
143}
alibi_sqlx/rate_limit.rs

alibi_sqlx/
rate_limit.rs
