quicproquo/crates/quicproquo-server/src/auth.rs

use std::net::IpAddr;
use std::sync::Arc;

use dashmap::DashMap;
use quicproquo_proto::node_capnp::auth;
use sha2::Digest;
use subtle::ConstantTimeEq;
use tokio::sync::Notify;
use zeroize::Zeroizing;

use crate::error_codes::*;

pub const SESSION_TTL_SECS: u64 = 24 * 60 * 60; // 24 hours
pub const PENDING_LOGIN_TTL_SECS: u64 = 300; // 5 minutes
pub const RATE_LIMIT_WINDOW_SECS: u64 = 60;
pub const RATE_LIMIT_MAX_ENQUEUES: u32 = 100;

#[derive(Clone)]
pub struct AuthConfig {
    /// Server bearer token — zeroized on drop to prevent memory disclosure.
    pub required_token: Option<Zeroizing<Vec<u8>>>,
    /// When true, a valid bearer token (no session) is accepted and the request's identity/key is used (dev/e2e only).
    /// CLI flag: --allow-insecure-auth / QPQ_ALLOW_INSECURE_AUTH.
    pub allow_insecure_identity_from_request: bool,
}

impl std::fmt::Debug for AuthConfig {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AuthConfig")
            .field("required_token", &self.required_token.as_ref().map(|_| "[REDACTED]"))
            .field("allow_insecure_identity_from_request", &self.allow_insecure_identity_from_request)
            .finish()
    }
}

impl AuthConfig {
    pub fn new(required_token: Option<String>, allow_insecure_identity_from_request: bool) -> Self {
        let required_token = required_token
            .filter(|s| !s.is_empty())
            .map(|s| Zeroizing::new(s.into_bytes()));
        Self {
            required_token,
            allow_insecure_identity_from_request,
        }
    }
}

#[derive(Clone)]
pub struct SessionInfo {
    #[allow(dead_code)]
    pub username: String,
    pub identity_key: Vec<u8>,
    #[allow(dead_code)]
    pub created_at: u64,
    pub expires_at: u64,
}

pub struct PendingLogin {
    pub state_bytes: Vec<u8>,
    pub created_at: u64,
}

pub struct RateEntry {
    pub count: u32,
    pub window_start: u64,
}

#[derive(Clone)]
pub struct AuthContext {
    pub token: Vec<u8>,
    pub identity_key: Option<Vec<u8>>,
}

pub fn current_timestamp() -> u64 {
    match std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH) {
        Ok(d) => d.as_secs(),
        Err(_) => {
            tracing::warn!("system time is before UNIX_EPOCH; using 0 for session/rate-limit timestamps");
            0
        }
    }
}

pub fn check_rate_limit(
    rate_limits: &DashMap<Vec<u8>, RateEntry>,
    token: &[u8],
) -> Result<(), capnp::Error> {
    let now = current_timestamp();
    let mut entry = rate_limits.entry(token.to_vec()).or_insert(RateEntry {
        count: 0,
        window_start: now,
    });

    if now - entry.window_start >= RATE_LIMIT_WINDOW_SECS {
        entry.count = 1;
        entry.window_start = now;
    } else {
        entry.count += 1;
        if entry.count > RATE_LIMIT_MAX_ENQUEUES {
            return Err(crate::error_codes::coded_error(
                E014_RATE_LIMITED,
                format!(
                    "rate limit exceeded: {} enqueues in {}s window",
                    RATE_LIMIT_MAX_ENQUEUES, RATE_LIMIT_WINDOW_SECS
                ),
            ));
        }
    }
    Ok(())
}

pub fn validate_auth(
    cfg: &AuthConfig,
    sessions: &DashMap<Vec<u8>, SessionInfo>,
    auth: Result<auth::Reader<'_>, capnp::Error>,
) -> Result<(), capnp::Error> {
    validate_auth_context(cfg, sessions, auth).map(|_| ())
}

pub fn validate_auth_context(
    cfg: &AuthConfig,
    sessions: &DashMap<Vec<u8>, SessionInfo>,
    auth: Result<auth::Reader<'_>, capnp::Error>,
) -> Result<AuthContext, capnp::Error> {
    let auth = auth?;
    let version = auth.get_version();

    if version != 1 {
        return Err(crate::error_codes::coded_error(
            E001_BAD_AUTH_VERSION,
            format!("unsupported auth version {} (expected 1)", version),
        ));
    }

    let token = auth
        .get_access_token()
        .map_err(|e| crate::error_codes::coded_error(E020_BAD_PARAMS, format!("auth.accessToken: {e}")))?
        .to_vec();

    if token.is_empty() {
        return Err(crate::error_codes::coded_error(
            E002_EMPTY_TOKEN,
            "auth.version=1 requires non-empty accessToken",
        ));
    }

    if let Some(expected) = &cfg.required_token {
        if expected.len() == token.len() && bool::from(expected.as_slice().ct_eq(&token)) {
            return Ok(AuthContext {
                token,
                identity_key: None,
            });
        }
    }

    if let Some(session) = sessions.get(&token) {
        let now = current_timestamp();
        if session.expires_at > now {
            let identity = if session.identity_key.is_empty() {
                None
            } else {
                Some(session.identity_key.clone())
            };

            return Ok(AuthContext {
                token,
                identity_key: identity,
            });
        }
        drop(session);
        sessions.remove(&token);
        return Err(crate::error_codes::coded_error(
            E017_SESSION_EXPIRED,
            "session token has expired",
        ));
    }

    Err(crate::error_codes::coded_error(E003_INVALID_TOKEN, "invalid accessToken"))
}

/// Validate a raw bearer token (no Cap'n Proto dependency).
/// Used by the WebSocket JSON-RPC bridge.
pub fn validate_token_raw(
    cfg: &AuthConfig,
    sessions: &DashMap<Vec<u8>, SessionInfo>,
    token: &[u8],
) -> Result<AuthContext, String> {
    if token.is_empty() {
        return Err("empty access token".to_string());
    }

    // Check static bearer token.
    if let Some(expected) = &cfg.required_token {
        if expected.len() == token.len() && bool::from(expected.as_slice().ct_eq(token)) {
            return Ok(AuthContext {
                token: token.to_vec(),
                identity_key: None,
            });
        }
    }

    // Check session tokens.
    if let Some(session) = sessions.get(token) {
        let now = current_timestamp();
        if session.expires_at > now {
            let identity = if session.identity_key.is_empty() {
                None
            } else {
                Some(session.identity_key.clone())
            };
            return Ok(AuthContext {
                token: token.to_vec(),
                identity_key: identity,
            });
        }
        drop(session);
        sessions.remove(token);
        return Err("session token has expired".to_string());
    }

    Err("invalid access token".to_string())
}

pub fn require_identity(auth_ctx: &AuthContext) -> Result<&[u8], capnp::Error> {
    match auth_ctx.identity_key.as_deref() {
        Some(ik) => Ok(ik),
        None => Err(crate::error_codes::coded_error(
            E003_INVALID_TOKEN,
            "access token is not identity-bound; login required",
        )),
    }
}

pub fn require_identity_match(auth_ctx: &AuthContext, expected: &[u8]) -> Result<(), capnp::Error> {
    let ik = require_identity(auth_ctx)?;
    if ik.len() != expected.len() || !bool::from(ik.ct_eq(expected)) {
        return Err(crate::error_codes::coded_error(
            E016_IDENTITY_MISMATCH,
            "access token is bound to a different identity",
        ));
    }
    Ok(())
}

/// When the token is a valid session, require it to match `request_identity`.
/// When the token is a bearer token (no identity) and `allow_insecure_identity_from_request` is true, accept the request identity (dev/e2e).
pub fn require_identity_or_request(
    auth_ctx: &AuthContext,
    request_identity: &[u8],
    allow_insecure: bool,
) -> Result<(), capnp::Error> {
    match auth_ctx.identity_key.as_deref() {
        Some(_) => require_identity_match(auth_ctx, request_identity),
        None if allow_insecure => Ok(()),
        None => Err(crate::error_codes::coded_error(
            E003_INVALID_TOKEN,
            "access token is not identity-bound; login required",
        )),
    }
}

pub fn fmt_hex(bytes: &[u8]) -> String {
    let hex: String = bytes.iter().map(|b| format!("{b:02x}")).collect();
    format!("{hex}…")
}

pub fn waiter(waiters: &DashMap<Vec<u8>, Arc<Notify>>, recipient_key: &[u8]) -> Arc<Notify> {
    waiters
        .entry(recipient_key.to_vec())
        .or_insert_with(|| Arc::new(Notify::new()))
        .clone()
}

pub const CONN_RATE_LIMIT_WINDOW_SECS: u64 = 60;
pub const CONN_RATE_LIMIT_MAX: u32 = 50;

/// Per-IP connection rate limiter. Returns `true` if the connection is allowed.
pub fn check_conn_rate_limit(
    conn_rate_limits: &DashMap<IpAddr, RateEntry>,
    ip: IpAddr,
) -> bool {
    let now = current_timestamp();
    let mut entry = conn_rate_limits.entry(ip).or_insert(RateEntry {
        count: 0,
        window_start: now,
    });

    if now - entry.window_start >= CONN_RATE_LIMIT_WINDOW_SECS {
        entry.count = 1;
        entry.window_start = now;
        true
    } else {
        entry.count += 1;
        entry.count <= CONN_RATE_LIMIT_MAX
    }
}

pub fn fingerprint(data: &[u8]) -> Vec<u8> {
    sha2::Sha256::digest(data).to_vec()
}

pub fn coded_error(code: &str, msg: impl std::fmt::Display) -> capnp::Error {
    crate::error_codes::coded_error(code, msg)
}