scufflecloud_core/
totp.rs

1use diesel::{ExpressionMethods, QueryDsl, SelectableHelper};
2use diesel_async::RunQueryDsl;
3use tonic_types::{ErrorDetails, StatusExt};
4use totp_rs::{Algorithm, TOTP, TotpUrlError};
5
6use crate::common;
7use crate::models::{MfaTotpCredential, UserId};
8use crate::schema::mfa_totp_credentials;
9use crate::std_ext::{DisplayExt, ResultExt};
10
11const ISSUER: &str = "scuffle.cloud";
12
13#[derive(Debug, thiserror::Error)]
14pub(crate) enum TotpError {
15    #[error("invalid TOTP secret: {0}")]
16    InvalidSecret(#[from] TotpUrlError),
17    #[error("system time error: {0}")]
18    SystemTime(#[from] std::time::SystemTimeError),
19    #[error("invalid TOTP token")]
20    InvalidToken,
21    #[error("failed to generate random secret: {0}")]
22    GenerateSecret(#[from] rand::Error),
23}
24
25pub(crate) fn new_token(account_name: String) -> Result<TOTP, TotpError> {
26    // Generate a new secret for TOTP
27    let secret = common::generate_random_bytes()?;
28    let totp = TOTP::new(
29        Algorithm::SHA1,
30        6,
31        1,
32        30,
33        secret.to_vec(),
34        Some(ISSUER.to_string()),
35        account_name,
36    )?;
37    Ok(totp)
38}
39
40pub(crate) fn verify_token(secret: Vec<u8>, token: &str) -> Result<(), TotpError> {
41    // https://docs.rs/totp-rs/5.7.0/totp_rs/struct.TOTP.html#fields
42    // account_name is not used in the verification process, so we can leave it empty.
43    let totp = TOTP::new(Algorithm::SHA1, 6, 1, 30, secret, Some(ISSUER.to_string()), String::new())?;
44
45    if totp.check_current(token)? {
46        Ok(())
47    } else {
48        Err(TotpError::InvalidToken)
49    }
50}
51
52pub(crate) async fn process_token(
53    tx: &mut diesel_async::AsyncPgConnection,
54    user_id: UserId,
55    token: &str,
56) -> Result<(), tonic::Status> {
57    let credentials = mfa_totp_credentials::dsl::mfa_totp_credentials
58        .filter(mfa_totp_credentials::dsl::user_id.eq(user_id))
59        .select(MfaTotpCredential::as_select())
60        .load::<MfaTotpCredential>(tx)
61        .await
62        .into_tonic_internal_err("failed to query TOTP secrets")?;
63
64    for credential in credentials {
65        match verify_token(credential.secret, token) {
66            Ok(_) => {
67                // Update the last used timestamp
68                diesel::update(mfa_totp_credentials::dsl::mfa_totp_credentials)
69                    .filter(mfa_totp_credentials::dsl::id.eq(credential.id))
70                    .set(mfa_totp_credentials::dsl::last_used_at.eq(chrono::Utc::now()))
71                    .execute(tx)
72                    .await
73                    .into_tonic_internal_err("failed to update TOTP credential")?;
74                return Ok(());
75            }
76            Err(TotpError::InvalidToken) => {} // Try the next secret
77            Err(e) => return Err(e.into_tonic_internal_err("failed to verify TOTP token")),
78        }
79    }
80
81    Err(tonic::Status::with_error_details(
82        tonic::Code::PermissionDenied,
83        "invalid TOTP token",
84        ErrorDetails::new(),
85    ))
86}