arcly_http_identity/

session.rs

//! Server-side token state: refresh-token rotation, pending MFA challenges, and
//! the per-user device/session index used for "log out everywhere" and
//! back-channel revocation (Phase 4).

use std::collections::HashMap;
use std::sync::RwLock;
use std::time::{SystemTime, UNIX_EPOCH};

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

use crate::error::Result;

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

// ─── Refresh token store (rotation) ─────────────────────────────────────────

/// Persists the live `jti` of each refresh token so rotation can atomically
/// invalidate the old one. Backed by Redis with per-key TTL in production.
#[async_trait]
pub trait RefreshStore: Send + Sync + 'static {
    /// Record that `jti` is a valid refresh token for `user_id`.
    async fn save(&self, jti: &str, user_id: &str, ttl_secs: u64) -> Result<()>;
    /// Atomically consume `jti`: return its `user_id` and delete it. `None` if
    /// unknown / already used / expired (rotation replay → reject + revoke).
    async fn take(&self, jti: &str) -> Result<Option<String>>;
    /// Revoke a specific `jti` (logout).
    async fn revoke(&self, jti: &str) -> Result<()>;
    /// Revoke every refresh token for a user ("log out everywhere").
    async fn revoke_all(&self, user_id: &str) -> Result<()>;
}

// ─── Pending MFA challenge store ────────────────────────────────────────────

/// After a correct password on an MFA-enrolled account, the engine parks a
/// short-lived challenge and returns its id; the second-factor step redeems it.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct PendingChallenge {
    pub user_id: String,
    pub expires_at: u64,
}

#[async_trait]
pub trait ChallengeStore: Send + Sync + 'static {
    async fn put(&self, challenge_id: &str, user_id: &str, ttl_secs: u64) -> Result<()>;
    async fn take(&self, challenge_id: &str) -> Result<Option<PendingChallenge>>;
}

// ─── Session / device index ─────────────────────────────────────────────────

/// One authenticated session/device row surfaced to the user for review/revoke.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionRecord {
    pub session_id: String,
    pub user_id: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub device: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub ip: Option<String>,
    pub created_at: u64,
    pub last_seen: u64,
}

/// Tracks active sessions per user for management + back-channel logout.
#[async_trait]
pub trait SessionIndex: Send + Sync + 'static {
    async fn register(&self, record: SessionRecord) -> Result<()>;
    async fn list(&self, user_id: &str) -> Result<Vec<SessionRecord>>;
    async fn revoke(&self, session_id: &str) -> Result<()>;
    async fn revoke_all(&self, user_id: &str) -> Result<()>;
}

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

#[derive(Default)]
pub struct MemoryRefreshStore {
    // jti -> (user_id, expires_at)
    inner: RwLock<HashMap<String, (String, u64)>>,
}
impl MemoryRefreshStore {
    pub fn new() -> Self {
        Self::default()
    }
}
#[async_trait]
impl RefreshStore for MemoryRefreshStore {
    async fn save(&self, jti: &str, user_id: &str, ttl_secs: u64) -> Result<()> {
        self.inner
            .write()
            .unwrap()
            .insert(jti.to_owned(), (user_id.to_owned(), unix_now() + ttl_secs));
        Ok(())
    }
    async fn take(&self, jti: &str) -> Result<Option<String>> {
        let removed = self.inner.write().unwrap().remove(jti);
        Ok(removed.and_then(|(uid, exp)| (unix_now() < exp).then_some(uid)))
    }
    async fn revoke(&self, jti: &str) -> Result<()> {
        self.inner.write().unwrap().remove(jti);
        Ok(())
    }
    async fn revoke_all(&self, user_id: &str) -> Result<()> {
        self.inner
            .write()
            .unwrap()
            .retain(|_, (uid, _)| uid != user_id);
        Ok(())
    }
}

#[derive(Default)]
pub struct MemoryChallengeStore {
    inner: RwLock<HashMap<String, PendingChallenge>>,
}
impl MemoryChallengeStore {
    pub fn new() -> Self {
        Self::default()
    }
}
#[async_trait]
impl ChallengeStore for MemoryChallengeStore {
    async fn put(&self, challenge_id: &str, user_id: &str, ttl_secs: u64) -> Result<()> {
        self.inner.write().unwrap().insert(
            challenge_id.to_owned(),
            PendingChallenge {
                user_id: user_id.to_owned(),
                expires_at: unix_now() + ttl_secs,
            },
        );
        Ok(())
    }
    async fn take(&self, challenge_id: &str) -> Result<Option<PendingChallenge>> {
        let taken = self.inner.write().unwrap().remove(challenge_id);
        Ok(taken.filter(|c| unix_now() < c.expires_at))
    }
}

#[derive(Default)]
pub struct MemorySessionIndex {
    inner: RwLock<HashMap<String, SessionRecord>>,
}
impl MemorySessionIndex {
    pub fn new() -> Self {
        Self::default()
    }
}
#[async_trait]
impl SessionIndex for MemorySessionIndex {
    async fn register(&self, record: SessionRecord) -> Result<()> {
        self.inner
            .write()
            .unwrap()
            .insert(record.session_id.clone(), record);
        Ok(())
    }
    async fn list(&self, user_id: &str) -> Result<Vec<SessionRecord>> {
        Ok(self
            .inner
            .read()
            .unwrap()
            .values()
            .filter(|r| r.user_id == user_id)
            .cloned()
            .collect())
    }
    async fn revoke(&self, session_id: &str) -> Result<()> {
        self.inner.write().unwrap().remove(session_id);
        Ok(())
    }
    async fn revoke_all(&self, user_id: &str) -> Result<()> {
        self.inner
            .write()
            .unwrap()
            .retain(|_, r| r.user_id != user_id);
        Ok(())
    }
}