alibi_sqlx/store/
api_keys.rs

1use super::entities::api_key::Model;
2use super::entities::api_key_start::ApiKeyStart;
3use super::{SqlxStore, lock_exclusive};
4use crate::error::{record_not_inserted, record_not_updated};
5use crate::model::{self, ActiveRow, SqlxModel};
6use crate::pool::Exec;
7use crate::schema::AuthSchema;
8use crate::sql::Sql;
9use crate::value::SqlxValue;
10use alibi_core::error::{AuthError, AuthResult};
11use alibi_core::store::adapter::parse_optional_rfc3339;
12use alibi_core::store::{ApiKeyStore, ConsumeApiKeyResult};
13use alibi_core::types::{ApiKey, CreateApiKey, UpdateApiKey};
14use async_trait::async_trait;
15use chrono::Utc;
16use uuid::Uuid;
17
18#[async_trait]
19impl<S> ApiKeyStore for SqlxStore<S>
20where
21    S: AuthSchema + Send + Sync,
22{
23    async fn consume_api_key_usage_from_snapshot(
24        &self,
25        observed: &ApiKey,
26        global_rate_limit_enabled: bool,
27    ) -> AuthResult<ConsumeApiKeyResult> {
28        self.consume_usage_phases(observed, global_rate_limit_enabled)
29            .await
30    }
31
32    async fn create_api_key(&self, input: CreateApiKey) -> AuthResult<ApiKey> {
33        let now = Utc::now();
34        let start = input
35            .start
36            .map(|start| ApiKeyStart::prepare(&start, self.exec().engine()))
37            .transpose()?;
38        let sqlite_cast = start
39            .as_ref()
40            .is_some_and(ApiKeyStart::requires_sqlite_cast);
41        let mut active = ActiveRow::new();
42        active.set(
43            "id",
44            self.generated_id(self.exec(), "apikey", "apikey", "id")
45                .await?
46                .unwrap_or_else(|| Uuid::new_v4().to_string()),
47        );
48        active.set("name", input.name);
49        active.set("start", start.into_sql_value());
50        active.set("prefix", input.prefix);
51        active.set("key", input.key_hash);
52        active.set("referenceId", input.reference_id);
53        active.set("configId", input.config_id);
54        active.set("refillInterval", input.refill_interval);
55        active.set("refillAmount", input.refill_amount);
56        active.set("lastRefillAt", None::<chrono::DateTime<Utc>>);
57        active.set("enabled", input.enabled);
58        active.set("rateLimitEnabled", input.rate_limit_enabled);
59        active.set("rateLimitTimeWindow", input.rate_limit_time_window);
60        active.set("rateLimitMax", input.rate_limit_max);
61        active.set("requestCount", Some(0.0_f64));
62        active.set("remaining", input.remaining);
63        active.set("lastRequest", None::<chrono::DateTime<Utc>>);
64        active.set(
65            "expiresAt",
66            parse_optional_rfc3339(input.expires_at.as_deref(), "expiresAt")?,
67        );
68        active.set("createdAt", now);
69        active.set("updatedAt", now);
70        active.set("permissions", input.permissions);
71        active.set("metadata", input.metadata);
72        let inserted = if sqlite_cast {
73            // Only an invalid-surrogate SQLite start needs a bytes-to-TEXT expression.
74            let mut sql = Sql::with(self.exec().engine(), "INSERT INTO ");
75            sql.ident(Model::TABLE);
76            sql.push(" (");
77            let present = active.present().collect::<Vec<_>>();
78            sql.column_list(&present.iter().map(|(name, _)| *name).collect::<Vec<_>>());
79            sql.push(") SELECT ");
80            for (index, (name, value)) in present.into_iter().enumerate() {
81                if index > 0 {
82                    sql.push(", ");
83                }
84                if name == "start" {
85                    sql.push("CAST(");
86                    sql.bind(value.clone());
87                    sql.push(" AS text)");
88                } else {
89                    sql.bind(value.clone());
90                }
91            }
92            model::returning::<Model>(&mut sql);
93            self.exec()
94                .fetch_optional::<Model>(sql)
95                .await?
96                .ok_or_else(record_not_inserted)?
97        } else {
98            model::insert::<Model>(self.exec(), &active).await?
99        };
100        Ok(ApiKey::from(&inserted))
101    }
102
103    async fn get_api_key_by_id(&self, id: &str) -> AuthResult<Option<ApiKey>> {
104        let mut sql = model::by_id::<Model>(self.exec(), id);
105        model::limit_one(&mut sql);
106        Ok(self
107            .exec()
108            .fetch_optional::<Model>(sql)
109            .await?
110            .map(|model| ApiKey::from(&model)))
111    }
112
113    async fn get_api_key_by_hash(&self, hash: &str) -> AuthResult<Option<ApiKey>> {
114        let mut sql = model::select_model::<Model>(self.exec());
115        sql.push(" WHERE ");
116        sql.compare(Model::TABLE, "key", " = ", hash);
117        model::limit_one(&mut sql);
118        Ok(self
119            .exec()
120            .fetch_optional::<Model>(sql)
121            .await?
122            .map(|model| ApiKey::from(&model)))
123    }
124
125    async fn list_api_keys_by_reference(&self, reference_id: &str) -> AuthResult<Vec<ApiKey>> {
126        // Explicit ASC order matches TS insertion order on every backend.
127        let mut sql = model::select_model::<Model>(self.exec());
128        sql.push(" WHERE ");
129        sql.compare(Model::TABLE, "referenceId", " = ", reference_id);
130        sql.push(" ORDER BY ");
131        sql.column(Model::TABLE, "createdAt");
132        sql.push(" ASC");
133        Ok(self
134            .exec()
135            .fetch_all::<Model>(sql)
136            .await?
137            .iter()
138            .map(ApiKey::from)
139            .collect())
140    }
141
142    async fn update_api_key(&self, id: &str, update: UpdateApiKey) -> AuthResult<ApiKey> {
143        let mut sql = model::by_id::<Model>(self.exec(), id);
144        model::limit_one(&mut sql);
145        let Some(model) = self.exec().fetch_optional::<Model>(sql).await? else {
146            return Err(AuthError::not_found("API Key not found"));
147        };
148
149        let active = apply_update_fields(model.into_active(), update)?;
150        model::update::<Model>(self.exec(), &active)
151            .await?
152            .map(|model| ApiKey::from(&model))
153            .ok_or_else(record_not_updated)
154    }
155
156    #[expect(
157        clippy::as_conversions,
158        clippy::cast_precision_loss,
159        reason = "JavaScript-compatible numbers deliberately retain IEEE754 rounding and guarded integer coercion at the wire or adapter boundary"
160    )]
161    async fn consume_api_key_usage(
162        &self,
163        id: &str,
164        global_rate_limit_enabled: bool,
165    ) -> AuthResult<ConsumeApiKeyResult> {
166        // SQLite must acquire its write reservation before reading usage counters.
167        let transaction = self.begin(true).await?;
168        let exec = Exec::tx(&transaction);
169        let id = id.to_owned();
170        let result = async {
171            let mut select = model::by_id::<Model>(exec, id.as_str());
172            model::limit_one(&mut select);
173            lock_exclusive(&mut select);
174            let Some(model) = exec.fetch_optional::<Model>(select).await? else {
175                return Err(AuthError::not_found("API Key not found"));
176            };
177
178            let now = Utc::now();
179            let mut update = UpdateApiKey::default();
180
181            if let Some(remaining) = model.remaining.map(|value| value.0) {
182                if let (Some(interval), Some(amount)) = (
183                    model.refill_interval.map(|value| value.0),
184                    model.refill_amount.map(|value| value.0),
185                ) && interval != 0.0
186                    && amount != 0.0
187                    && (now.timestamp_millis()
188                        - model
189                            .last_refill_at
190                            .unwrap_or(model.created_at)
191                            .timestamp_millis()) as f64
192                        > interval
193                {
194                    update.remaining = Some(amount - 1.0);
195                    update.last_refill_at = Some(Some(now.to_rfc3339()));
196                } else if remaining > 0.0 {
197                    update.remaining = Some(remaining - 1.0);
198                } else {
199                    return Ok(ConsumeApiKeyResult::UsageExhausted);
200                }
201            }
202
203            if global_rate_limit_enabled && model.rate_limit_enabled.unwrap_or(false) {
204                if let (Some(window), Some(max)) = (
205                    model.rate_limit_time_window.map(|value| value.0),
206                    model.rate_limit_max.map(|value| value.0),
207                ) {
208                    let elapsed = model
209                        .last_request
210                        .map(|last| (now.timestamp_millis() - last.timestamp_millis()) as f64);
211                    if let Some(elapsed) = elapsed
212                        && elapsed <= window
213                        && model.request_count.map_or(0.0, |value| value.0) >= max
214                    {
215                        // TS consumes quota before rejecting a rate-limited request.
216                        // A rejection preserves the rate-limit window and updated_at.
217                        if update.remaining.is_some() {
218                            let updated_at = model.updated_at;
219                            let mut active = apply_update_fields(model.into_active(), update)?;
220                            active.set("updatedAt", updated_at);
221                            _ = model::update::<Model>(exec, &active)
222                                .await?
223                                .ok_or_else(record_not_updated)?;
224                        }
225                        return Ok(ConsumeApiKeyResult::RateLimited {
226                            try_again_in: (window - elapsed).ceil(),
227                        });
228                    }
229
230                    update.request_count =
231                        Some(if elapsed.is_none_or(|elapsed| elapsed > window) {
232                            1.0
233                        } else {
234                            model.request_count.map_or(0.0, |value| value.0) + 1.0
235                        });
236                    update.last_request = Some(Some(now.to_rfc3339()));
237                }
238            } else {
239                update.last_request = Some(Some(now.to_rfc3339()));
240            }
241
242            let active = apply_update_fields(model.into_active(), update)?;
243            let updated = model::update::<Model>(exec, &active)
244                .await?
245                .ok_or_else(record_not_updated)?;
246            Ok(ConsumeApiKeyResult::Allowed(Box::new(ApiKey::from(
247                &updated,
248            ))))
249        }
250        .await;
251        if result.is_ok() {
252            transaction.commit().await?;
253        } else {
254            transaction.rollback().await?;
255        }
256        result
257    }
258
259    async fn delete_api_key(&self, id: &str) -> AuthResult<()> {
260        self.exec()
261            .execute(model::delete_by_id::<Model>(self.exec(), id))
262            .await
263            .map(|_| ())
264    }
265
266    async fn delete_expired_api_keys(&self) -> AuthResult<usize> {
267        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
268        sql.ident(Model::TABLE);
269        sql.push(" WHERE ");
270        sql.column(Model::TABLE, "expiresAt");
271        sql.push(" IS NOT NULL AND ");
272        sql.compare(Model::TABLE, "expiresAt", " < ", Utc::now());
273        let deleted = self.exec().execute(sql).await?;
274        usize::try_from(deleted)
275            .map_err(|_error| AuthError::internal("Affected row count exceeds usize"))
276    }
277}
278
279/// Apply `UpdateApiKey` fields to an API key row.
280fn apply_update_fields(mut active: ActiveRow, update: UpdateApiKey) -> AuthResult<ActiveRow> {
281    if let Some(name) = update.name {
282        active.set("name", Some(name));
283    }
284    if let Some(enabled) = update.enabled {
285        active.set("enabled", enabled);
286    }
287    if let Some(remaining) = update.remaining {
288        active.set("remaining", Some(remaining));
289    }
290    if let Some(rate_limit_enabled) = update.rate_limit_enabled {
291        active.set("rateLimitEnabled", rate_limit_enabled);
292    }
293    if let Some(rate_limit_time_window) = update.rate_limit_time_window {
294        active.set("rateLimitTimeWindow", Some(rate_limit_time_window));
295    }
296    if let Some(rate_limit_max) = update.rate_limit_max {
297        active.set("rateLimitMax", Some(rate_limit_max));
298    }
299    if let Some(refill_interval) = update.refill_interval {
300        active.set("refillInterval", Some(refill_interval));
301    }
302    if let Some(refill_amount) = update.refill_amount {
303        active.set("refillAmount", Some(refill_amount));
304    }
305    if let Some(permissions) = update.permissions {
306        active.set("permissions", Some(permissions));
307    }
308    if let Some(metadata) = update.metadata {
309        active.set("metadata", Some(metadata));
310    }
311    if let Some(expires_at) = update.expires_at {
312        active.set(
313            "expiresAt",
314            parse_optional_rfc3339(expires_at.as_deref(), "expiresAt")?,
315        );
316    }
317    if let Some(last_request) = update.last_request {
318        active.set(
319            "lastRequest",
320            parse_optional_rfc3339(last_request.as_deref(), "lastRequest")?,
321        );
322    }
323    if let Some(request_count) = update.request_count {
324        active.set("requestCount", Some(request_count));
325    }
326    if let Some(last_refill_at) = update.last_refill_at {
327        active.set(
328            "lastRefillAt",
329            parse_optional_rfc3339(last_refill_at.as_deref(), "lastRefillAt")?,
330        );
331    }
332    active.set("updatedAt", Utc::now());
333    Ok(active)
334}
alibi_sqlx/store/api_keys.rs

alibi_sqlx/store/
api_keys.rs
