auth.rs

//! Token-based authentication with argon2 hashing and per-key rate limiting.

use std::fmt;
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};

use argon2::{
    password_hash::{rand_core::OsRng, SaltString},
    Argon2, PasswordHash, PasswordHasher, PasswordVerifier,
};
use dashmap::DashMap;
use sha2::{Digest, Sha256};
use tonic::{Request, Status};
use tracing::{info, warn};

// ---------------------------------------------------------------------------
// AuthError
// ---------------------------------------------------------------------------

/// Errors that can occur during token hashing or hash parsing.
#[derive(Debug)]
pub enum AuthError {
    HashingFailed(String),
    // Reserved for future use when parsing stored hashes from config.
    #[allow(dead_code)]
    InvalidHash(String),
}

impl fmt::Display for AuthError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            AuthError::HashingFailed(msg) => write!(f, "token hashing failed: {msg}"),
            AuthError::InvalidHash(msg) => write!(f, "invalid token hash: {msg}"),
        }
    }
}

impl std::error::Error for AuthError {}

// ---------------------------------------------------------------------------
// ClientIdentity
// ---------------------------------------------------------------------------

/// Identifies an authenticated client and the method used.
#[derive(Debug, Clone)]
pub struct ClientIdentity {
    /// Human-readable client identifier (from cert CN or a generated id).
    pub client_id: String,
    /// How the client was authenticated.
    // Used in tests and will be consumed by future audit logging.
    #[allow(dead_code)]
    pub auth_method: AuthMethod,
}

/// The authentication method that was used for a request.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AuthMethod {
    /// Bearer token validated against an argon2 hash.
    Token,
    /// Mutual TLS — client presented a valid certificate.
    MutualTls { cert_fingerprint: String },
    /// Authentication is disabled.
    None,
}

// ---------------------------------------------------------------------------
// AuthInterceptor
// ---------------------------------------------------------------------------

/// gRPC request authenticator.
///
/// Supports two authentication methods that can be used simultaneously:
///
/// 1. **Bearer token** — `Authorization: Bearer <token>` header validated
///    against an argon2 hash.
/// 2. **Mutual TLS** — the client presents a certificate during the TLS
///    handshake; the server verifies it against the CA.
///
/// When `require_mtls` is `true`, only mTLS is accepted (token auth is
/// rejected).  When both `token_hash` and mTLS are configured, either
/// method is sufficient.  When neither is configured, all requests are
/// accepted.
///
/// # Token verification performance
///
/// Argon2 is intentionally slow (~100 ms per call). The bearer token is
/// static for the server's lifetime, so once a token is successfully
/// verified its SHA-256 digest is cached in `verified_token_digest`.
/// Subsequent requests skip the KDF entirely and compare SHA-256 digests
/// in constant time instead (< 1 µs). Failed verification attempts are
/// never cached, so brute-force attempts still incur the full Argon2 cost.
///
/// A digest is cached rather than the plaintext to avoid retaining the raw
/// token in memory beyond the initial verification call, keeping the blast
/// radius of any memory disclosure consistent with the pre-cache state
/// (only the Argon2 hash was retained before).
pub struct AuthInterceptor {
    token_hash: Option<String>,
    require_mtls: bool,
    /// SHA-256 digest of the first successfully verified bearer token.
    ///
    /// Populated on the first successful `verify_token` call. Subsequent
    /// calls short-circuit by comparing the SHA-256 digest of the incoming
    /// token against this cached digest in constant time, bypassing Argon2.
    ///
    /// # Concurrency note
    ///
    /// Multiple concurrent first requests may each run Argon2 before one
    /// wins the `OnceLock::set` race.  This is intentional: the race is
    /// benign because all concurrent winners verified the same correct token
    /// and would store the same digest.  Once the lock is set, all
    /// subsequent calls take the fast path.  The "at most once" guarantee
    /// holds for steady-state traffic; only a narrow startup burst can
    /// observe multiple Argon2 runs, and only if those requests happen to
    /// arrive before the lock is set.
    verified_token_digest: OnceLock<[u8; 32]>,
}

impl AuthInterceptor {
    /// Create a new interceptor.
    ///
    /// * `token_hash` — argon2 PHC hash of the bearer token, or `None` to
    ///   skip token-based auth.
    /// * `require_mtls` — when `true`, only requests authenticated via mTLS
    ///   are accepted.
    pub fn new(token_hash: Option<String>, require_mtls: bool) -> Self {
        Self {
            token_hash,
            require_mtls,
            verified_token_digest: OnceLock::new(),
        }
    }

    /// Whether this interceptor has any authentication requirement.
    pub fn auth_enabled(&self) -> bool {
        self.token_hash.is_some() || self.require_mtls
    }

    /// Verify a raw bearer token against the stored argon2 hash.
    ///
    /// Returns `true` when:
    /// - token auth is not configured (`token_hash` is `None`), or
    /// - the token matches the stored hash.
    ///
    /// On the first successful verification the SHA-256 digest of the token
    /// is cached so that subsequent calls skip the expensive Argon2 KDF
    /// entirely.
    pub fn verify_token(&self, token: &str) -> bool {
        match &self.token_hash {
            None => true,
            Some(hash) => {
                let incoming_digest = sha256(token.as_bytes());

                // Fast path: a token was already verified — compare digests
                // in constant time without invoking the Argon2 KDF.
                if let Some(cached) = self.verified_token_digest.get() {
                    return constant_time_eq(&incoming_digest, cached);
                }

                // Slow path: first verification — run Argon2.
                let parsed = match PasswordHash::new(hash) {
                    Ok(h) => h,
                    Err(e) => {
                        warn!(error = %e, "stored token hash is malformed");
                        return false;
                    }
                };
                let ok = Argon2::default()
                    .verify_password(token.as_bytes(), &parsed)
                    .is_ok();

                if ok {
                    // Cache the digest of the verified token. If another
                    // thread raced us here, the winner's value is canonical
                    // — both verified the same correct token so the stored
                    // digest is identical either way.
                    let _ = self.verified_token_digest.set(incoming_digest);
                }

                ok
            }
        }
    }

    /// Authenticate a gRPC request.
    ///
    /// Checks for a valid mTLS client certificate first (via the
    /// `x-client-cert-cn` and `x-client-cert-fingerprint` metadata headers
    /// injected by the TLS layer), then falls back to bearer token.
    ///
    /// Returns a [`ClientIdentity`] on success.
    pub fn check_request<T>(&self, request: &Request<T>) -> Result<ClientIdentity, Status> {
        // No auth configured — accept everyone.
        if !self.auth_enabled() {
            return Ok(ClientIdentity {
                client_id: "anonymous".into(),
                auth_method: AuthMethod::None,
            });
        }

        let remote_addr = request
            .remote_addr()
            .map_or_else(|| "<unknown>".to_owned(), |a| a.to_string());

        // 1. Try mTLS (peer certs are exposed via tonic's TLS info).
        if let Some(identity) = self.check_mtls_metadata(request) {
            info!(
                client_ip = %remote_addr,
                client_id = %identity.client_id,
                "mTLS auth success",
            );
            return Ok(identity);
        }

        // 2. If mTLS is required and wasn't present, reject.
        if self.require_mtls {
            warn!(
                client_ip = %remote_addr,
                "auth failure: mTLS required but no client certificate",
            );
            return Err(Status::unauthenticated(
                "Authentication failed. This server requires a mutual TLS client certificate. \
                 Please configure your client certificate and try again.",
            ));
        }

        // 3. Fall back to bearer token.
        if self.token_hash.is_some() {
            let token = Self::extract_bearer(request).ok_or_else(|| {
                warn!(
                    client_ip = %remote_addr,
                    "auth failure: missing or malformed authorization header",
                );
                Status::unauthenticated(
                    "Authentication required. Please provide a valid Bearer token in the Authorization header."
                )
            })?;

            if self.verify_token(token) {
                info!(
                    client_ip = %remote_addr,
                    "token auth success",
                );
                return Ok(ClientIdentity {
                    client_id: format!("token-{}", &remote_addr),
                    auth_method: AuthMethod::Token,
                });
            }

            warn!(
                client_ip = %remote_addr,
                "auth failure: invalid token",
            );
            return Err(Status::unauthenticated(
                "Authentication failed. The provided token is invalid. Please check your token and try again."
            ));
        }

        // No valid auth method succeeded.
        Err(Status::unauthenticated(
            "Authentication required. Please provide a valid Bearer token or client certificate.",
        ))
    }

    /// Check for mTLS identity carried through gRPC metadata.
    ///
    /// When mTLS is configured at the transport level, the server's TLS
    /// acceptor verifies the client certificate against the CA.  We expose
    /// the verified client CN and fingerprint as metadata so the
    /// application layer can use them for identity.
    ///
    /// Metadata headers (set by the mTLS-aware server bootstrap):
    /// - `x-client-cert-cn`          — Common Name from the client cert
    /// - `x-client-cert-fingerprint` — SHA-256 fingerprint
    fn check_mtls_metadata<T>(&self, request: &Request<T>) -> Option<ClientIdentity> {
        let cn = request.metadata().get("x-client-cert-cn")?.to_str().ok()?;
        let fingerprint = request
            .metadata()
            .get("x-client-cert-fingerprint")?
            .to_str()
            .ok()?;

        Some(ClientIdentity {
            client_id: cn.to_owned(),
            auth_method: AuthMethod::MutualTls {
                cert_fingerprint: fingerprint.to_owned(),
            },
        })
    }

    /// Hash a raw token for storage in the runner config.
    ///
    /// Uses argon2id with a random salt. The returned string is a PHC-format
    /// hash suitable for storage in `AuthConfig::token_hash`.
    pub fn hash_token(token: &str) -> Result<String, AuthError> {
        let salt = SaltString::generate(&mut OsRng);
        let hash = Argon2::default()
            .hash_password(token.as_bytes(), &salt)
            .map_err(|e| AuthError::HashingFailed(e.to_string()))?;
        Ok(hash.to_string())
    }

    // ------------------------------------------------------------------
    // Private helpers
    // ------------------------------------------------------------------

    /// Extract the bearer token from `authorization` metadata.
    ///
    /// Returns `None` when the header is absent, non-ASCII, or does not start
    /// with `"Bearer "`.
    fn extract_bearer<T>(request: &Request<T>) -> Option<&str> {
        let value = request.metadata().get("authorization")?.to_str().ok()?;
        value.strip_prefix("Bearer ")
    }

    /// Wrap `self` in an `Arc`.
    #[allow(dead_code)]
    pub fn into_arc(self) -> Arc<Self> {
        Arc::new(self)
    }
}

/// Compute the SHA-256 digest of `data`.
fn sha256(data: &[u8]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(data);
    hasher.finalize().into()
}

/// Compare two equal-length byte slices in constant time.
///
/// Returns `true` only when `a` and `b` are identical in content.
///
/// # Constant-time guarantee
///
/// This function assumes both slices have the same length (callers must check
/// lengths independently if needed). Given equal-length inputs it visits
/// every byte pair regardless of where a mismatch occurs, so the execution
/// time does not reveal the position of the first differing byte.
///
/// Length comparison itself is **not** constant-time — if you need to compare
/// values whose lengths must also be kept secret, pad them to a fixed size
/// before calling this function.
fn constant_time_eq(a: &[u8; 32], b: &[u8; 32]) -> bool {
    // XOR each pair of bytes and OR the results together.  Any non-zero
    // accumulator means at least one byte differed.
    let diff = a
        .iter()
        .zip(b.iter())
        .fold(0u8, |acc, (x, y)| acc | (x ^ y));
    diff == 0
}

// ---------------------------------------------------------------------------
// RateLimiter
// ---------------------------------------------------------------------------

/// Per-key fixed-window rate limiter backed by a lock-free [`DashMap`].
///
/// Each key (typically a client IP address) is tracked independently.
/// Within each window the counter increments up to `max_requests`; once the
/// window expires the counter resets.
pub struct RateLimiter {
    limits: DashMap<String, RateState>,
    max_requests: u32,
    window: Duration,
}

struct RateState {
    count: u32,
    window_start: Instant,
}

impl RateLimiter {
    /// Create a new rate limiter.
    ///
    /// # Parameters
    /// - `max_requests`: maximum requests allowed per `window`.
    /// - `window`: length of each fixed window.
    pub fn new(max_requests: u32, window: Duration) -> Self {
        Self {
            limits: DashMap::new(),
            max_requests,
            window,
        }
    }

    /// Create a rate limiter with the default connection rate:
    /// 10 requests per 60 seconds per key.
    #[allow(dead_code)]
    pub fn default_connection_rate() -> Self {
        Self::new(10, Duration::from_secs(60))
    }

    /// Check whether a request from `key` should be allowed.
    ///
    /// Returns:
    /// - `Ok(())` when the request is within the rate limit.
    /// - `Err(retry_after)` when the limit has been exceeded; `retry_after`
    ///   is the remaining time until the current window resets.
    pub fn check(&self, key: &str) -> Result<(), Duration> {
        let now = Instant::now();

        let mut entry = self.limits.entry(key.to_owned()).or_insert(RateState {
            count: 0,
            window_start: now,
        });

        let state = entry.value_mut();

        // Reset counter if the window has expired.
        if now.duration_since(state.window_start) >= self.window {
            state.count = 0;
            state.window_start = now;
        }

        if state.count >= self.max_requests {
            let elapsed = now.duration_since(state.window_start);
            let retry_after = self.window.saturating_sub(elapsed);
            return Err(retry_after);
        }

        state.count += 1;
        Ok(())
    }

    /// Remove entries whose window has expired.
    ///
    /// Call this periodically (e.g. every few minutes) to prevent unbounded
    /// memory growth from keys that are no longer active.
    #[allow(dead_code)]
    pub fn cleanup(&self) {
        let now = Instant::now();
        self.limits
            .retain(|_, state| now.duration_since(state.window_start) < self.window);
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // -- AuthInterceptor ------------------------------------------------------

    #[test]
    fn auth_disabled_accepts_all() {
        let interceptor = AuthInterceptor::new(None, false);
        assert!(interceptor.verify_token("anything"));
        assert!(interceptor.verify_token(""));
    }

    #[test]
    fn hash_and_verify_roundtrip() {
        let token = "super-secret-token-42";
        let hash = AuthInterceptor::hash_token(token).expect("hash_token failed");
        let interceptor = AuthInterceptor::new(Some(hash), false);
        assert!(interceptor.verify_token(token));
        assert!(!interceptor.verify_token("wrong-token"));
    }

    #[test]
    fn verify_rejects_wrong_token() {
        let hash = AuthInterceptor::hash_token("correct").unwrap();
        let interceptor = AuthInterceptor::new(Some(hash), false);
        assert!(!interceptor.verify_token("incorrect"));
    }

    #[test]
    fn verify_uses_cache_on_second_correct_call() {
        let token = "cached-token";
        let hash = AuthInterceptor::hash_token(token).unwrap();
        let interceptor = AuthInterceptor::new(Some(hash), false);

        // First call populates the cache via Argon2.
        assert!(interceptor.verify_token(token));
        assert!(interceptor.verified_token_digest.get().is_some());

        // Second call must hit the cache (token still valid).
        assert!(interceptor.verify_token(token));
    }

    #[test]
    fn verify_wrong_token_does_not_populate_cache() {
        let hash = AuthInterceptor::hash_token("real-token").unwrap();
        let interceptor = AuthInterceptor::new(Some(hash), false);

        // Wrong token — should not populate cache.
        assert!(!interceptor.verify_token("wrong-token"));
        assert!(interceptor.verified_token_digest.get().is_none());

        // Correct token still works after a failed attempt.
        assert!(interceptor.verify_token("real-token"));
    }

    #[test]
    fn verify_cached_rejects_different_token() {
        let token = "my-token";
        let hash = AuthInterceptor::hash_token(token).unwrap();
        let interceptor = AuthInterceptor::new(Some(hash), false);

        // Populate cache.
        assert!(interceptor.verify_token(token));

        // A different token must be rejected even after cache is warm.
        assert!(!interceptor.verify_token("other-token"));
    }

    #[test]
    fn check_request_passes_when_auth_disabled() {
        let interceptor = AuthInterceptor::new(None, false);
        let req = Request::new(());
        let identity = interceptor.check_request(&req).unwrap();
        assert_eq!(identity.auth_method, AuthMethod::None);
    }

    #[test]
    fn check_request_missing_header_returns_unauthenticated() {
        let hash = AuthInterceptor::hash_token("tok").unwrap();
        let interceptor = AuthInterceptor::new(Some(hash), false);
        let req = Request::new(());
        let err = interceptor.check_request(&req).unwrap_err();
        assert_eq!(err.code(), tonic::Code::Unauthenticated);
        assert!(err.message().contains("Authentication required"));
    }

    #[test]
    fn check_request_invalid_token_returns_unauthenticated() {
        let hash = AuthInterceptor::hash_token("correct").unwrap();
        let interceptor = AuthInterceptor::new(Some(hash), false);
        let mut req = Request::new(());
        req.metadata_mut()
            .insert("authorization", "Bearer wrong".parse().unwrap());
        let err = interceptor.check_request(&req).unwrap_err();
        assert_eq!(err.code(), tonic::Code::Unauthenticated);
        assert!(err.message().contains("Authentication failed"));
    }

    #[test]
    fn check_request_valid_token_succeeds() {
        let token = "my-token";
        let hash = AuthInterceptor::hash_token(token).unwrap();
        let interceptor = AuthInterceptor::new(Some(hash), false);
        let mut req = Request::new(());
        req.metadata_mut()
            .insert("authorization", format!("Bearer {token}").parse().unwrap());
        let identity = interceptor.check_request(&req).unwrap();
        assert_eq!(identity.auth_method, AuthMethod::Token);
    }

    #[test]
    fn check_request_mtls_metadata_succeeds() {
        let interceptor = AuthInterceptor::new(None, true);
        let mut req = Request::new(());
        req.metadata_mut()
            .insert("x-client-cert-cn", "test-client".parse().unwrap());
        req.metadata_mut()
            .insert("x-client-cert-fingerprint", "AA:BB:CC".parse().unwrap());
        let identity = interceptor.check_request(&req).unwrap();
        assert_eq!(identity.client_id, "test-client");
        assert_eq!(
            identity.auth_method,
            AuthMethod::MutualTls {
                cert_fingerprint: "AA:BB:CC".into()
            }
        );
    }

    #[test]
    fn check_request_require_mtls_rejects_token() {
        let hash = AuthInterceptor::hash_token("tok").unwrap();
        let interceptor = AuthInterceptor::new(Some(hash), true);
        let mut req = Request::new(());
        req.metadata_mut()
            .insert("authorization", "Bearer tok".parse().unwrap());
        // No mTLS metadata — should be rejected even with valid token.
        let err = interceptor.check_request(&req).unwrap_err();
        assert_eq!(err.code(), tonic::Code::Unauthenticated);
    }

    // -- sha256 / constant_time_eq --------------------------------------------

    #[test]
    fn sha256_is_deterministic() {
        assert_eq!(sha256(b"hello"), sha256(b"hello"));
        assert_ne!(sha256(b"hello"), sha256(b"world"));
    }

    #[test]
    fn constant_time_eq_equal() {
        let a = sha256(b"same");
        let b = sha256(b"same");
        assert!(constant_time_eq(&a, &b));
    }

    #[test]
    fn constant_time_eq_different() {
        let a = sha256(b"hello");
        let b = sha256(b"world");
        assert!(!constant_time_eq(&a, &b));
    }

    // -- RateLimiter ----------------------------------------------------------

    #[test]
    fn rate_limiter_allows_up_to_max() {
        let rl = RateLimiter::new(3, Duration::from_secs(60));
        assert!(rl.check("client").is_ok());
        assert!(rl.check("client").is_ok());
        assert!(rl.check("client").is_ok());
        assert!(rl.check("client").is_err());
    }

    #[test]
    fn rate_limiter_independent_keys() {
        let rl = RateLimiter::new(1, Duration::from_secs(60));
        assert!(rl.check("a").is_ok());
        assert!(rl.check("b").is_ok()); // independent key — not limited
        assert!(rl.check("a").is_err());
    }

    #[test]
    fn rate_limiter_retry_after_is_positive() {
        let rl = RateLimiter::new(1, Duration::from_secs(60));
        rl.check("x").unwrap();
        let retry_after = rl.check("x").unwrap_err();
        assert!(retry_after > Duration::ZERO);
        assert!(retry_after <= Duration::from_secs(60));
    }

    #[test]
    fn rate_limiter_cleanup_removes_expired() {
        // Use a very short window to force immediate expiry.
        let rl = RateLimiter::new(1, Duration::from_nanos(1));
        rl.check("key").unwrap();
        // The entry should be present.
        assert!(rl.limits.contains_key("key"));
        // After sleeping (even 0ms) the window expires; cleanup removes it.
        std::thread::sleep(Duration::from_millis(1));
        rl.cleanup();
        assert!(!rl.limits.contains_key("key"));
    }
}