Skip to main content

authx_plugins/totp/
service.rs

1use std::sync::Arc;
2
3use rand::Rng;
4use totp_rs::{Algorithm, Secret, TOTP};
5use tracing::instrument;
6use uuid::Uuid;
7
8use authx_core::{
9    crypto::sha256_hex,
10    error::{AuthError, Result},
11    models::{CreateCredential, CredentialKind},
12};
13use authx_storage::ports::{CredentialRepository, UserRepository};
14
15/// TOTP service — handles setup, verification, and backup codes.
16///
17/// Secrets are stored as `CredentialKind::Passkey` entries (reusing the
18/// credential row with kind `"passkey"` to avoid schema changes).
19/// The secret is stored as a base32-encoded string in `credential_hash`.
20///
21/// Backup codes are SHA-256 hashed and stored in metadata as a JSON array.
22pub struct TotpService<S> {
23    storage: S,
24    app_name: Arc<str>,
25}
26
27/// Returned when the user first enables TOTP.
28#[derive(Debug)]
29pub struct TotpSetup {
30    /// Base32-encoded secret — store server-side and show once to user.
31    pub secret_base32: String,
32    /// `otpauth://totp/...` URI for QR code generation.
33    pub otpauth_uri: String,
34    /// One-time backup codes (show to user, hash before storing).
35    pub backup_codes: Vec<String>,
36}
37
38pub struct TotpVerifyRequest {
39    pub user_id: Uuid,
40    pub code: String,
41}
42
43impl<S> TotpService<S>
44where
45    S: UserRepository + CredentialRepository + Clone + Send + Sync + 'static,
46{
47    pub fn new(storage: S, app_name: impl Into<Arc<str>>) -> Self {
48        Self {
49            storage,
50            app_name: app_name.into(),
51        }
52    }
53
54    /// Generate a new TOTP secret and backup codes for a user.
55    /// Call `confirm_setup` with the first code before persisting.
56    #[instrument(skip(self), fields(user_id = %user_id))]
57    pub async fn begin_setup(&self, user_id: Uuid) -> Result<TotpSetup> {
58        let user = UserRepository::find_by_id(&self.storage, user_id)
59            .await?
60            .ok_or(AuthError::UserNotFound)?;
61
62        let secret = Secret::generate_secret();
63        let secret_base32 = secret.to_encoded().to_string();
64
65        let totp = build_totp(&secret_base32, &user.email, &self.app_name)?;
66        let otpauth_uri = totp.get_url();
67
68        let backup_codes = generate_backup_codes(8);
69
70        tracing::info!(user_id = %user_id, "totp setup initiated");
71        Ok(TotpSetup {
72            secret_base32,
73            otpauth_uri,
74            backup_codes,
75        })
76    }
77
78    /// Confirm the user can produce a valid code, then persist the secret.
79    #[instrument(skip(self, setup, code), fields(user_id = %user_id))]
80    pub async fn confirm_setup(&self, user_id: Uuid, setup: &TotpSetup, code: &str) -> Result<()> {
81        let user = UserRepository::find_by_id(&self.storage, user_id)
82            .await?
83            .ok_or(AuthError::UserNotFound)?;
84
85        let totp = build_totp(&setup.secret_base32, &user.email, &self.app_name)?;
86        if !totp
87            .check_current(code)
88            .map_err(|_| AuthError::InvalidToken)?
89        {
90            return Err(AuthError::InvalidToken);
91        }
92
93        let hashed_codes: Vec<String> = setup
94            .backup_codes
95            .iter()
96            .map(|c| sha256_hex(c.as_bytes()))
97            .collect();
98
99        CredentialRepository::create(
100            &self.storage,
101            CreateCredential {
102                user_id,
103                kind: CredentialKind::Passkey,
104                credential_hash: setup.secret_base32.clone(),
105                metadata: Some(serde_json::json!({ "backup_codes": hashed_codes })),
106            },
107        )
108        .await?;
109
110        tracing::info!(user_id = %user_id, "totp enabled");
111        Ok(())
112    }
113
114    /// Verify a TOTP code (or a backup code) during sign-in.
115    #[instrument(skip(self, req), fields(user_id = %req.user_id))]
116    pub async fn verify(&self, req: TotpVerifyRequest) -> Result<()> {
117        let user = UserRepository::find_by_id(&self.storage, req.user_id)
118            .await?
119            .ok_or(AuthError::UserNotFound)?;
120
121        let cred = CredentialRepository::find_by_user_and_kind(
122            &self.storage,
123            req.user_id,
124            CredentialKind::Passkey,
125        )
126        .await?
127        .ok_or(AuthError::InvalidToken)?;
128
129        let totp = build_totp(&cred.credential_hash, &user.email, &self.app_name)?;
130        if totp
131            .check_current(&req.code)
132            .map_err(|_| AuthError::InvalidToken)?
133        {
134            tracing::info!(user_id = %req.user_id, "totp verified");
135            return Ok(());
136        }
137
138        // Try backup codes — single-use. Remove the consumed code from storage.
139        let code_hash = sha256_hex(req.code.as_bytes());
140        let mut codes: Vec<String> = cred
141            .metadata
142            .get("backup_codes")
143            .and_then(|v| serde_json::from_value(v.clone()).ok())
144            .unwrap_or_default();
145
146        if let Some(pos) = codes.iter().position(|c| c == &code_hash) {
147            codes.remove(pos);
148
149            // Persist the updated (shorter) backup-code list.
150            CredentialRepository::delete_by_user_and_kind(
151                &self.storage,
152                req.user_id,
153                CredentialKind::Passkey,
154            )
155            .await?;
156            CredentialRepository::create(
157                &self.storage,
158                CreateCredential {
159                    user_id: req.user_id,
160                    kind: CredentialKind::Passkey,
161                    credential_hash: cred.credential_hash,
162                    metadata: Some(serde_json::json!({ "backup_codes": codes })),
163                },
164            )
165            .await?;
166
167            tracing::info!(user_id = %req.user_id, remaining = codes.len(), "totp: backup code consumed");
168            return Ok(());
169        }
170
171        tracing::warn!(user_id = %req.user_id, "totp verification failed");
172        Err(AuthError::InvalidToken)
173    }
174
175    /// Remove TOTP from a user account.
176    #[instrument(skip(self), fields(user_id = %user_id))]
177    pub async fn disable(&self, user_id: Uuid) -> Result<()> {
178        CredentialRepository::delete_by_user_and_kind(
179            &self.storage,
180            user_id,
181            CredentialKind::Passkey,
182        )
183        .await?;
184        tracing::info!(user_id = %user_id, "totp disabled");
185        Ok(())
186    }
187
188    /// Returns `true` if the user has TOTP enabled.
189    pub async fn is_enabled(&self, user_id: Uuid) -> Result<bool> {
190        Ok(CredentialRepository::find_by_user_and_kind(
191            &self.storage,
192            user_id,
193            CredentialKind::Passkey,
194        )
195        .await?
196        .is_some())
197    }
198}
199
200// ── Helpers ───────────────────────────────────────────────────────────────────
201
202fn build_totp(secret_base32: &str, email: &str, app_name: &str) -> Result<TOTP> {
203    let secret = Secret::Encoded(secret_base32.to_owned());
204    TOTP::new(
205        Algorithm::SHA1,
206        6,
207        1,
208        30,
209        secret.to_bytes().map_err(|_| AuthError::InvalidToken)?,
210        Some(app_name.to_owned()),
211        email.to_owned(),
212    )
213    .map_err(|e| AuthError::Internal(format!("totp init: {e}")))
214}
215
216fn generate_backup_codes(count: usize) -> Vec<String> {
217    let mut rng = rand::thread_rng();
218    (0..count)
219        .map(|_| {
220            (0..8)
221                .map(|_| {
222                    let idx = rng.gen_range(0..36u8);
223                    (if idx < 10 {
224                        b'0' + idx
225                    } else {
226                        b'A' + idx - 10
227                    }) as char
228                })
229                .collect::<String>()
230        })
231        .collect()
232}