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
14enum Assignment {
16 Value(&'static str, SqlValue),
17 Increment,
18}
19
20fn 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 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}