alibi_sqlx/store/
two_factor.rs

1use super::SqlxStore;
2use super::entities::two_factor::Model;
3use crate::model::{self, ActiveRow, SqlxModel};
4use crate::schema::AuthSchema;
5use crate::sql::Sql;
6use crate::value::SqlValue;
7use alibi_core::error::AuthResult;
8use alibi_core::store::TwoFactorStore;
9use alibi_core::types::{CreateTwoFactor, TwoFactor, UpdateTwoFactor};
10use async_trait::async_trait;
11use chrono::{DateTime, Utc};
12use uuid::Uuid;
13
14/// One assignment in a factor update: a bound value or a SQL expression.
15enum Assignment {
16    Value(&'static str, SqlValue),
17    Increment,
18}
19
20/// Every factor column; SQLite INTEGER affinity retains integral storage, so the
21/// counter is projected as REAL for the typed `f64` decoder under its own name.
22fn factor_columns(sql: &mut Sql, qualified: bool) {
23    for (index, column) in Model::COLUMN_NAMES.iter().copied().enumerate() {
24        if index > 0 {
25            sql.push(", ");
26        }
27        if column == "failed_verification_count" {
28            if qualified {
29                sql.push("CAST(");
30                sql.column(Model::TABLE, column);
31                sql.push(" AS DOUBLE PRECISION) AS ");
32                sql.ident(column);
33            } else {
34                sql.push("CAST(\"failed_verification_count\" AS DOUBLE PRECISION) AS \"failed_verification_count\"");
35            }
36        } else if qualified {
37            sql.column(Model::TABLE, column);
38            sql.push(" AS ");
39            sql.ident(column);
40        } else {
41            sql.ident(column);
42            sql.push(" AS ");
43            sql.ident(column);
44        }
45    }
46}
47
48fn factor_select(sql: &mut Sql) {
49    sql.push("SELECT ");
50    factor_columns(sql, true);
51    sql.push(" FROM ");
52    sql.ident(Model::TABLE);
53}
54
55impl<S: AuthSchema + Send + Sync> SqlxStore<S> {
56    /// RETURNING captures the row in the conditional atomic write, rather
57    /// than a later read that can observe another request's reset/rotation.
58    async fn apply_factor_update(
59        &self,
60        assignments: Vec<Assignment>,
61        filters: Vec<(&'static str, &'static str, SqlValue)>,
62    ) -> AuthResult<Option<TwoFactor>> {
63        let mut sql = Sql::with(self.exec().engine(), "UPDATE ");
64        sql.ident(Model::TABLE);
65        sql.push(" SET ");
66        for (index, assignment) in assignments.into_iter().enumerate() {
67            if index > 0 {
68                sql.push(", ");
69            }
70            match assignment {
71                Assignment::Value(column, value) => {
72                    sql.assign(column, value);
73                }
74                Assignment::Increment => {
75                    sql.ident("failed_verification_count");
76                    sql.push(" = \"failed_verification_count\" + 1");
77                }
78            }
79        }
80        for (index, (column, operator, value)) in filters.into_iter().enumerate() {
81            sql.push(if index == 0 { " WHERE " } else { " AND " });
82            sql.column(Model::TABLE, column);
83            sql.push(operator);
84            sql.bind(value);
85        }
86        sql.push(" RETURNING ");
87        factor_columns(&mut sql, false);
88        Ok(self
89            .exec()
90            .fetch_optional::<Model>(sql)
91            .await?
92            .map(|row| TwoFactor::from(&row)))
93    }
94
95    async fn find_factor(&self, column: &str, value: &str) -> AuthResult<Option<Model>> {
96        let mut sql = Sql::new(self.exec().engine());
97        factor_select(&mut sql);
98        sql.push(" WHERE ");
99        sql.compare(Model::TABLE, column, " = ", value);
100        model::limit_one(&mut sql);
101        self.exec().fetch_optional(sql).await
102    }
103}
104
105#[async_trait]
106impl<S> TwoFactorStore for SqlxStore<S>
107where
108    S: AuthSchema + Send + Sync,
109{
110    async fn create_two_factor(&self, two_factor: CreateTwoFactor) -> AuthResult<TwoFactor> {
111        let now = Utc::now();
112        let mut active = ActiveRow::new();
113        active.set(
114            "id",
115            self.generated_id(self.exec(), "twoFactor", "two_factor", "id")
116                .await?
117                .unwrap_or_else(|| Uuid::new_v4().to_string()),
118        );
119        active.set("secret", two_factor.secret);
120        active.set("backup_codes", two_factor.backup_codes);
121        active.set("user_id", two_factor.user_id);
122        active.set("verified", two_factor.verified);
123        active.set(
124            "failed_verification_count",
125            two_factor.failed_verification_count,
126        );
127        active.set("locked_until", two_factor.locked_until);
128        active.set("created_at", now);
129        active.set("updated_at", now);
130        let mut sql = Sql::with(self.exec().engine(), "INSERT INTO ");
131        sql.ident(Model::TABLE);
132        sql.push(" (");
133        let present = active.present().collect::<Vec<_>>();
134        sql.column_list(&present.iter().map(|(name, _)| *name).collect::<Vec<_>>());
135        sql.push(") VALUES ");
136        sql.bind_list(present.into_iter().map(|(_, value)| value.clone()));
137        sql.push(" RETURNING ");
138        factor_columns(&mut sql, false);
139        self.exec()
140            .fetch_optional::<Model>(sql)
141            .await?
142            .map(|model| TwoFactor::from(&model))
143            .ok_or_else(|| {
144                alibi_core::error::AuthError::internal("Factor insertion returned no row")
145            })
146    }
147
148    async fn get_two_factor_by_user_id(&self, user_id: &str) -> AuthResult<Option<TwoFactor>> {
149        Ok(self
150            .find_factor("user_id", user_id)
151            .await?
152            .map(|model| TwoFactor::from(&model)))
153    }
154
155    async fn update_two_factor(
156        &self,
157        id: &str,
158        update: UpdateTwoFactor,
159    ) -> AuthResult<Option<TwoFactor>> {
160        let unchanged =
161            update.secret.is_none() && update.backup_codes.is_none() && update.verified.is_none();
162        let mut assignments = Vec::new();
163        if let Some(secret) = update.secret {
164            assignments.push(Assignment::Value("secret", secret.into()));
165        }
166        if let Some(codes) = update.backup_codes {
167            assignments.push(Assignment::Value("backup_codes", codes.into()));
168        }
169        if let Some(verified) = update.verified {
170            assignments.push(Assignment::Value("verified", verified.into()));
171        }
172        if unchanged {
173            return Ok(self
174                .find_factor("id", id)
175                .await?
176                .map(|row| TwoFactor::from(&row)));
177        }
178        self.apply_factor_update(assignments, vec![("id", " = ", id.into())])
179            .await
180    }
181
182    async fn increment_two_factor_failure(&self, id: &str) -> AuthResult<Option<TwoFactor>> {
183        self.apply_factor_update(vec![Assignment::Increment], vec![("id", " = ", id.into())])
184            .await
185    }
186
187    async fn set_two_factor_lock_if_count_at_least(
188        &self,
189        id: &str,
190        threshold: f64,
191        until: DateTime<Utc>,
192    ) -> AuthResult<Option<TwoFactor>> {
193        self.apply_factor_update(
194            vec![Assignment::Value("locked_until", until.into())],
195            vec![
196                ("id", " = ", id.into()),
197                ("failed_verification_count", " >= ", threshold.into()),
198            ],
199        )
200        .await
201    }
202
203    async fn clear_expired_two_factor_lock(
204        &self,
205        id: &str,
206        now: DateTime<Utc>,
207    ) -> AuthResult<Option<TwoFactor>> {
208        self.apply_factor_update(
209            vec![
210                Assignment::Value("failed_verification_count", 0.0_f64.into()),
211                Assignment::Value("locked_until", None::<DateTime<Utc>>.into()),
212            ],
213            vec![
214                ("id", " = ", id.into()),
215                ("locked_until", " <= ", now.into()),
216            ],
217        )
218        .await
219    }
220
221    async fn reset_two_factor_failures(&self, id: &str) -> AuthResult<()> {
222        _ = self
223            .apply_factor_update(
224                vec![
225                    Assignment::Value("failed_verification_count", 0.0_f64.into()),
226                    Assignment::Value("locked_until", None::<DateTime<Utc>>.into()),
227                ],
228                vec![("id", " = ", id.into())],
229            )
230            .await?;
231        Ok(())
232    }
233
234    async fn compare_and_swap_two_factor_backup_codes(
235        &self,
236        id: &str,
237        expected: &str,
238        replacement: &str,
239    ) -> AuthResult<bool> {
240        self.apply_factor_update(
241            vec![Assignment::Value("backup_codes", replacement.into())],
242            vec![
243                ("id", " = ", id.into()),
244                ("backup_codes", " = ", expected.into()),
245            ],
246        )
247        .await
248        .map(|row| row.is_some())
249    }
250
251    async fn update_two_factor_backup_codes(
252        &self,
253        user_id: &str,
254        backup_codes: &str,
255    ) -> AuthResult<TwoFactor> {
256        let Some(model) = self.find_factor("user_id", user_id).await? else {
257            return Err(alibi_core::error::AuthError::not_found(
258                "Two-factor settings not found",
259            ));
260        };
261
262        self.apply_factor_update(
263            vec![
264                Assignment::Value("backup_codes", backup_codes.into()),
265                Assignment::Value("updated_at", Utc::now().into()),
266            ],
267            vec![("id", " = ", model.id.into())],
268        )
269        .await?
270        .ok_or_else(|| alibi_core::error::AuthError::not_found("Two-factor settings not found"))
271    }
272
273    async fn delete_two_factor(&self, user_id: &str) -> AuthResult<()> {
274        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
275        sql.ident(Model::TABLE);
276        sql.push(" WHERE ");
277        sql.compare(Model::TABLE, "user_id", " = ", user_id);
278        self.exec().execute(sql).await.map(|_| ())
279    }
280}
alibi_sqlx/store/two_factor.rs

alibi_sqlx/store/
two_factor.rs
