Skip to main content

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 == "failedVerificationCount" {
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(\"failedVerificationCount\" AS DOUBLE PRECISION) AS \"failedVerificationCount\"");
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("failedVerificationCount");
76                    sql.push(" = \"failedVerificationCount\" + 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 mut active = ActiveRow::new();
112        active.set(
113            "id",
114            self.generated_id(self.exec(), "twoFactor", "twoFactor", "id")
115                .await?
116                .unwrap_or_else(|| Uuid::new_v4().to_string()),
117        );
118        active.set("secret", two_factor.secret);
119        active.set("backupCodes", two_factor.backup_codes);
120        active.set("userId", two_factor.user_id);
121        active.set("verified", two_factor.verified);
122        active.set(
123            "failedVerificationCount",
124            two_factor.failed_verification_count,
125        );
126        active.set("lockedUntil", two_factor.locked_until);
127
128        let mut sql = Sql::with(self.exec().engine(), "INSERT INTO ");
129        sql.ident(Model::TABLE);
130        sql.push(" (");
131        let present = active.present().collect::<Vec<_>>();
132        sql.column_list(&present.iter().map(|(name, _)| *name).collect::<Vec<_>>());
133        sql.push(") VALUES ");
134        sql.bind_list(present.into_iter().map(|(_, value)| value.clone()));
135        sql.push(" RETURNING ");
136        factor_columns(&mut sql, false);
137        self.exec()
138            .fetch_optional::<Model>(sql)
139            .await?
140            .map(|model| TwoFactor::from(&model))
141            .ok_or_else(|| {
142                alibi_core::error::AuthError::internal("Factor insertion returned no row")
143            })
144    }
145
146    async fn get_two_factor_by_user_id(&self, user_id: &str) -> AuthResult<Option<TwoFactor>> {
147        Ok(self
148            .find_factor("userId", user_id)
149            .await?
150            .map(|model| TwoFactor::from(&model)))
151    }
152
153    async fn update_two_factor(
154        &self,
155        id: &str,
156        update: UpdateTwoFactor,
157    ) -> AuthResult<Option<TwoFactor>> {
158        let unchanged =
159            update.secret.is_none() && update.backup_codes.is_none() && update.verified.is_none();
160        let mut assignments = Vec::new();
161        if let Some(secret) = update.secret {
162            assignments.push(Assignment::Value("secret", secret.into()));
163        }
164        if let Some(codes) = update.backup_codes {
165            assignments.push(Assignment::Value("backupCodes", codes.into()));
166        }
167        if let Some(verified) = update.verified {
168            assignments.push(Assignment::Value("verified", verified.into()));
169        }
170        if unchanged {
171            return Ok(self
172                .find_factor("id", id)
173                .await?
174                .map(|row| TwoFactor::from(&row)));
175        }
176        self.apply_factor_update(assignments, vec![("id", " = ", id.into())])
177            .await
178    }
179
180    async fn increment_two_factor_failure(&self, id: &str) -> AuthResult<Option<TwoFactor>> {
181        self.apply_factor_update(vec![Assignment::Increment], vec![("id", " = ", id.into())])
182            .await
183    }
184
185    async fn set_two_factor_lock_if_count_at_least(
186        &self,
187        id: &str,
188        threshold: f64,
189        until: DateTime<Utc>,
190    ) -> AuthResult<Option<TwoFactor>> {
191        self.apply_factor_update(
192            vec![Assignment::Value("lockedUntil", until.into())],
193            vec![
194                ("id", " = ", id.into()),
195                ("failedVerificationCount", " >= ", threshold.into()),
196            ],
197        )
198        .await
199    }
200
201    async fn clear_expired_two_factor_lock(
202        &self,
203        id: &str,
204        now: DateTime<Utc>,
205    ) -> AuthResult<Option<TwoFactor>> {
206        self.apply_factor_update(
207            vec![
208                Assignment::Value("failedVerificationCount", 0.0_f64.into()),
209                Assignment::Value("lockedUntil", None::<DateTime<Utc>>.into()),
210            ],
211            vec![
212                ("id", " = ", id.into()),
213                ("lockedUntil", " <= ", now.into()),
214            ],
215        )
216        .await
217    }
218
219    async fn reset_two_factor_failures(&self, id: &str) -> AuthResult<()> {
220        _ = self
221            .apply_factor_update(
222                vec![
223                    Assignment::Value("failedVerificationCount", 0.0_f64.into()),
224                    Assignment::Value("lockedUntil", None::<DateTime<Utc>>.into()),
225                ],
226                vec![("id", " = ", id.into())],
227            )
228            .await?;
229        Ok(())
230    }
231
232    async fn compare_and_swap_two_factor_backup_codes(
233        &self,
234        id: &str,
235        expected: &str,
236        replacement: &str,
237    ) -> AuthResult<bool> {
238        self.apply_factor_update(
239            vec![Assignment::Value("backupCodes", replacement.into())],
240            vec![
241                ("id", " = ", id.into()),
242                ("backupCodes", " = ", expected.into()),
243            ],
244        )
245        .await
246        .map(|row| row.is_some())
247    }
248
249    async fn update_two_factor_backup_codes(
250        &self,
251        user_id: &str,
252        backup_codes: &str,
253    ) -> AuthResult<TwoFactor> {
254        let Some(model) = self.find_factor("userId", user_id).await? else {
255            return Err(alibi_core::error::AuthError::not_found(
256                "Two-factor settings not found",
257            ));
258        };
259
260        self.apply_factor_update(
261            vec![Assignment::Value("backupCodes", backup_codes.into())],
262            vec![("id", " = ", model.id.into())],
263        )
264        .await?
265        .ok_or_else(|| alibi_core::error::AuthError::not_found("Two-factor settings not found"))
266    }
267
268    async fn delete_two_factor(&self, user_id: &str) -> AuthResult<()> {
269        let mut sql = Sql::with(self.exec().engine(), "DELETE FROM ");
270        sql.ident(Model::TABLE);
271        sql.push(" WHERE ");
272        sql.compare(Model::TABLE, "userId", " = ", user_id);
273        self.exec().execute(sql).await.map(|_| ())
274    }
275}