arcly-http-identity 0.9.0

//! RFC 6238 TOTP (HMAC-SHA1, 30s step, 6 digits) with enrolment, verification
//! (±1 step clock-skew tolerance), and single-use recovery codes.

use std::time::{SystemTime, UNIX_EPOCH};

use async_trait::async_trait;
use hmac::{Hmac, Mac};
use rand::RngCore;
use sha1::Sha1;
use sha2::{Digest, Sha256};
use subtle::ConstantTimeEq;

use crate::error::{IdentityError, Result};
use crate::mfa::MfaVerifier;

type HmacSha1 = Hmac<Sha1>;

const STEP_SECS: u64 = 30;
const DIGITS: u32 = 6;
const SKEW_STEPS: i64 = 1; // accept the previous/next step for clock drift

/// The result of starting TOTP enrolment: hand `secret_base32` (or the
/// `otpauth_uri` as a QR code) to the user's authenticator app, then confirm
/// with a live code before marking the factor enrolled.
#[derive(Debug, Clone)]
pub struct TotpEnrollment {
    pub secret_base32: String,
    pub otpauth_uri: String,
    /// One-time recovery codes (plaintext, shown once). Only their hashes persist.
    pub recovery_codes: Vec<String>,
}

/// Persistence for a user's TOTP secret and recovery-code hashes.
#[async_trait]
pub trait MfaSecretStore: Send + Sync + 'static {
    async fn set_totp_secret(&self, user_id: &str, secret: &[u8]) -> Result<()>;
    async fn get_totp_secret(&self, user_id: &str) -> Result<Option<Vec<u8>>>;
    async fn clear_totp(&self, user_id: &str) -> Result<()>;
    /// Store recovery-code SHA-256 hashes (replaces any existing set).
    async fn set_recovery_hashes(&self, user_id: &str, hashes: Vec<[u8; 32]>) -> Result<()>;
    /// Atomically consume a matching recovery-code hash; `true` if one was found.
    async fn take_recovery_hash(&self, user_id: &str, hash: &[u8; 32]) -> Result<bool>;
    async fn recovery_count(&self, user_id: &str) -> Result<u32>;
}

/// TOTP factor + recovery codes, verifying both `"totp"` and `"recovery"`.
pub struct TotpService {
    store: std::sync::Arc<dyn MfaSecretStore>,
    issuer: String,
}

impl TotpService {
    pub fn new(store: std::sync::Arc<dyn MfaSecretStore>, issuer: impl Into<String>) -> Self {
        Self {
            store,
            issuer: issuer.into(),
        }
    }

    /// Begin enrolment: generate a fresh secret + recovery codes and persist the
    /// secret and recovery hashes. The factor is not "enrolled" on the identity
    /// until [`confirm`](Self::confirm) succeeds.
    pub async fn begin_enrollment(
        &self,
        user_id: &str,
        account_label: &str,
        recovery_count: usize,
    ) -> Result<TotpEnrollment> {
        let mut secret = [0u8; 20]; // 160-bit, RFC 4226 recommended
        rand::thread_rng().fill_bytes(&mut secret);
        let secret_base32 = base32::encode(base32::Alphabet::RFC4648 { padding: false }, &secret);

        let otpauth_uri = format!(
            "otpauth://totp/{issuer}:{label}?secret={secret}&issuer={issuer}&digits={digits}&period={period}",
            issuer = urlencode(&self.issuer),
            label = urlencode(account_label),
            secret = secret_base32,
            digits = DIGITS,
            period = STEP_SECS,
        );

        let (codes, hashes) = generate_recovery_codes(recovery_count);
        self.store.set_totp_secret(user_id, &secret).await?;
        self.store.set_recovery_hashes(user_id, hashes).await?;

        Ok(TotpEnrollment {
            secret_base32,
            otpauth_uri,
            recovery_codes: codes,
        })
    }

    /// Confirm enrolment by checking a live code. Returns `Ok(())` when valid so
    /// the caller can flip `MfaState::totp_enrolled`.
    pub async fn confirm(&self, user_id: &str, code: &str) -> Result<()> {
        if self.verify_totp(user_id, code).await? {
            Ok(())
        } else {
            Err(IdentityError::MfaFailed)
        }
    }

    /// Remove the TOTP factor entirely.
    pub async fn disable(&self, user_id: &str) -> Result<()> {
        self.store.clear_totp(user_id).await
    }

    async fn verify_totp(&self, user_id: &str, code: &str) -> Result<bool> {
        let Some(secret) = self.store.get_totp_secret(user_id).await? else {
            return Ok(false);
        };
        let now = unix_now();
        let counter = now / STEP_SECS;
        for delta in -SKEW_STEPS..=SKEW_STEPS {
            let c = (counter as i64 + delta) as u64;
            let expected = hotp(&secret, c);
            // Constant-time compare of the zero-padded decimal strings.
            let candidate = format!("{expected:0width$}", width = DIGITS as usize);
            if candidate.as_bytes().ct_eq(code.trim().as_bytes()).into() {
                return Ok(true);
            }
        }
        Ok(false)
    }

    async fn verify_recovery(&self, user_id: &str, code: &str) -> Result<bool> {
        let hash = sha256(code.trim().as_bytes());
        self.store.take_recovery_hash(user_id, &hash).await
    }
}

#[async_trait]
impl MfaVerifier for TotpService {
    async fn verify(&self, user_id: &str, method: &str, code: &str) -> Result<bool> {
        match method {
            "totp" => self.verify_totp(user_id, code).await,
            "recovery" => self.verify_recovery(user_id, code).await,
            _ => Ok(false),
        }
    }
}

/// RFC 4226 HOTP truncation over an 8-byte big-endian counter.
fn hotp(secret: &[u8], counter: u64) -> u32 {
    let mut mac = HmacSha1::new_from_slice(secret).expect("HMAC accepts any key length");
    mac.update(&counter.to_be_bytes());
    let digest = mac.finalize().into_bytes();
    let offset = (digest[digest.len() - 1] & 0x0f) as usize;
    let bin = ((u32::from(digest[offset]) & 0x7f) << 24)
        | (u32::from(digest[offset + 1]) << 16)
        | (u32::from(digest[offset + 2]) << 8)
        | u32::from(digest[offset + 3]);
    bin % 10u32.pow(DIGITS)
}

fn generate_recovery_codes(count: usize) -> (Vec<String>, Vec<[u8; 32]>) {
    let mut codes = Vec::with_capacity(count);
    let mut hashes = Vec::with_capacity(count);
    for _ in 0..count {
        let mut b = [0u8; 8];
        rand::thread_rng().fill_bytes(&mut b);
        // 16-hex-char grouped code, e.g. "3f9a-1c04-7b22-e8d1"
        let raw = b.iter().map(|x| format!("{x:02x}")).collect::<String>();
        let grouped = raw
            .as_bytes()
            .chunks(4)
            .map(|c| std::str::from_utf8(c).unwrap())
            .collect::<Vec<_>>()
            .join("-");
        hashes.push(sha256(grouped.as_bytes()));
        codes.push(grouped);
    }
    (codes, hashes)
}

fn sha256(bytes: &[u8]) -> [u8; 32] {
    let mut h = Sha256::new();
    h.update(bytes);
    h.finalize().into()
}

fn unix_now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

/// Minimal percent-encoding for the otpauth label/issuer (space + a few reserved
/// chars); authenticator apps are lenient but colons and spaces must be escaped.
fn urlencode(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for b in s.bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                out.push(b as char)
            }
            _ => out.push_str(&format!("%{b:02X}")),
        }
    }
    out
}

// ─── In-memory MfaSecretStore (tests / dev) ─────────────────────────────────

/// In-memory [`MfaSecretStore`] — tests / local dev only.
#[derive(Default)]
pub struct MemoryMfaSecretStore {
    secrets: std::sync::RwLock<std::collections::HashMap<String, Vec<u8>>>,
    recovery: std::sync::RwLock<std::collections::HashMap<String, Vec<[u8; 32]>>>,
}

impl MemoryMfaSecretStore {
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl MfaSecretStore for MemoryMfaSecretStore {
    async fn set_totp_secret(&self, user_id: &str, secret: &[u8]) -> Result<()> {
        self.secrets
            .write()
            .unwrap()
            .insert(user_id.to_owned(), secret.to_vec());
        Ok(())
    }
    async fn get_totp_secret(&self, user_id: &str) -> Result<Option<Vec<u8>>> {
        Ok(self.secrets.read().unwrap().get(user_id).cloned())
    }
    async fn clear_totp(&self, user_id: &str) -> Result<()> {
        self.secrets.write().unwrap().remove(user_id);
        Ok(())
    }
    async fn set_recovery_hashes(&self, user_id: &str, hashes: Vec<[u8; 32]>) -> Result<()> {
        self.recovery
            .write()
            .unwrap()
            .insert(user_id.to_owned(), hashes);
        Ok(())
    }
    async fn take_recovery_hash(&self, user_id: &str, hash: &[u8; 32]) -> Result<bool> {
        let mut g = self.recovery.write().unwrap();
        let Some(v) = g.get_mut(user_id) else {
            return Ok(false);
        };
        if let Some(pos) = v.iter().position(|h| h == hash) {
            v.remove(pos);
            Ok(true)
        } else {
            Ok(false)
        }
    }
    async fn recovery_count(&self, user_id: &str) -> Result<u32> {
        Ok(self
            .recovery
            .read()
            .unwrap()
            .get(user_id)
            .map(|v| v.len() as u32)
            .unwrap_or(0))
    }
}

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

    #[test]
    fn hotp_matches_rfc4226_vectors() {
        // RFC 4226 Appendix D test secret "12345678901234567890".
        let secret = b"12345678901234567890";
        assert_eq!(hotp(secret, 0), 755224);
        assert_eq!(hotp(secret, 1), 287082);
        assert_eq!(hotp(secret, 2), 359152);
    }
}