umbral_auth/

challenge.rs

//! Short-lived, single-use, hashed-at-rest secrets for the email-verification
//! and password-reset flows. One table, discriminated by `purpose`.

use crate::AuthUser;
use crate::mailer::{OutgoingMail, active_mailer};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use chrono::{DateTime, Utc};
use rand::{Rng, RngCore};
use serde::{Deserialize, Serialize};
use std::time::Duration;
use umbral::db::transaction;
use umbral::orm::{F, ForeignKey};
use umbral::templates::{context, render};

/// Stored discriminator values for [`AuthChallenge::purpose`].
pub const PURPOSE_EMAIL_VERIFY: &str = "email_verify";
pub const PURPOSE_PASSWORD_RESET: &str = "password_reset";

/// One pending challenge. The plaintext (6-digit code or opaque token) is
/// never stored — only `base64(sha256(plaintext))`. Single-use (`used_at`),
/// time-boxed (`expires_at`), and (for codes) attempt-capped (`attempts`).
#[derive(Debug, Clone, sqlx::FromRow, Serialize, Deserialize, umbral::orm::Model)]
pub struct AuthChallenge {
    pub id: i64,
    #[umbral(on_delete = "cascade")]
    pub user_id: ForeignKey<AuthUser>,
    #[umbral(max_length = 32)]
    pub purpose: String,
    #[umbral(max_length = 64)]
    pub secret_hash: String,
    pub expires_at: DateTime<Utc>,
    pub attempts: i32,
    pub used_at: Option<DateTime<Utc>>,
    pub created_at: DateTime<Utc>,
}

// =========================================================================
// Generation helpers
// =========================================================================

/// Generate a zero-padded 6-digit verification code.
///
/// Returns a `String` like `"048392"`. The uniform padding ensures the
/// string is always 6 characters regardless of the random value.
///
/// # Called by
///
/// `start_email_verification` (email verification flow).
pub(crate) fn generate_code() -> String {
    let n: u32 = rand::rngs::OsRng.gen_range(0..1_000_000);
    format!("{n:06}")
}

/// Generate an opaque password-reset token.
///
/// Shape: `umbral_` + 43 URL-safe base64 chars (32 bytes of OS entropy,
/// no padding). The prefix mirrors [`crate::token::TOKEN_PREFIX`] and lets
/// log scrubbers grep for accidentally-logged tokens. The 43-char body
/// gives 256 bits of entropy — brute-force infeasible.
///
/// # Called by
///
/// [`start_password_reset`] (password-reset flow).
pub(crate) fn generate_reset_token() -> String {
    let mut buf = [0u8; 32];
    rand::rngs::OsRng.fill_bytes(&mut buf);
    format!("umbral_{}", URL_SAFE_NO_PAD.encode(buf))
}

/// SHA-256 the plaintext and return a URL-safe base64 digest (43 chars,
/// no padding). Delegates to [`crate::token::digest_token`] so the
/// hashing algorithm is the same as bearer tokens — one implementation,
/// one test surface.
pub(crate) fn hash_secret(plaintext: &str) -> String {
    crate::token::digest_token(plaintext)
}

// =========================================================================
// AuthChallenge ORM methods
// =========================================================================

impl AuthChallenge {
    /// Create and persist a new challenge for `user_id`.
    ///
    /// The `plaintext` (a 6-digit code or opaque token generated by the
    /// caller) is **not** stored — only its `hash_secret` digest reaches
    /// the database. The challenge is valid for `ttl`; `attempts` starts
    /// at 0; `used_at` is `None`.
    ///
    /// # Called by
    ///
    /// Tasks 8 (email verify) and 9 (password reset) call this.
    pub async fn issue(
        user_id: i64,
        purpose: &str,
        plaintext: &str,
        ttl: Duration,
    ) -> Result<AuthChallenge, crate::AuthError> {
        let now = Utc::now();
        let expires_at =
            now + chrono::Duration::from_std(ttl).unwrap_or_else(|_| chrono::Duration::minutes(15));
        let row = AuthChallenge::objects()
            .create(AuthChallenge {
                id: 0, // assigned by the DB
                user_id: ForeignKey::new(user_id),
                purpose: purpose.to_string(),
                secret_hash: hash_secret(plaintext),
                expires_at,
                attempts: 0,
                used_at: None,
                created_at: now,
            })
            .await?;
        Ok(row)
    }

    /// `true` when the challenge has not been consumed (`used_at IS NULL`)
    /// and has not expired (`expires_at > now`). The in-memory check
    /// complements the SQL `USED_AT IS NULL` filter: expired-but-unused
    /// rows are filtered out by this method even if they slipped through
    /// (e.g. a race where `expires_at` passed between the query and the
    /// caller).
    pub fn is_live(&self) -> bool {
        self.used_at.is_none() && self.expires_at > Utc::now()
    }

    /// Return the most-recently-issued active challenge for `(user_id,
    /// purpose)`, or `None` if no unused, unexpired row exists.
    ///
    /// "Active" = `USED_AT IS NULL` in SQL **and** `is_live()` in Rust
    /// (the Rust guard catches challenges whose `expires_at` passed between
    /// the query and this call).
    ///
    /// # ORM note
    ///
    /// `auth_challenge::USED_AT.is_null()` is the SQL `IS NULL` predicate
    /// on the nullable `used_at` column. `order_by` takes a single
    /// `OrderExpr<T>` (not a slice) and can be called multiple times to
    /// append clauses. Mirror the single-arg idiom from `umbral-tasks`.
    pub async fn find_active_for_user(
        user_id: i64,
        purpose: &str,
    ) -> Result<Option<AuthChallenge>, crate::AuthError> {
        let row = AuthChallenge::objects()
            .filter(
                auth_challenge::USER_ID.eq(user_id)
                    & auth_challenge::PURPOSE.eq(purpose)
                    & auth_challenge::USED_AT.is_null(),
            )
            .order_by(auth_challenge::CREATED_AT.desc())
            .first()
            .await?;
        // `Option::filter` keeps the row only if it passes the in-Rust
        // liveness check (handles the expires_at race window).
        Ok(row.filter(|c| c.is_live()))
    }

    /// Return the challenge whose stored hash matches `hash_secret(plaintext)`
    /// for the given `purpose`, or `None` if no active match exists.
    ///
    /// Used by the verification handler: the handler receives the raw
    /// plaintext (code or token), this method hashes it and looks up the
    /// digest — the plaintext never touches the WHERE clause.
    pub async fn find_active_by_secret(
        plaintext: &str,
        purpose: &str,
    ) -> Result<Option<AuthChallenge>, crate::AuthError> {
        let row = AuthChallenge::objects()
            .filter(
                auth_challenge::SECRET_HASH.eq(hash_secret(plaintext))
                    & auth_challenge::PURPOSE.eq(purpose)
                    & auth_challenge::USED_AT.is_null(),
            )
            .first()
            .await?;
        Ok(row.filter(|c| c.is_live()))
    }

    /// Stamp `used_at` with the current time. After this call,
    /// `find_active_for_user` and `find_active_by_secret` will not return
    /// this challenge. Idempotent in the DB sense (a second call just
    /// overwrites with a slightly later timestamp), though callers should
    /// treat the challenge as finished after the first call.
    pub async fn mark_used(&self) -> Result<(), crate::AuthError> {
        let mut delta = serde_json::Map::new();
        delta.insert("used_at".to_string(), serde_json::json!(Utc::now()));
        AuthChallenge::objects()
            .filter(auth_challenge::ID.eq(self.id))
            .update_values(delta)
            .await?;
        Ok(())
    }

    /// Increment `attempts` by 1. Call this on each failed verification
    /// attempt so the caller can gate on a maximum before locking the
    /// challenge. The in-memory `self.attempts` is NOT updated; re-fetch
    /// the row to get the current count.
    ///
    /// The increment is an atomic server-side `SET attempts = attempts + 1`
    /// (via `update_expr`) — two concurrent calls cannot both read the same
    /// stale value and both write the same result, so the attempt cap used
    /// by the brute-force guard in Task 8 cannot be under-counted.
    pub async fn bump_attempts(&self) -> Result<(), crate::AuthError> {
        AuthChallenge::objects()
            .filter(auth_challenge::ID.eq(self.id))
            .update_expr("attempts", F::col("attempts").add(1))
            .await?;
        Ok(())
    }
}

// =========================================================================
// Email-verification flow helpers (Task 8)
// =========================================================================

/// How long a verification code stays live after issue.
const CODE_TTL: Duration = Duration::from_secs(15 * 60);

/// Maximum failed attempts before the challenge is burned.
const MAX_CODE_ATTEMPTS: i32 = 5;

/// Issue a 6-digit verification code for `user`, render
/// `auth/email/verify_code.{html,txt}`, and send via the ambient mailer.
///
/// On success the code is stored (hashed) in `auth_challenge` and the
/// user receives an email with the plaintext code. The code expires in
/// 15 minutes. Calling this multiple times issues fresh challenges
/// (the old row stays in the table but `find_active_for_user` returns
/// the most-recent one by `created_at DESC`).
pub async fn start_email_verification(user: &AuthUser) -> Result<(), crate::AuthError> {
    let code = generate_code();
    AuthChallenge::issue(user.id, PURPOSE_EMAIL_VERIFY, &code, CODE_TTL).await?;
    let ctx = context! { code => code.clone(), username => user.username.clone() };
    let html = render("auth/email/verify_code.html", &ctx)
        .map_err(|e| crate::AuthError::Template(e.to_string()))?;
    let text = render("auth/email/verify_code.txt", &ctx)
        .map_err(|e| crate::AuthError::Template(e.to_string()))?;
    active_mailer()
        .send(OutgoingMail {
            to: user.email.clone(),
            username: user.username.clone(),
            kind: crate::mailer::MailKind::EmailVerification { code },
            subject: "Verify your email".into(),
            html,
            text,
        })
        .await
        .map_err(|e| crate::AuthError::Mail(e.to_string()))?;
    Ok(())
}

/// Verify `email` against the submitted `code`.
///
/// Look up the user by email (None → `InvalidChallenge`); find their
/// active `email_verify` challenge (None → `InvalidChallenge`); gate on
/// the attempt cap (≥ 5 → burn + `InvalidChallenge`); compare hashes
/// (mismatch → bump attempts + `InvalidChallenge`); on match: mark used
/// and stamp `email_verified_at = now` on the user row.
///
/// All failure arms return the same opaque `InvalidChallenge` error —
/// no account enumeration through the verification surface.
pub async fn verify_email(email: &str, code: &str) -> Result<(), crate::AuthError> {
    let Some(user) = AuthUser::objects()
        .filter(crate::auth_user::EMAIL.eq(email.to_string()))
        .first()
        .await?
    else {
        return Err(crate::AuthError::InvalidChallenge);
    };

    let Some(challenge) =
        AuthChallenge::find_active_for_user(user.id, PURPOSE_EMAIL_VERIFY).await?
    else {
        return Err(crate::AuthError::InvalidChallenge);
    };

    if challenge.attempts >= MAX_CODE_ATTEMPTS {
        challenge.mark_used().await?; // burn the exhausted challenge
        return Err(crate::AuthError::InvalidChallenge);
    }

    if hash_secret(code) != challenge.secret_hash {
        challenge.bump_attempts().await?;
        return Err(crate::AuthError::InvalidChallenge);
    }

    // Consume the challenge and stamp email_verified_at in one atomic
    // transaction. Without this, a crash between the two writes would
    // leave either (a) a still-live, replayable challenge while the
    // user is already verified, or (b) a consumed challenge while the
    // user is still unverified. Both are security bugs; the transaction
    // prevents both.
    let challenge_id = challenge.id;
    let user_id = user.id;
    transaction(|tx| {
        Box::pin(async move {
            let mut mark_delta = serde_json::Map::new();
            mark_delta.insert("used_at".to_string(), serde_json::json!(Utc::now()));
            AuthChallenge::objects()
                .filter(auth_challenge::ID.eq(challenge_id))
                .on_tx(tx)
                .update_values(mark_delta)
                .await?;

            let mut verify_delta = serde_json::Map::new();
            verify_delta.insert(
                "email_verified_at".to_string(),
                serde_json::json!(Utc::now()),
            );
            AuthUser::objects()
                .filter(crate::auth_user::ID.eq(user_id))
                .on_tx(tx)
                .update_values(verify_delta)
                .await?;

            Ok::<_, crate::AuthError>(())
        })
    })
    .await?;

    Ok(())
}

// =========================================================================
// Password-reset flow helpers (Task 9)
// =========================================================================

/// How long a password-reset token stays valid after issue.
const RESET_TTL: Duration = Duration::from_secs(60 * 60); // 1 hour

/// Issue a password-reset token for the account that owns `email`,
/// render `auth/email/reset_link.{html,txt}` with the full reset URL,
/// and send via the ambient mailer.
///
/// If no account exists for `email`, returns `Ok(())` silently — the
/// caller gets no information about whether an account was found. This
/// prevents account enumeration through the password-reset surface: an
/// attacker who submits an arbitrary email address sees the same outcome
/// as a legitimate user.
///
/// The token is an opaque [`generate_reset_token`] value stored
/// (hashed) in `auth_challenge`. It expires in 1 hour.
pub async fn start_password_reset(
    email: &str,
    reset_url_base: &str,
) -> Result<(), crate::AuthError> {
    // Silent on unknown email — never reveal whether an account exists.
    let Some(user) = crate::AuthUser::objects()
        .filter(crate::auth_user::EMAIL.eq(email.to_string()))
        .first()
        .await?
    else {
        return Ok(());
    };
    let token = generate_reset_token();
    AuthChallenge::issue(user.id, PURPOSE_PASSWORD_RESET, &token, RESET_TTL).await?;
    let reset_url = format!("{reset_url_base}?token={token}");
    let ctx = context! { reset_url => reset_url.clone(), username => user.username.clone() };
    let html = render("auth/email/reset_link.html", &ctx)
        .map_err(|e| crate::AuthError::Template(e.to_string()))?;
    let text = render("auth/email/reset_link.txt", &ctx)
        .map_err(|e| crate::AuthError::Template(e.to_string()))?;
    active_mailer()
        .send(OutgoingMail {
            to: user.email.clone(),
            username: user.username.clone(),
            kind: crate::mailer::MailKind::PasswordReset { reset_url },
            subject: "Reset your password".into(),
            html,
            text,
        })
        .await
        .map_err(|e| crate::AuthError::Mail(e.to_string()))?;
    Ok(())
}

/// Consume a password-reset token and set the account's password to
/// `new_password`.
///
/// # Validation
///
/// The candidate password must pass the ambient [`crate::PasswordPolicy`]
/// (checked via [`crate::validate_password`]). On failure the error is
/// [`crate::AuthError::WeakPassword`] with every rejection reason; the
/// challenge is NOT consumed so the user can retry with a stronger password.
///
/// # Atomicity
///
/// The password-hash update and the challenge consume are written in a
/// single transaction so a crash between the two cannot leave a window
/// where the old password still works on an already-consumed token or,
/// conversely, a live token against an already-updated password.
///
/// # Revocations (best-effort, post-commit)
///
/// After the transaction commits, all bearer tokens and sessions for the
/// user are revoked ("log out everywhere"). A revocation failure does
/// **not** roll back the password change — the password is already
/// updated at that point, and failing the call would confuse the caller.
/// Revocation failures are logged at ERROR level (via `tracing::error!`) so
/// they are observable in prod without aborting the successful reset.
///
/// # Errors
///
/// Returns [`crate::AuthError::InvalidChallenge`] for ALL failure arms
/// that involve the token (not found, expired, already used, no user) —
/// the same opaque error prevents oracle attacks.
pub async fn reset_password(token: &str, new_password: &str) -> Result<(), crate::AuthError> {
    let Some(challenge) =
        AuthChallenge::find_active_by_secret(token, PURPOSE_PASSWORD_RESET).await?
    else {
        return Err(crate::AuthError::InvalidChallenge);
    };
    // ForeignKey<AuthUser>::id() returns the i64 primary key.
    let user_id: i64 = challenge.user_id.id();
    let Some(user) = crate::AuthUser::objects()
        .filter(crate::auth_user::ID.eq(user_id))
        .first()
        .await?
    else {
        // The user row was deleted after the challenge was issued.
        return Err(crate::AuthError::InvalidChallenge);
    };
    // Enforce the strength policy before touching the DB.
    crate::validate_password(
        new_password,
        &crate::PasswordContext::new(Some(&user.username), Some(&user.email)),
    )
    .map_err(crate::AuthError::WeakPassword)?;
    let hash = crate::hash_password(new_password)?;

    // Atomically: update the password hash AND consume the challenge.
    let challenge_id = challenge.id;
    transaction(|tx| {
        let hash = hash.clone();
        Box::pin(async move {
            let mut pw_delta = serde_json::Map::new();
            pw_delta.insert("password_hash".to_string(), serde_json::json!(hash));
            crate::AuthUser::objects()
                .filter(crate::auth_user::ID.eq(user_id))
                .on_tx(tx)
                .update_values(pw_delta)
                .await?;

            let mut mark_delta = serde_json::Map::new();
            mark_delta.insert("used_at".to_string(), serde_json::json!(Utc::now()));
            AuthChallenge::objects()
                .filter(auth_challenge::ID.eq(challenge_id))
                .on_tx(tx)
                .update_values(mark_delta)
                .await?;

            Ok::<_, crate::AuthError>(())
        })
    })
    .await?;

    // Post-commit best-effort revocations. A reset implies possible account
    // compromise; "log out everywhere" is the safe response. Failures here
    // do NOT un-change the password — the hash is already updated. Errors are
    // logged so a failed revocation is observable in production.
    if let Err(e) = crate::token::AuthToken::objects()
        .filter(crate::token::auth_token::USER_ID.eq(user_id))
        .delete()
        .await
    {
        tracing::error!(user_id, error = %e, "password reset: failed to revoke bearer tokens");
    }
    if let Err(e) = umbral_sessions::revoke_user_sessions(&user_id.to_string()).await {
        tracing::error!(user_id, error = %e, "password reset: failed to revoke sessions");
    }

    Ok(())
}

// =========================================================================
// Unit tests — keep pub(crate) helpers exercised inside the module
// =========================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn generate_code_is_six_ascii_digits() {
        for _ in 0..20 {
            let code = generate_code();
            assert_eq!(
                code.len(),
                6,
                "generate_code must produce exactly 6 chars; got {code:?}"
            );
            assert!(
                code.chars().all(|c| c.is_ascii_digit()),
                "generate_code must contain only ASCII digits; got {code:?}"
            );
        }
    }

    #[test]
    fn generate_code_zero_pads_small_numbers() {
        // Deterministic proof that the {:06} format spec zero-pads small values.
        assert_eq!(format!("{:06}", 48u32), "000048");
        // Also confirm the invariant holds across random samples.
        let code = generate_code();
        assert_eq!(code.len(), 6);
    }

    #[test]
    fn generate_reset_token_has_prefix_and_length() {
        let tok = generate_reset_token();
        assert!(
            tok.starts_with("umbral_"),
            "token must start with 'umbral_'; got {tok:?}"
        );
        // prefix (7) + base64(32 bytes, no padding) = 7 + 43 = 50
        assert_eq!(
            tok.len(),
            50,
            "token must be 50 chars (7 prefix + 43 base64); got {tok:?}"
        );
    }

    #[test]
    fn generate_reset_token_is_unique_per_call() {
        let a = generate_reset_token();
        let b = generate_reset_token();
        assert_ne!(a, b, "two consecutive tokens must not collide");
    }

    #[test]
    fn hash_secret_is_deterministic_and_delegates_to_digest_token() {
        let a = hash_secret("483920");
        let b = hash_secret("483920");
        let c = hash_secret("999999");
        assert_eq!(a, b, "hash_secret must be deterministic");
        assert_ne!(a, c, "different inputs must produce different digests");
        assert_eq!(
            a.len(),
            43,
            "SHA-256 in URL-safe base64 (no pad) is 43 chars"
        );
    }
}