eric / rustwarden

//! TOTP two-factor authentication service
//!
//! Provides methods for generating, verifying, and managing TOTP 2FA
//! including backup codes and login challenge flow.

use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
use std::time::Instant;

use chrono::{Duration, Utc};
use sea_orm::sea_query::Expr;
use sea_orm::{
    ActiveModelTrait, ColumnTrait, Condition, DatabaseConnection, EntityTrait, NotSet, QueryFilter,
    Set,
};
use totp_rs::{Algorithm, Secret, TOTP};
use uuid::Uuid;

use crate::error::AppError;
use crate::models::{totp_challenge, user_totp};
use crate::services::crypto_service::CryptoService;

const BACKUP_CODE_COUNT: usize = 10;
const BACKUP_CODE_LENGTH: usize = 8;
const CHALLENGE_EXPIRY_MINUTES: i64 = 5;
const MAX_CHALLENGE_ATTEMPTS: i32 = 5;
/// Cap of *summed* challenge attempts per user over 15 minutes.
const MAX_ATTEMPTS_PER_WINDOW: i32 = 15;
const REPLAY_GUARD_TTL_SECS: u64 = 90;

/// Returns true when the summed challenge attempts reach [`MAX_ATTEMPTS_PER_WINDOW`].
pub(crate) fn totp_window_exceeded(attempts: impl Iterator<Item = i32>) -> bool {
    attempts.sum::<i32>() >= MAX_ATTEMPTS_PER_WINDOW
}

/// Returns the matching TOTP time-step (`unix_time / 30`) if `code` is valid
/// for the current step or skew ±1.
pub(crate) fn matching_totp_step(secret_base32: &str, code: &str) -> Option<u64> {
    let secret = Secret::Encoded(secret_base32.to_string());
    let totp = TOTP::new(
        Algorithm::SHA1,
        6,
        1,
        30,
        secret.to_bytes().ok()?,
        None,
        String::new(),
    )
    .ok()?;

    // `check` accounts for skew ±1 and compares with constant_time_eq.
    // Identify the matching window only after that reject path.
    let now = Utc::now().timestamp().max(0) as u64;
    if !totp.check(code, now) {
        return None;
    }
    for delta in [-1_i64, 0, 1] {
        let Some(step) = (now / 30).checked_add_signed(delta) else {
            continue;
        };
        if totp.generate(step * 30) == code {
            return Some(step);
        }
    }
    None
}

/// Accepts a `(user_id, step)` pair once per 90s in-process window.
///
/// Returns `false` when the same step was already accepted (replay).
pub(crate) fn totp_replay_guard_accept(user_id: Uuid, step: u64) -> bool {
    static CACHE: OnceLock<Mutex<HashMap<(Uuid, u64), Instant>>> = OnceLock::new();
    let mut cache = CACHE
        .get_or_init(|| Mutex::new(HashMap::new()))
        .lock()
        .unwrap_or_else(|e| e.into_inner());

    let ttl = std::time::Duration::from_secs(REPLAY_GUARD_TTL_SECS);
    cache.retain(|_, instant| instant.elapsed() < ttl);

    let key = (user_id, step);
    if cache.contains_key(&key) {
        return false;
    }

    cache.insert(key, Instant::now());
    true
}

/// Finds the first bcrypt-matching backup-code hash index off the async runtime.
///
/// Short-circuits on the first match: the index is not a secret (the attacker
/// already knows the plaintext code they submitted). Join errors become
/// [`AppError::Internal`]. A bcrypt verification failure is a non-match.
async fn find_matching_backup_index(
    code: String,
    hashed_codes: Vec<String>,
) -> Result<Option<usize>, AppError> {
    tokio::task::spawn_blocking(move || {
        let mut found = None;
        for (i, hashed) in hashed_codes.iter().enumerate() {
            if bcrypt::verify(&code, hashed).unwrap_or(false) {
                found = Some(i);
                break; // allowed: index is not a secret; attacker knows their code
            }
        }
        found
    })
    .await
    .map_err(|e| AppError::Internal(format!("Blocking task failed: {e}")))
}

/// Verifies a plaintext backup code against a single bcrypt hash off the async runtime.
///
/// Join errors from the blocking thread pool become [`AppError::Internal`].
/// A bcrypt verification failure is treated as a non-match (`false`).
#[cfg(test)]
pub(crate) async fn verify_one_backup(code: String, hashed: String) -> Result<bool, AppError> {
    Ok(find_matching_backup_index(code, vec![hashed])
        .await?
        .is_some())
}

/// Response from TOTP setup containing data needed to display QR code
pub struct TotpSetupResponse {
    /// Base32-encoded secret for manual entry
    pub secret_base32: String,
    /// otpauth:// URI for QR code generation
    pub otpauth_uri: String,
    /// QR code as base64-encoded PNG data URI
    pub qr_code_data_uri: String,
}

/// TOTP two-factor authentication service
pub struct TotpService;

impl TotpService {
    /// Generate a new TOTP secret and create a setup row (enabled=false).
    ///
    /// Returns setup data including QR code for the authenticator app.
    /// The user must verify with a valid code before 2FA is enabled.
    pub async fn setup_totp(
        db: &DatabaseConnection,
        user_id: Uuid,
        username: &str,
        issuer: &str,
    ) -> Result<TotpSetupResponse, AppError> {
        // Check if TOTP is already enabled
        if let Some(existing) = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user_id))
            .one(db)
            .await?
        {
            if existing.enabled {
                return Err(AppError::TotpAlreadyEnabled);
            }
            // Delete incomplete setup to start fresh
            let model: user_totp::ActiveModel = existing.into();
            model.delete(db).await?;
        }

        // Generate TOTP secret
        let secret = Secret::generate_secret();
        let secret_base32 = secret.to_encoded().to_string();

        let totp = TOTP::new(
            Algorithm::SHA1,
            6,
            1,
            30,
            secret
                .to_bytes()
                .map_err(|e| AppError::Internal(format!("Failed to convert TOTP secret: {}", e)))?,
            Some(issuer.to_string()),
            username.to_string(),
        )
        .map_err(|e| AppError::Internal(format!("Failed to create TOTP: {}", e)))?;

        let otpauth_uri = totp.get_url();

        // Generate QR code as base64 PNG data URI
        let qr_code_data_uri = totp
            .get_qr_base64()
            .map_err(|e| AppError::Internal(format!("Failed to generate QR code: {}", e)))?;
        let qr_code_data_uri = format!("data:image/png;base64,{}", qr_code_data_uri);

        // Encrypt secret before storage
        let secret_encrypted = CryptoService::encrypt(&secret_base32)?;

        // Create user_totp row with enabled=false
        let now = Utc::now();
        let model = user_totp::ActiveModel {
            id: Set(Uuid::new_v4()),
            user_id: Set(user_id),
            secret_encrypted: Set(secret_encrypted),
            enabled: Set(false),
            backup_codes_encrypted: Set(None),
            backup_codes_remaining: Set(0),
            enabled_at: Set(None),
            last_used_step: NotSet,
            created_at: Set(now),
            updated_at: Set(now),
        };
        model.insert(db).await?;

        Ok(TotpSetupResponse {
            secret_base32,
            otpauth_uri,
            qr_code_data_uri,
        })
    }

    /// Verify a TOTP code during setup and enable 2FA.
    ///
    /// On success, generates backup codes and returns them as plaintext (one-time display).
    ///
    /// The matched time-step is consumed via the persistent anti-replay
    /// guard (RFC 0037) **before** 2FA is enabled, so the setup code cannot
    /// be replayed as an immediate login on any replica. A replayed or
    /// invalid code fails with [`AppError::TotpFailed`] (no replay oracle).
    pub async fn verify_and_enable(
        db: &DatabaseConnection,
        user_id: Uuid,
        code: &str,
    ) -> Result<Vec<String>, AppError> {
        let record = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user_id))
            .one(db)
            .await?
            .ok_or_else(|| AppError::NotFound("TOTP setup not found".to_string()))?;

        if record.enabled {
            return Err(AppError::TotpAlreadyEnabled);
        }

        // Decrypt secret and identify the matching step (RFC 0037: store the
        // step that actually matched, never `now / 30`).
        let secret_base32 = CryptoService::decrypt(&record.secret_encrypted)?;
        let Some(step) = matching_totp_step(&secret_base32, code) else {
            return Err(AppError::TotpFailed);
        };

        // Consume the step before enabling 2FA: the setup code cannot serve
        // as an immediate login on another replica (RFC 0037).
        if !Self::claim_totp_step(db, user_id, step).await? {
            return Err(AppError::TotpFailed);
        }

        // Generate backup codes
        let (plaintext_codes, encrypted_codes) = Self::generate_backup_codes_internal()?;

        // Enable TOTP
        let now = Utc::now();
        let mut model: user_totp::ActiveModel = record.into();
        model.enabled = Set(true);
        model.enabled_at = Set(Some(now));
        model.backup_codes_encrypted = Set(Some(encrypted_codes));
        model.backup_codes_remaining = Set(BACKUP_CODE_COUNT as i32);
        model.updated_at = Set(now);
        model.update(db).await?;

        Ok(plaintext_codes)
    }

    /// Verify a TOTP code for an enabled user (login flow).
    ///
    /// A valid code is accepted **once**: the matched time-step is persisted
    /// in `user_totp.last_used_step` through a conditional compare-and-set
    /// `UPDATE` (RFC 0037), so a replayed code — on this replica or any
    /// other — returns `Ok(false)`, indistinguishable from an invalid code
    /// (no replay oracle). Database errors are propagated as [`AppError`];
    /// the guard never fails open.
    pub async fn verify_code(
        db: &DatabaseConnection,
        user_id: Uuid,
        code: &str,
    ) -> Result<bool, AppError> {
        let record = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user_id))
            .filter(user_totp::Column::Enabled.eq(true))
            .one(db)
            .await?
            .ok_or_else(|| AppError::TotpNotEnabled)?;

        let secret_base32 = CryptoService::decrypt(&record.secret_encrypted)?;
        let Some(step) = matching_totp_step(&secret_base32, code) else {
            return Ok(false);
        };
        // Authoritative cross-replica guard: consume the matched step in DB.
        if !Self::claim_totp_step(db, user_id, step).await? {
            return Ok(false);
        }
        // In-process guard kept as defense-in-depth (arbitrated decision,
        // RFC 0037): practically unreachable after a won CAS and fails
        // closed if it ever disagrees. The DB compare-and-set above remains
        // authoritative.
        if !totp_replay_guard_accept(user_id, step) {
            return Ok(false);
        }
        Ok(true)
    }

    /// Verify and consume a backup code.
    ///
    /// Returns true if the code was valid and consumed.
    pub async fn verify_backup_code(
        db: &DatabaseConnection,
        user_id: Uuid,
        code: &str,
    ) -> Result<bool, AppError> {
        let record = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user_id))
            .filter(user_totp::Column::Enabled.eq(true))
            .one(db)
            .await?
            .ok_or_else(|| AppError::TotpNotEnabled)?;

        let encrypted_codes = match &record.backup_codes_encrypted {
            Some(c) => c.clone(),
            None => return Ok(false),
        };

        // Decrypt the JSON array of hashed codes
        let codes_json = CryptoService::decrypt(&encrypted_codes)?;
        let mut hashed_codes: Vec<String> = serde_json::from_str(&codes_json)
            .map_err(|e| AppError::Internal(format!("Failed to parse backup codes: {}", e)))?;

        // bcrypt::verify is CPU-bound; run the scan off the async runtime.
        let found_index =
            find_matching_backup_index(code.to_string(), hashed_codes.clone()).await?;

        let Some(index) = found_index else {
            return Ok(false);
        };

        // Remove used code
        hashed_codes.remove(index);

        // Re-encrypt and save
        let updated_json = serde_json::to_string(&hashed_codes)
            .map_err(|e| AppError::Internal(format!("Failed to serialize backup codes: {}", e)))?;
        let updated_encrypted = CryptoService::encrypt(&updated_json)?;

        let mut model: user_totp::ActiveModel = record.into();
        model.backup_codes_encrypted = Set(Some(updated_encrypted));
        model.backup_codes_remaining = Set(hashed_codes.len() as i32);
        model.updated_at = Set(Utc::now());
        model.update(db).await?;

        Ok(true)
    }

    /// Disable TOTP for a user (deletes the user_totp row).
    pub async fn disable_totp(db: &DatabaseConnection, user_id: Uuid) -> Result<(), AppError> {
        let record = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user_id))
            .one(db)
            .await?
            .ok_or_else(|| AppError::TotpNotEnabled)?;

        let model: user_totp::ActiveModel = record.into();
        model.delete(db).await?;
        Ok(())
    }

    /// Regenerate backup codes. Returns new plaintext codes.
    ///
    /// Replaces all existing backup codes.
    pub async fn regenerate_backup_codes(
        db: &DatabaseConnection,
        user_id: Uuid,
    ) -> Result<Vec<String>, AppError> {
        let record = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user_id))
            .filter(user_totp::Column::Enabled.eq(true))
            .one(db)
            .await?
            .ok_or_else(|| AppError::TotpNotEnabled)?;

        let (plaintext_codes, encrypted_codes) = Self::generate_backup_codes_internal()?;

        let mut model: user_totp::ActiveModel = record.into();
        model.backup_codes_encrypted = Set(Some(encrypted_codes));
        model.backup_codes_remaining = Set(BACKUP_CODE_COUNT as i32);
        model.updated_at = Set(Utc::now());
        model.update(db).await?;

        Ok(plaintext_codes)
    }

    /// Check if a user has TOTP enabled.
    pub async fn is_totp_enabled(db: &DatabaseConnection, user_id: Uuid) -> Result<bool, AppError> {
        let record = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user_id))
            .filter(user_totp::Column::Enabled.eq(true))
            .one(db)
            .await?;
        Ok(record.is_some())
    }

    /// Get TOTP info for a user (for display in settings/admin).
    pub async fn get_totp_info(
        db: &DatabaseConnection,
        user_id: Uuid,
    ) -> Result<Option<user_totp::Model>, AppError> {
        let record = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user_id))
            .one(db)
            .await?;
        Ok(record)
    }

    // --- Challenge management (for 2FA login flow) ---

    /// Create a 2FA challenge for the login flow.
    ///
    /// Returns the challenge token to send to the client.
    pub async fn create_challenge(
        db: &DatabaseConnection,
        user_id: Uuid,
    ) -> Result<String, AppError> {
        // Summed attempt cap across all challenges in the last 15 minutes (E-3).
        let recent_cutoff = Utc::now() - Duration::minutes(15);
        let recent = totp_challenge::Entity::find()
            .filter(totp_challenge::Column::UserId.eq(user_id))
            .filter(totp_challenge::Column::CreatedAt.gte(recent_cutoff))
            .all(db)
            .await?;

        if totp_window_exceeded(recent.iter().map(|c| c.attempts)) {
            tracing::warn!(%user_id, "TOTP global attempt cap reached");
            return Err(AppError::TotpMaxAttemptsExceeded);
        }

        let challenge_token = Uuid::new_v4().to_string();
        let now = Utc::now();

        let model = totp_challenge::ActiveModel {
            id: Set(Uuid::new_v4()),
            user_id: Set(user_id),
            challenge_token: Set(challenge_token.clone()),
            attempts: Set(0),
            max_attempts: Set(MAX_CHALLENGE_ATTEMPTS),
            expires_at: Set(now + Duration::minutes(CHALLENGE_EXPIRY_MINUTES)),
            created_at: Set(now),
        };
        model.insert(db).await?;

        Ok(challenge_token)
    }

    /// Validate a challenge token. Returns the user_id if valid.
    ///
    /// Checks expiration and max attempts.
    pub async fn validate_challenge(
        db: &DatabaseConnection,
        challenge_token: &str,
    ) -> Result<Uuid, AppError> {
        let challenge = totp_challenge::Entity::find()
            .filter(totp_challenge::Column::ChallengeToken.eq(challenge_token))
            .one(db)
            .await?
            .ok_or(AppError::TotpChallengeInvalid)?;

        // Check expiration
        if Utc::now() > challenge.expires_at {
            // Delete expired challenge
            let model: totp_challenge::ActiveModel = challenge.into();
            model.delete(db).await?;
            return Err(AppError::TotpChallengeInvalid);
        }

        // Check max attempts
        if challenge.attempts >= challenge.max_attempts {
            return Err(AppError::TotpMaxAttemptsExceeded);
        }

        Ok(challenge.user_id)
    }

    /// Increment attempt counter for a challenge.
    pub async fn record_challenge_attempt(
        db: &DatabaseConnection,
        challenge_token: &str,
    ) -> Result<(), AppError> {
        let challenge = totp_challenge::Entity::find()
            .filter(totp_challenge::Column::ChallengeToken.eq(challenge_token))
            .one(db)
            .await?
            .ok_or(AppError::TotpChallengeInvalid)?;

        let mut model: totp_challenge::ActiveModel = challenge.into();
        model.attempts = Set(model.attempts.unwrap() + 1);
        model.update(db).await?;
        Ok(())
    }

    /// Delete a challenge after successful verification.
    pub async fn consume_challenge(
        db: &DatabaseConnection,
        challenge_token: &str,
    ) -> Result<(), AppError> {
        let challenge = totp_challenge::Entity::find()
            .filter(totp_challenge::Column::ChallengeToken.eq(challenge_token))
            .one(db)
            .await?;

        if let Some(c) = challenge {
            let model: totp_challenge::ActiveModel = c.into();
            model.delete(db).await?;
        }
        Ok(())
    }

    /// Cleanup expired challenges (background task).
    pub async fn cleanup_expired_challenges(db: &DatabaseConnection) -> Result<u64, AppError> {
        let now = Utc::now();
        let result = totp_challenge::Entity::delete_many()
            .filter(totp_challenge::Column::ExpiresAt.lt(now))
            .exec(db)
            .await?;
        Ok(result.rows_affected)
    }

    // --- Internal helpers ---

    /// Atomically consume TOTP time-step `step` for `user_id` (RFC 0037).
    ///
    /// Runs a conditional `UPDATE` (compare-and-set) that only succeeds when
    /// `last_used_step` is `NULL` or **strictly** lower than `step`, so the
    /// same step can never be consumed twice — even across replicas.
    /// Returns `Ok(true)` iff this caller won the CAS
    /// (`rows_affected == 1`). A step that does not fit in `i64` is treated
    /// as an invalid code (`Ok(false)`); database errors are propagated
    /// (never fail-open).
    async fn claim_totp_step(
        db: &DatabaseConnection,
        user_id: Uuid,
        step: u64,
    ) -> Result<bool, AppError> {
        // Unreachable for realistic clocks (a unix_time / 30 step fits in
        // i64 for centuries); treat overflow as an invalid code, no panic.
        let Ok(step_i64) = i64::try_from(step) else {
            return Ok(false);
        };

        let result = user_totp::Entity::update_many()
            .col_expr(user_totp::Column::LastUsedStep, Expr::value(step_i64))
            .col_expr(user_totp::Column::UpdatedAt, Expr::value(Utc::now()))
            .filter(user_totp::Column::UserId.eq(user_id))
            .filter(
                Condition::any()
                    .add(user_totp::Column::LastUsedStep.is_null())
                    .add(user_totp::Column::LastUsedStep.lt(step_i64)),
            )
            .exec(db)
            .await?;

        Ok(result.rows_affected == 1)
    }

    /// Verify a TOTP code against a base32-encoded secret.
    ///
    /// Test-only shim over [`matching_totp_step`]: production paths need the
    /// matched step itself to feed the persistent anti-replay guard.
    #[cfg(test)]
    fn verify_totp_code(secret_base32: &str, code: &str) -> Result<bool, AppError> {
        Ok(matching_totp_step(secret_base32, code).is_some())
    }

    /// Generate backup codes: returns (plaintext_codes, encrypted_json_of_hashed_codes).
    fn generate_backup_codes_internal() -> Result<(Vec<String>, String), AppError> {
        use rand::Rng;

        let mut rng = rand::rng();
        let charset: &[u8] = b"ABCDEFGHJKLMNPQRSTUVWXYZ23456789"; // No I/O/0/1 for readability

        let mut plaintext_codes = Vec::with_capacity(BACKUP_CODE_COUNT);
        let mut hashed_codes = Vec::with_capacity(BACKUP_CODE_COUNT);

        for _ in 0..BACKUP_CODE_COUNT {
            let code: String = (0..BACKUP_CODE_LENGTH)
                .map(|_| {
                    let idx = rng.random_range(0..charset.len());
                    charset[idx] as char
                })
                .collect();

            let hashed = bcrypt::hash(&code, 10)
                .map_err(|e| AppError::Internal(format!("Failed to hash backup code: {}", e)))?;

            plaintext_codes.push(code);
            hashed_codes.push(hashed);
        }

        let codes_json = serde_json::to_string(&hashed_codes)
            .map_err(|e| AppError::Internal(format!("Failed to serialize backup codes: {}", e)))?;
        let encrypted = CryptoService::encrypt(&codes_json)?;

        Ok((plaintext_codes, encrypted))
    }
}

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

    #[test]
    fn test_verify_totp_code_with_known_secret() {
        // Generate a secret and create a valid code, then verify
        let secret = Secret::generate_secret();
        let secret_base32 = secret.to_encoded().to_string();

        let totp = TOTP::new(
            Algorithm::SHA1,
            6,
            1,
            30,
            secret.to_bytes().unwrap(),
            None,
            String::new(),
        )
        .unwrap();

        let valid_code = totp.generate_current().unwrap();
        let result = TotpService::verify_totp_code(&secret_base32, &valid_code).unwrap();
        assert!(result, "Valid TOTP code should verify");

        let result = TotpService::verify_totp_code(&secret_base32, "000000").unwrap();
        // May or may not be valid depending on timing, but "000000" is rarely valid
        // So we just check it doesn't error
        let _ = result;
    }

    /// Generates and decrypts backup codes against a fixed test key.
    ///
    /// Backup-code generation calls `CryptoService::encrypt`, which reads
    /// `ENCRYPTION_KEY` from the process environment. `env::set_var` is
    /// process-wide, so concurrent tests (e.g. `crypto_service::test_roundtrip`)
    /// can observe a partial update. `#[serial(encryption_key)]` from
    /// `serial_test` serialises every test sharing this key, removing the race.
    #[test]
    #[serial_test::serial(encryption_key)]
    fn test_generate_backup_codes() {
        // SAFETY: serialised by `#[serial(encryption_key)]`; no other test
        // touches ENCRYPTION_KEY while this block runs.
        unsafe {
            std::env::set_var(
                "ENCRYPTION_KEY",
                "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
            );
        }

        let (plaintext, encrypted) = TotpService::generate_backup_codes_internal().unwrap();

        assert_eq!(plaintext.len(), BACKUP_CODE_COUNT);
        assert!(encrypted.starts_with("enc:"));

        // Each code should be 8 characters
        for code in &plaintext {
            assert_eq!(code.len(), BACKUP_CODE_LENGTH);
        }

        // Decrypt and verify hashes match
        let decrypted_json = CryptoService::decrypt(&encrypted).unwrap();
        let hashed_codes: Vec<String> = serde_json::from_str(&decrypted_json).unwrap();
        assert_eq!(hashed_codes.len(), BACKUP_CODE_COUNT);

        // Verify each plaintext code matches its hash
        for (plain, hashed) in plaintext.iter().zip(hashed_codes.iter()) {
            assert!(bcrypt::verify(plain, hashed).unwrap());
        }
    }

    #[test]
    fn four_challenges_of_four_attempts_are_blocked() {
        assert!(totp_window_exceeded([4, 4, 4, 4].into_iter())); // 16
    }

    #[test]
    fn fourteen_attempts_still_allowed() {
        assert!(!totp_window_exceeded([4, 4, 4, 2].into_iter())); // 14
    }

    #[test]
    fn fifteen_attempts_are_blocked() {
        assert!(totp_window_exceeded([4, 4, 4, 3].into_iter())); // 15
    }

    #[test]
    fn same_step_replay_is_rejected() {
        // Dedicated id so parallel tests cannot collide on Uuid::nil().
        let uid = Uuid::from_u128(0x0707_051E_01A7_u128);
        assert!(totp_replay_guard_accept(uid, 42));
        assert!(!totp_replay_guard_accept(uid, 42));
        assert!(totp_replay_guard_accept(uid, 43));
    }

    #[test]
    fn matching_totp_step_accepts_current_and_previous_windows() {
        let secret = Secret::generate_secret();
        let secret_base32 = secret.to_encoded().to_string();
        let totp = TOTP::new(
            Algorithm::SHA1,
            6,
            1,
            30,
            secret.to_bytes().unwrap(),
            None,
            String::new(),
        )
        .unwrap();

        let now = Utc::now().timestamp().max(0) as u64;
        let current_step = now / 30;
        let current_code = totp.generate(current_step * 30);
        assert_eq!(
            matching_totp_step(&secret_base32, &current_code),
            Some(current_step)
        );

        if current_step > 0 && totp.ttl().unwrap_or(0) > 1 {
            let previous_code = totp.generate((current_step - 1) * 30);
            assert_eq!(
                matching_totp_step(&secret_base32, &previous_code),
                Some(current_step - 1)
            );
        }

        assert_eq!(matching_totp_step(&secret_base32, "1"), None);
    }

    /// Asserts each generated backup code is unique.
    ///
    /// Same rationale as `test_generate_backup_codes`: the test reads
    /// `ENCRYPTION_KEY` indirectly via `CryptoService::encrypt`, so it shares
    /// the `encryption_key` serialisation tag with every other env-mutating
    /// crypto test in this crate.
    #[test]
    #[serial_test::serial(encryption_key)]
    fn test_backup_codes_are_unique() {
        // SAFETY: serialised by `#[serial(encryption_key)]`.
        unsafe {
            std::env::set_var(
                "ENCRYPTION_KEY",
                "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
            );
        }

        let (codes, _) = TotpService::generate_backup_codes_internal().unwrap();
        let mut unique = std::collections::HashSet::new();
        for code in &codes {
            assert!(unique.insert(code.clone()), "Backup codes should be unique");
        }
    }

    #[tokio::test]
    async fn verify_one_backup_matches_known_hash_and_rejects_mismatch() {
        let hashed = bcrypt::hash("ABCD2345", 4).expect("hash backup code");

        assert!(
            verify_one_backup("ABCD2345".to_string(), hashed.clone())
                .await
                .expect("spawn_blocking join"),
            "matching backup code should verify"
        );
        assert!(
            !verify_one_backup("ZZZZ9999".to_string(), hashed)
                .await
                .expect("spawn_blocking join"),
            "mismatched backup code should not verify"
        );
    }

    // ── RFC 0037: persistent anti-replay (`last_used_step`) DB tests ───────
    //
    // These tests require DATABASE_URL (dev PostgreSQL) with the
    // `m20260918_000010_totp_last_used_step` migration applied. They only
    // delete rows they created themselves.

    use crate::models::user;
    use crate::services::user_service::UserService;

    async fn test_db() -> DatabaseConnection {
        dotenvy::dotenv().ok();
        let url =
            std::env::var("DATABASE_URL").expect("DATABASE_URL required for totp_service DB tests");
        sea_orm::Database::connect(&url)
            .await
            .expect("connect test DB")
    }

    /// Sets the fixed test encryption key used by every crypto-touching test.
    fn set_test_encryption_key() {
        // SAFETY: callers are serialised by `#[serial(encryption_key)]`; no
        // other test touches ENCRYPTION_KEY while this block runs.
        unsafe {
            std::env::set_var(
                "ENCRYPTION_KEY",
                "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
            );
        }
    }

    /// Creates a fresh user plus a `user_totp` row and returns the
    /// plaintext base32 secret. Rows belong to this test only.
    async fn create_totp_user(db: &DatabaseConnection, enabled: bool) -> (user::Model, String) {
        let unique = Uuid::new_v4().simple().to_string();
        let user = UserService::create_user(
            db,
            format!("t37-{}", &unique[..12]),
            format!("t37-{}@example.com", &unique[..12]),
            "not-a-real-hash".to_string(),
            false,
        )
        .await
        .expect("create test user");

        let secret = Secret::generate_secret();
        let secret_base32 = secret.to_encoded().to_string();
        let encrypted = CryptoService::encrypt(&secret_base32).expect("encrypt test secret");
        let now = Utc::now();
        user_totp::ActiveModel {
            id: Set(Uuid::new_v4()),
            user_id: Set(user.id),
            secret_encrypted: Set(encrypted),
            enabled: Set(enabled),
            backup_codes_encrypted: Set(None),
            backup_codes_remaining: Set(0),
            enabled_at: Set(enabled.then_some(now)),
            last_used_step: NotSet,
            created_at: Set(now),
            updated_at: Set(now),
        }
        .insert(db)
        .await
        .expect("insert test user_totp row");

        (user, secret_base32)
    }

    /// Deletes only the rows created by the calling test (its `user_totp`
    /// row, then the user itself).
    async fn cleanup_totp_user(db: &DatabaseConnection, user: user::Model) {
        if let Some(row) = user_totp::Entity::find()
            .filter(user_totp::Column::UserId.eq(user.id))
            .one(db)
            .await
            .expect("load test user_totp row for cleanup")
        {
            let model: user_totp::ActiveModel = row.into();
            model.delete(db).await.expect("delete test user_totp row");
        }
        let model: user::ActiveModel = user.into();
        model.delete(db).await.expect("delete test user");
    }

    /// Generates the current 6-digit code for a base32 secret.
    fn current_code(secret_base32: &str) -> String {
        let secret = Secret::Encoded(secret_base32.to_string());
        let totp = TOTP::new(
            Algorithm::SHA1,
            6,
            1,
            30,
            secret.to_bytes().expect("decode test secret"),
            None,
            String::new(),
        )
        .expect("build test TOTP");
        totp.generate_current().expect("generate current code")
    }

    /// RFC 0037 — error case 1: malformed codes never touch the database.
    #[tokio::test]
    #[serial_test::serial(encryption_key)]
    async fn verify_code_invalid_code_writes_nothing() {
        let db = test_db().await;
        set_test_encryption_key();
        let (user, _secret) = create_totp_user(&db, true).await;

        assert!(
            !TotpService::verify_code(&db, user.id, "1")
                .await
                .expect("verify '1'"),
            "'1' must be rejected"
        );
        assert!(
            !TotpService::verify_code(&db, user.id, "12345")
                .await
                .expect("verify 5 digits"),
            "5-digit code must be rejected"
        );
        let row = TotpService::get_totp_info(&db, user.id)
            .await
            .expect("reload row")
            .expect("row exists");
        assert_eq!(
            row.last_used_step, None,
            "invalid codes must not write last_used_step"
        );

        cleanup_totp_user(&db, user).await;
    }

    /// RFC 0037 — happy path + error case 2: first code accepted and persisted,
    /// same code replayed → `Ok(false)` with `last_used_step` unchanged.
    #[tokio::test]
    #[serial_test::serial(encryption_key)]
    async fn verify_code_persists_matched_step_and_rejects_replay() {
        let db = test_db().await;
        set_test_encryption_key();
        let (user, secret_base32) = create_totp_user(&db, true).await;

        let code = current_code(&secret_base32);
        let expected_step =
            matching_totp_step(&secret_base32, &code).expect("generated code must match a step");

        assert!(
            TotpService::verify_code(&db, user.id, &code)
                .await
                .expect("first verify"),
            "first presentation of a valid code must be accepted"
        );
        let row = TotpService::get_totp_info(&db, user.id)
            .await
            .expect("reload row")
            .expect("row exists");
        assert_eq!(
            row.last_used_step,
            Some(i64::try_from(expected_step).expect("step fits in i64")),
            "the matched step (not now/30) must be persisted"
        );

        assert!(
            !TotpService::verify_code(&db, user.id, &code)
                .await
                .expect("replay verify"),
            "replaying the same code must be rejected"
        );
        let row_after = TotpService::get_totp_info(&db, user.id)
            .await
            .expect("reload row")
            .expect("row exists");
        assert_eq!(
            row_after.last_used_step, row.last_used_step,
            "a rejected replay must not change last_used_step"
        );

        cleanup_totp_user(&db, user).await;
    }

    /// RFC 0037 — concurrency: two concurrent claims of the same
    /// `(user, step)` — exactly one wins the CAS. Also covers error case 3:
    /// an older step can no longer be claimed afterwards (strict `<`).
    #[tokio::test]
    #[serial_test::serial(encryption_key)]
    async fn concurrent_claims_exactly_one_wins() {
        let db = test_db().await;
        set_test_encryption_key();
        let (user, _secret) = create_totp_user(&db, true).await;

        let step: u64 = 61_234_567;
        let (a, b) = tokio::join!(
            TotpService::claim_totp_step(&db, user.id, step),
            TotpService::claim_totp_step(&db, user.id, step),
        );
        let (a, b) = (a.expect("claim a"), b.expect("claim b"));
        assert!(
            a ^ b,
            "exactly one concurrent claim must win the CAS (a={a}, b={b})"
        );

        // Same step again → replay, strict `<` rejects it.
        assert!(
            !TotpService::claim_totp_step(&db, user.id, step)
                .await
                .expect("replay claim"),
            "re-claiming the same step must fail"
        );
        // Older window (error case 3) → rejected too.
        assert!(
            !TotpService::claim_totp_step(&db, user.id, step - 1)
                .await
                .expect("older claim"),
            "claiming an older step must fail"
        );
        // Newer window still passes.
        assert!(
            TotpService::claim_totp_step(&db, user.id, step + 1)
                .await
                .expect("newer claim"),
            "claiming a strictly newer step must succeed"
        );

        cleanup_totp_user(&db, user).await;
    }

    /// RFC 0037 — `verify_and_enable` consumes the matched step, so the
    /// setup code cannot be replayed as an immediate login.
    #[tokio::test]
    #[serial_test::serial(encryption_key)]
    async fn verify_and_enable_consumes_step_blocking_login_replay() {
        let db = test_db().await;
        set_test_encryption_key();

        let unique = Uuid::new_v4().simple().to_string();
        let username = format!("t37e-{}", &unique[..12]);
        let user = UserService::create_user(
            &db,
            username.clone(),
            format!("{username}@example.com"),
            "not-a-real-hash".to_string(),
            false,
        )
        .await
        .expect("create test user");

        let setup = TotpService::setup_totp(&db, user.id, &username, "rustwarden-test")
            .await
            .expect("setup TOTP");
        let code = current_code(&setup.secret_base32);

        let backup_codes = TotpService::verify_and_enable(&db, user.id, &code)
            .await
            .expect("verify_and_enable with a valid code");
        assert_eq!(backup_codes.len(), BACKUP_CODE_COUNT);

        let row = TotpService::get_totp_info(&db, user.id)
            .await
            .expect("reload row")
            .expect("row exists");
        assert!(row.enabled, "2FA must be enabled");
        assert!(
            row.last_used_step.is_some(),
            "verify_and_enable must persist the consumed step"
        );

        assert!(
            !TotpService::verify_code(&db, user.id, &code)
                .await
                .expect("immediate login replay"),
            "the setup code must not be reusable as an immediate login"
        );

        cleanup_totp_user(&db, user).await;
    }

    /// RFC 0037 — error case 4: a `DbErr` from the conditional `UPDATE` is
    /// propagated as `Err(_)` — never a fail-open `Ok(true)`.
    #[tokio::test]
    async fn claim_totp_step_propagates_db_error_never_fails_open() {
        use sea_orm::{DatabaseBackend, DbErr, MockDatabase};

        let db = MockDatabase::new(DatabaseBackend::Postgres)
            .append_exec_errors([DbErr::Custom("injected UPDATE failure".to_string())])
            .into_connection();

        let result = TotpService::claim_totp_step(&db, Uuid::new_v4(), 42).await;
        assert!(
            result.is_err(),
            "a DbErr from the CAS UPDATE must propagate, not fail open: {result:?}"
        );
    }
}