shared/models/
user_recovery_code.rs

1use crate::prelude::*;
2use rand::distr::SampleString;
3use serde::{Deserialize, Serialize};
4use sqlx::{Row, postgres::PgRow};
5use std::{collections::BTreeMap, sync::LazyLock};
6
7#[derive(Serialize, Deserialize)]
8pub struct UserRecoveryCode {
9    pub code: compact_str::CompactString,
10
11    pub created: chrono::NaiveDateTime,
12
13    extension_data: super::ModelExtensionData,
14}
15
16impl BaseModel for UserRecoveryCode {
17    const NAME: &'static str = "user_recovery_code";
18
19    fn get_extension_list() -> &'static super::ModelExtensionList {
20        static EXTENSIONS: LazyLock<super::ModelExtensionList> =
21            LazyLock::new(|| parking_lot::RwLock::new(Vec::new()));
22
23        &EXTENSIONS
24    }
25
26    fn get_extension_data(&self) -> &super::ModelExtensionData {
27        &self.extension_data
28    }
29
30    #[inline]
31    fn base_columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
32        let prefix = prefix.unwrap_or_default();
33
34        BTreeMap::from([
35            (
36                "user_recovery_codes.code",
37                compact_str::format_compact!("{prefix}code"),
38            ),
39            (
40                "user_recovery_codes.created",
41                compact_str::format_compact!("{prefix}created"),
42            ),
43        ])
44    }
45
46    #[inline]
47    fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError> {
48        let prefix = prefix.unwrap_or_default();
49
50        Ok(Self {
51            code: row.try_get(compact_str::format_compact!("{prefix}code").as_str())?,
52            created: row.try_get(compact_str::format_compact!("{prefix}created").as_str())?,
53            extension_data: Self::map_extensions(prefix, row)?,
54        })
55    }
56}
57
58impl UserRecoveryCode {
59    pub async fn create_all(
60        database: &crate::database::Database,
61        user_uuid: uuid::Uuid,
62    ) -> Result<Vec<String>, crate::database::DatabaseError> {
63        let mut codes = Vec::new();
64        codes.reserve_exact(10);
65
66        let mut transaction = database.write().begin().await?;
67
68        for _ in 0..10 {
69            let code = rand::distr::Alphanumeric.sample_string(&mut rand::rng(), 10);
70
71            sqlx::query(
72                r#"
73                INSERT INTO user_recovery_codes (user_uuid, code, created)
74                VALUES ($1, $2, NOW())
75                "#,
76            )
77            .bind(user_uuid)
78            .bind(&code)
79            .execute(&mut *transaction)
80            .await?;
81
82            codes.push(code);
83        }
84
85        transaction.commit().await?;
86
87        Ok(codes)
88    }
89
90    /// Enabling a second two factor method must not invalidate codes the user already wrote down.
91    pub async fn create_all_if_absent(
92        database: &crate::database::Database,
93        user_uuid: uuid::Uuid,
94    ) -> Result<Vec<String>, crate::database::DatabaseError> {
95        let existing: Vec<String> = sqlx::query(
96            r#"
97            SELECT user_recovery_codes.code
98            FROM user_recovery_codes
99            WHERE user_recovery_codes.user_uuid = $1
100            "#,
101        )
102        .bind(user_uuid)
103        .fetch_all(database.read())
104        .await?
105        .into_iter()
106        .map(|row| row.get::<String, _>("code"))
107        .collect();
108
109        if !existing.is_empty() {
110            return Ok(existing);
111        }
112
113        Self::create_all(database, user_uuid).await
114    }
115
116    pub async fn delete_by_user_uuid_code(
117        database: &crate::database::Database,
118        user_uuid: uuid::Uuid,
119        code: &str,
120    ) -> Result<Option<Self>, crate::database::DatabaseError> {
121        let row = sqlx::query(sqlx::AssertSqlSafe(format!(
122            r#"
123            DELETE FROM user_recovery_codes
124            WHERE user_recovery_codes.user_uuid = $1 AND user_recovery_codes.code = $2
125            RETURNING {}
126            "#,
127            Self::columns_sql(None)
128        )))
129        .bind(user_uuid)
130        .bind(code)
131        .fetch_optional(database.write())
132        .await?;
133
134        row.try_map(|row| Self::map(None, &row))
135    }
136
137    pub async fn delete_by_user_uuid(
138        database: &crate::database::Database,
139        user_uuid: uuid::Uuid,
140    ) -> Result<(), crate::database::DatabaseError> {
141        sqlx::query(
142            r#"
143            DELETE FROM user_recovery_codes
144            WHERE user_recovery_codes.user_uuid = $1
145            "#,
146        )
147        .bind(user_uuid)
148        .execute(database.write())
149        .await?;
150
151        Ok(())
152    }
153}
shared/models/user_recovery_code.rs

shared/models/
user_recovery_code.rs
