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

alibi_sqlx/store/
two_factor.rs
