alibi_sqlx/
rate_limit.rs

1//! Atomic rolling rate limits persisted independently of an application's auth schema.
2
3use crate::SqlxValue;
4use crate::pool::{Exec, SqlxPool};
5use crate::sql::Sql;
6use crate::store::migrator::{Migration, apply};
7use alibi_core::middleware::rate_limit::bucket::{self, Step};
8use alibi_core::store::SchemaMigrator;
9use alibi_core::{AuthResult, EndpointRateLimit, RateLimitDecision, RateLimitStorage};
10use async_trait::async_trait;
11
12const RATE_LIMIT_SCHEMA: Migration = Migration {
13    name: "m20261013_000003_typescript_rate_limit_schema",
14    sqlite: alibi_schema_registry::SQLITE_RATE_LIMIT_SCHEMA,
15    postgres: alibi_schema_registry::POSTGRES_RATE_LIMIT_SCHEMA,
16};
17
18/// One stored quota bucket.
19#[derive(Clone, Debug, PartialEq, sqlx::FromRow)]
20pub struct RateLimitRow {
21    pub key: String,
22    pub count: crate::Number,
23    #[sqlx(rename = "lastRequest")]
24    pub last_request: 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}
34
35impl SqlxRateLimitStorage {
36    #[must_use]
37    pub fn new(pool: impl Into<SqlxPool>) -> Self {
38        Self { pool: pool.into() }
39    }
40
41    const fn exec(&self) -> Exec<'_> {
42        Exec::pool(&self.pool)
43    }
44}
45
46/// Installs the opt-in rate-limit table using its own migration ledger.
47/// Migration errors are returned without serving requests on missing storage.
48#[async_trait]
49impl SchemaMigrator for SqlxRateLimitStorage {
50    async fn migrate(&self) -> AuthResult<()> {
51        apply(
52            &self.pool,
53            "better_auth_rate_limit_migrations",
54            &[RATE_LIMIT_SCHEMA],
55        )
56        .await
57    }
58}
59
60#[async_trait]
61impl RateLimitStorage for SqlxRateLimitStorage {
62    async fn consume(&self, key: &str, rule: &EndpointRateLimit) -> AuthResult<RateLimitDecision> {
63        let backend = self.exec().engine();
64        loop {
65            let mut select = Sql::with(
66                backend,
67                "SELECT \"rateLimit\".\"key\", \"rateLimit\".\"count\", \"rateLimit\".\"lastRequest\" FROM \"rateLimit\" WHERE ",
68            );
69            select.compare("rateLimit", "key", " = ", key);
70            select.push(" LIMIT 1");
71            let observed = self.exec().fetch_optional::<RateLimitRow>(select).await?;
72            let now = chrono::Utc::now().timestamp_millis();
73            let Some(observed) = observed else {
74                let mut insert = Sql::with(
75                    backend,
76                    "INSERT INTO \"rateLimit\" (\"id\", \"key\", \"count\", \"lastRequest\") VALUES (",
77                );
78                insert.bind(uuid::Uuid::new_v4().to_string());
79                insert.push(", ");
80                insert.bind(key);
81                insert.push(", ");
82                insert.bind(1_i32);
83                insert.push(", ");
84                insert.bind(now);
85                insert.push(") ON CONFLICT (\"key\") DO NOTHING");
86                if self.exec().execute(insert).await? == 1 {
87                    return Ok(RateLimitDecision::Allowed);
88                }
89                continue;
90            };
91            let count = match bucket::step(observed.count.0, observed.last_request, now, rule) {
92                Step::Blocked { retry_after } => {
93                    return Ok(RateLimitDecision::Blocked { retry_after });
94                }
95                Step::Allowed { count, .. } => count,
96            };
97            let mut update = Sql::with(backend, "UPDATE \"rateLimit\" SET \"count\" = ");
98            update.bind(crate::Number(count).into_sql_value());
99            update.push(", \"lastRequest\" = ");
100            update.bind(now);
101            update.push(" WHERE ");
102            update.compare("rateLimit", "key", " = ", key);
103            update.push(" AND ");
104            update.compare("rateLimit", "lastRequest", " = ", observed.last_request);
105            update.push(" AND ");
106            update.compare("rateLimit", "count", " = ", observed.count.into_sql_value());
107            if self.exec().execute(update).await? == 1 {
108                return Ok(RateLimitDecision::Allowed);
109            }
110        }
111    }
112}
alibi_sqlx/rate_limit.rs

alibi_sqlx/
rate_limit.rs
