shared/models/
user_two_factor_code.rs

1use crate::prelude::*;
2use rand::RngExt;
3use serde::{Deserialize, Serialize};
4use sqlx::{Row, postgres::PgRow};
5use std::{collections::BTreeMap, sync::LazyLock};
6
7pub const CODE_VALIDITY_MINUTES: i32 = 10;
8pub const RESEND_COOLDOWN_SECONDS: f64 = 60.0;
9pub const MAX_ATTEMPTS: i32 = 5;
10
11#[derive(Serialize, Deserialize)]
12pub struct UserTwoFactorCode {
13    pub uuid: uuid::Uuid,
14
15    pub attempts: i32,
16
17    pub created: chrono::NaiveDateTime,
18
19    extension_data: super::ModelExtensionData,
20}
21
22impl BaseModel for UserTwoFactorCode {
23    const NAME: &'static str = "user_two_factor_code";
24
25    fn get_extension_list() -> &'static super::ModelExtensionList {
26        static EXTENSIONS: LazyLock<super::ModelExtensionList> =
27            LazyLock::new(|| parking_lot::RwLock::new(Vec::new()));
28
29        &EXTENSIONS
30    }
31
32    fn get_extension_data(&self) -> &super::ModelExtensionData {
33        &self.extension_data
34    }
35
36    #[inline]
37    fn base_columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
38        let prefix = prefix.unwrap_or_default();
39
40        BTreeMap::from([
41            (
42                "user_two_factor_codes.uuid",
43                compact_str::format_compact!("{prefix}uuid"),
44            ),
45            (
46                "user_two_factor_codes.attempts",
47                compact_str::format_compact!("{prefix}attempts"),
48            ),
49            (
50                "user_two_factor_codes.created",
51                compact_str::format_compact!("{prefix}created"),
52            ),
53        ])
54    }
55
56    #[inline]
57    fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
58        let prefix = prefix.unwrap_or_default();
59
60        Ok(Self {
61            uuid: row.try_get(compact_str::format_compact!("{prefix}uuid").as_str())?,
62            attempts: row.try_get(compact_str::format_compact!("{prefix}attempts").as_str())?,
63            created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
64            extension_data: Self::map_extensions(prefix, row)?,
65        })
66    }
67}
68
69impl UserTwoFactorCode {
70    /// Returns the plaintext code; the database only ever holds a bcrypt hash of it.
71    pub async fn create(
72        database: &crate::database::Database,
73        user_uuid: uuid::Uuid,
74    ) -> Result<compact_str::CompactString, anyhow::Error> {
75        let mut transaction = database.write().begin().await?;
76
77        let existing = sqlx::query(
78            r#"
79            SELECT COUNT(*)
80            FROM user_two_factor_codes
81            WHERE user_two_factor_codes.user_uuid = $1
82            AND user_two_factor_codes.created > NOW() - make_interval(secs => $2)
83            "#,
84        )
85        .bind(user_uuid)
86        .bind(RESEND_COOLDOWN_SECONDS)
87        .fetch_optional(&mut *transaction)
88        .await?;
89
90        if let Some(row) = existing
91            && row.get::<i64, _>(0) > 0
92        {
93            return Err(anyhow::anyhow!("a code was already requested recently"));
94        }
95
96        sqlx::query(
97            r#"
98            DELETE FROM user_two_factor_codes
99            WHERE user_two_factor_codes.user_uuid = $1
100            "#,
101        )
102        .bind(user_uuid)
103        .execute(&mut *transaction)
104        .await?;
105
106        let code = compact_str::format_compact!("{:06}", rand::rng().random_range(0..1_000_000));
107
108        sqlx::query(
109            r#"
110            INSERT INTO user_two_factor_codes (user_uuid, code, created)
111            VALUES ($1, crypt($2, gen_salt('bf', 12)), NOW())
112            "#,
113        )
114        .bind(user_uuid)
115        .bind(code.as_str())
116        .execute(&mut *transaction)
117        .await?;
118
119        transaction.commit().await?;
120
121        Ok(code)
122    }
123
124    /// A wrong guess burns an attempt and the code dies at [`MAX_ATTEMPTS`], so the six digit space
125    /// cannot be walked inside the validity window. The row is locked so concurrent guesses cannot
126    /// share an attempt.
127    pub async fn consume(
128        database: &crate::database::Database,
129        user_uuid: uuid::Uuid,
130        code: &str,
131    ) -> Result<bool, crate::database::DatabaseError> {
132        let mut transaction = database.write().begin().await?;
133
134        let row = sqlx::query(
135            r#"
136            SELECT
137                user_two_factor_codes.uuid,
138                user_two_factor_codes.attempts,
139                user_two_factor_codes.code = crypt($2, user_two_factor_codes.code) AS matched
140            FROM user_two_factor_codes
141            WHERE user_two_factor_codes.user_uuid = $1
142            AND user_two_factor_codes.created > NOW() - make_interval(mins => $3)
143            ORDER BY user_two_factor_codes.created DESC
144            LIMIT 1
145            FOR UPDATE
146            "#,
147        )
148        .bind(user_uuid)
149        .bind(code)
150        .bind(CODE_VALIDITY_MINUTES)
151        .fetch_optional(&mut *transaction)
152        .await?;
153
154        let Some(row) = row else {
155            transaction.commit().await?;
156
157            return Ok(false);
158        };
159
160        let uuid: uuid::Uuid = row.try_get("uuid")?;
161        let attempts: i32 = row.try_get("attempts")?;
162        let matched: bool = row.try_get::<Option<bool>, _>("matched")?.unwrap_or(false);
163
164        if matched || attempts + 1 >= MAX_ATTEMPTS {
165            sqlx::query(
166                r#"
167                DELETE FROM user_two_factor_codes
168                WHERE user_two_factor_codes.uuid = $1
169                "#,
170            )
171            .bind(uuid)
172            .execute(&mut *transaction)
173            .await?;
174        } else {
175            sqlx::query(
176                r#"
177                UPDATE user_two_factor_codes
178                SET attempts = attempts + 1
179                WHERE user_two_factor_codes.uuid = $1
180                "#,
181            )
182            .bind(uuid)
183            .execute(&mut *transaction)
184            .await?;
185        }
186
187        transaction.commit().await?;
188
189        Ok(matched)
190    }
191
192    pub async fn delete_by_user_uuid(
193        database: &crate::database::Database,
194        user_uuid: uuid::Uuid,
195    ) -> Result<(), crate::database::DatabaseError> {
196        sqlx::query(
197            r#"
198            DELETE FROM user_two_factor_codes
199            WHERE user_two_factor_codes.user_uuid = $1
200            "#,
201        )
202        .bind(user_uuid)
203        .execute(database.write())
204        .await?;
205
206        Ok(())
207    }
208
209    pub async fn delete_expired(database: &crate::database::Database) -> Result<u64, sqlx::Error> {
210        Ok(sqlx::query(
211            r#"
212            DELETE FROM user_two_factor_codes
213            WHERE user_two_factor_codes.created < NOW() - make_interval(mins => $1)
214            "#,
215        )
216        .bind(CODE_VALIDITY_MINUTES)
217        .execute(database.write())
218        .await?
219        .rows_affected())
220    }
221}
shared/models/user_two_factor_code.rs

shared/models/
user_two_factor_code.rs
