rifts 0.3.10

//! # Authentication
//!
//! This module provides a pluggable authentication framework for the
//! protocol. The core abstraction is the [`AuthProvider`] trait, which
//! backends implement to verify client credentials and produce an
//! [`AuthContext`] on success.
//!
//! ## Trait Design
//!
//! [`AuthProvider`] is an async trait so that implementations may
//! perform network calls to external identity providers (OAuth2 servers,
//! LDAP directories, SAML IdPs, etc.) without blocking the runtime.
//!
//! ## Built-in Implementations
//!
//! * [`TokenAuth`] -- An in-memory bearer token store suitable for tests
//!   and single-process deployments. Tokens are registered ahead of time
//!   via [`TokenAuth::register`] and looked up during authentication.
//!
//! * [`AllowAllAuth`] -- A no-op provider that accepts every request as
//!   an anonymous principal. Intended exclusively for unit tests.
//!
//! ## Usage
//!
//! ```rust,no_run
//! use rifts::session::auth::{TokenAuth, AuthContext, AuthProvider, shared};
//! use rifts::session::session::ClientId;
//! use rifts::protocol::hello::AuthMode;
//!
//! let auth = TokenAuth::new();
//! auth.register("secret-token", AuthContext {
//!     client_id: ClientId::new("user-42"),
//!     claims: serde_json::json!({"sub": "user-42"}),
//!     mode: AuthMode::Bearer,
//!     hints: Default::default(),
//! });
//! ```

use std::collections::HashMap;
use std::sync::Arc;

use async_trait::async_trait;
use parking_lot::RwLock;

use crate::error::{Result, RiftError};
use crate::protocol::hello::AuthMode;
use crate::session::session::ClientId;

/// Authenticated principal returned to the server after successful
/// authentication.
///
/// Contains the caller's identity, arbitrary claims from the identity
/// provider, the authentication mode that was used, and optional
/// authorization hints (region, device, risk score).
#[derive(Debug, Clone)]
pub struct AuthContext {
    /// Unique client identifier assigned by the authentication backend.
    /// This is the primary key used to associate the authenticated
    /// principal with a [`Session`](crate::session::session::Session).
    pub client_id: ClientId,

    /// Arbitrary claims returned by the identity provider. The structure
    /// is backend-specific; for JWT-based backends this would typically
    /// contain the decoded JWT payload.
    pub claims: serde_json::Value,

    /// The authentication mode that was used to produce this context
    /// (e.g., `Bearer`, `Anonymous`).
    pub mode: AuthMode,

    /// Free-form authorization hints such as geographic region, device
    /// identifier, or risk score. These are defined in the protocol
    /// specification section 17.2 and may be used for access-control
    /// decisions downstream.
    pub hints: AuthHints,
}

/// Optional authorization context hints carried alongside an
/// [`AuthContext`].
///
/// These fields provide supplementary information that the server may
/// use for fine-grained access control, rate limiting, or audit logging.
/// All fields are optional and have sensible defaults.
#[derive(Debug, Clone, Default)]
pub struct AuthHints {
    /// Geographic or cloud region the client reports (e.g., `"us-east-1"`,
    /// `"eu-west"`). Used for regional routing or compliance checks.
    pub region: Option<String>,

    /// A device identifier such as a hardware serial number or a UUID
    /// generated by the client application. Used for device-level
    /// authorization policies.
    pub device: Option<String>,

    /// An integer risk score in the range 0-100 where higher values
    /// indicate a riskier session (e.g., unfamiliar IP address, expired
    /// certificate). A score of 0 means trusted/low-risk.
    pub risk: Option<u32>,
}

/// Trait implemented by authentication backends.
///
/// Backends are responsible for verifying credentials (bearer tokens,
/// anonymous access, client certificates, etc.) and returning an
/// [`AuthContext`] on success. The trait is `Send + Sync` so that a
/// single provider instance can be shared across async tasks.
///
/// Implementors that need to call out to an external identity provider
/// can do so inside the async methods without blocking the Tokio runtime.
#[async_trait]
pub trait AuthProvider: Send + Sync {
    /// Authenticate the caller given an [`AuthMode`] and an optional
    /// bearer token.
    ///
    /// For [`AuthMode::Anonymous`], the `token` parameter is expected to
    /// be `None`; the backend should synthesize an anonymous principal.
    /// For [`AuthMode::Bearer`], the token must be present and valid.
    ///
    /// # Errors
    ///
    /// Returns [`RiftError::Auth`] if the
    /// credentials are missing, invalid, or if the backend itself
    /// encounters an error.
    async fn authenticate(&self, mode: AuthMode, token: Option<&str>) -> Result<AuthContext>;

    /// Revoke all credentials associated with a given [`ClientId`].
    ///
    /// Called on logout, token expiry, or administrative lockout. After
    /// a successful revocation, subsequent calls to [`authenticate`]
    /// with previously valid tokens must fail.
    ///
    /// # Errors
    ///
    /// Returns an error if the revocation could not be completed (e.g.,
    /// the identity provider is unreachable).
    async fn revoke(&self, client_id: &ClientId) -> Result<()>;

    /// Return the list of [`AuthMode`] values this provider supports.
    ///
    /// Used during the Hello handshake to match the client's offered
    /// auth modes against the server's capabilities. The first mode in
    /// the client's list that also appears in the server's list wins.
    /// Default: `[Bearer, Anonymous]`.
    fn supported_modes(&self) -> Vec<AuthMode> {
        vec![AuthMode::Bearer, AuthMode::Anonymous]
    }
}

/// A trivial in-memory token store.
///
/// Suitable for tests and single-process deployments. Tokens are
/// registered ahead of time via [`register`](Self::register) and
/// looked up during [`authenticate`](AuthProvider::authenticate).
/// Production code should plug in a real identity provider.
pub struct TokenAuth {
    /// Maps bearer tokens to their associated [`AuthContext`].
    tokens: RwLock<HashMap<String, AuthContext>>,

    /// Reverse index mapping each `ClientId` to the set of active
    /// tokens associated with it. Used to efficiently revoke all
    /// tokens for a given client. Uses `HashSet` so duplicate-token
    /// registrations do not bloat the index.
    by_client: RwLock<HashMap<String, std::collections::HashSet<String>>>,
}

impl TokenAuth {
    /// Create an empty token store with no registered tokens.
    pub fn new() -> Self {
        Self {
            tokens: RwLock::new(HashMap::new()),
            by_client: RwLock::new(HashMap::new()),
        }
    }

    /// Register a bearer token and its associated [`AuthContext`].
    ///
    /// The token can be any string (e.g., `"my-secret-token"`). If the
    /// same token is registered multiple times, the later registration
    /// overwrites the earlier one.
    ///
    /// This method also updates the internal reverse index so that
    /// [`revoke`](AuthProvider::revoke) can efficiently remove all
    /// tokens for a given client.
    pub fn register(&self, token: impl Into<String>, ctx: AuthContext) {
        let token = token.into();
        let client_id = ctx.client_id.0.clone();
        // Insert into the forward map first; the reverse index
        // records the token only if the forward map actually
        // changed (avoids pushing duplicate tokens on re-registration).
        let inserted = self.tokens.write().insert(token.clone(), ctx).is_none();
        if inserted {
            self.by_client
                .write()
                .entry(client_id)
                .or_default()
                .insert(token);
        }
    }
}

impl Default for TokenAuth {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl AuthProvider for TokenAuth {
    /// Authenticate a bearer token by looking it up in the in-memory store.
    ///
    /// For [`AuthMode::Anonymous`], a synthetic anonymous [`AuthContext`]
    /// with a ULID-based client identifier is returned without consulting
    /// the token store.
    ///
    /// # Errors
    ///
    /// Returns [`AuthReject::Required`](crate::error::AuthReject::Required)
    /// if the mode is `Bearer` but no token was provided. Returns
    /// [`AuthReject::Invalid`](crate::error::AuthReject::Invalid) if the
    /// token is not found in the store.
    async fn authenticate(&self, mode: AuthMode, token: Option<&str>) -> Result<AuthContext> {
        if mode == AuthMode::Anonymous {
            return Ok(AuthContext {
                client_id: ClientId::new(format!("anon-{}", rand_suffix())),
                claims: serde_json::json!({}),
                mode,
                hints: AuthHints::default(),
            });
        }
        let token = token.ok_or_else(|| RiftError::Auth(crate::error::AuthReject::Required))?;
        self.tokens.read().get(token).cloned().ok_or_else(|| {
            RiftError::Auth(crate::error::AuthReject::Invalid("unknown token".into()))
        })
    }

    /// Revoke all tokens associated with the given [`ClientId`].
    ///
    /// Removes entries from both the forward token map and the reverse
    /// client-to-tokens index. After this call, any previously valid
    /// tokens for this client will fail authentication.
    async fn revoke(&self, client_id: &ClientId) -> Result<()> {
        // Pop the client's token set atomically; if new tokens
        // are registered concurrently, they will live past this
        // revocation, which matches the caller-visible semantics
        // ("revoke all tokens currently associated with the
        // client at the time of the call").
        let tokens = {
            let mut bc = self.by_client.write();
            bc.remove(client_id.as_str()).unwrap_or_default()
        };
        let mut t = self.tokens.write();
        for token in &tokens {
            t.remove(token);
        }
        Ok(())
    }
}

/// Generate a random suffix for anonymous client identifiers using ULID.
fn rand_suffix() -> String {
    ulid::Ulid::new().to_string()
}

/// A no-op authentication provider that accepts every request.
///
/// Always returns an anonymous [`AuthContext`] with client id `"anonymous"`
/// and [`AuthMode::Anonymous`]. Intended exclusively for unit tests where
/// authentication logic is not under test.
pub struct AllowAllAuth;

#[async_trait]
impl AuthProvider for AllowAllAuth {
    /// Always succeeds, returning an anonymous [`AuthContext`].
    async fn authenticate(&self, _mode: AuthMode, _token: Option<&str>) -> Result<AuthContext> {
        Ok(AuthContext {
            client_id: ClientId::new("anonymous"),
            claims: serde_json::json!({}),
            mode: AuthMode::Anonymous,
            hints: AuthHints::default(),
        })
    }

    /// No-op revocation -- always succeeds without side effects.
    async fn revoke(&self, _client_id: &ClientId) -> Result<()> {
        Ok(())
    }
}

/// Wrap any [`AuthProvider`] implementation in an [`Arc`] for shared
/// ownership across async tasks.
///
/// This is a convenience helper; callers can also construct the `Arc`
/// manually if preferred.
pub fn shared<A: AuthProvider + 'static>(a: A) -> Arc<dyn AuthProvider> {
    Arc::new(a)
}

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

    #[tokio::test]
    async fn token_auth_round_trip() {
        let auth = TokenAuth::new();
        auth.register(
            "tok-1",
            AuthContext {
                client_id: ClientId::new("user-1"),
                claims: serde_json::json!({"sub": "user-1"}),
                mode: AuthMode::Bearer,
                hints: AuthHints::default(),
            },
        );
        let ctx = auth
            .authenticate(AuthMode::Bearer, Some("tok-1"))
            .await
            .unwrap();
        assert_eq!(ctx.client_id.as_str(), "user-1");
        assert!(
            auth.authenticate(AuthMode::Bearer, Some("nope"))
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn anonymous_auth_works_without_token() {
        let auth = TokenAuth::new();
        let ctx = auth.authenticate(AuthMode::Anonymous, None).await.unwrap();
        assert_eq!(ctx.mode, AuthMode::Anonymous);
    }

    #[tokio::test]
    async fn revoke_removes_ability_to_authenticate() {
        let auth = TokenAuth::new();
        let client = ClientId::new("user-1");
        auth.register(
            "tok-1",
            AuthContext {
                client_id: client.clone(),
                claims: serde_json::json!({}),
                mode: AuthMode::Bearer,
                hints: AuthHints::default(),
            },
        );
        assert!(
            auth.authenticate(AuthMode::Bearer, Some("tok-1"))
                .await
                .is_ok()
        );
        auth.revoke(&client).await.unwrap();
        assert!(
            auth.authenticate(AuthMode::Bearer, Some("tok-1"))
                .await
                .is_err()
        );
    }
}