PiCloud/crates/manager-core/src/auth_middleware.rs

//! Authentication middleware — resolves the caller's `Principal` from
//! either a session cookie / Bearer session-token OR an API key
//! (`Authorization: Bearer pic_…`). Both paths converge on the same
//! request extension so downstream handlers see one shape.
//!
//! Capability checks live in `crate::authz` and are called per-handler
//! (after the relevant resource is loaded, so the capability binds to
//! the actual resource's `app_id`). This middleware is gate-only: it
//! ensures *some* `Principal` is attached, or returns 401.
//!
//! Token discriminator: the `pic_` prefix on a Bearer value selects
//! the API-key path; anything else (raw 32-byte base64-url-encoded
//! string) takes the session path. The session cookie can only ever
//! carry a session token (cookies are never API keys).

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use axum::body::Body;
use axum::extract::{Request, State};
use axum::http::{header, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Json, Response};
use chrono::Utc;
use picloud_shared::{AdminUserId, Principal};
use serde_json::json;

use crate::admin_session_repo::AdminSessionRepository;
use crate::admin_user_repo::AdminUserRepository;
use crate::api_key_repo::{ApiKeyRepository, ApiKeyVerification};
use crate::auth::{hash_token, verify_password};

/// F-P-009: short-lived cache of resolved `(token → principal)` pairs.
/// Two reasons it's load-bearing:
/// 1. `verify_api_key` Argon2-verifies *every* candidate sharing the
///    8-char prefix per request — caching cuts that to once per token
///    per TTL window.
/// 2. `attach_principal_if_present` runs on every data-plane request
///    including paths that may not even need authz; even the session
///    path costs three round-trips (lookup + user-get + touch).
///
/// Keyed on the hashed token (sha256 in `hash_token`); value is the
/// resolved Principal + the instant it landed in the cache. TTL is
/// short — the audit cited 60-300s.
#[derive(Debug, Default)]
pub struct PrincipalCache {
    inner: Mutex<HashMap<String, (Instant, Principal)>>,
}

/// TTL on cached principals. Short enough that role / scope changes
/// take effect quickly; long enough to amortize Argon2 over many
/// adjacent requests on a hot key.
const PRINCIPAL_CACHE_TTL: Duration = Duration::from_secs(60);

/// F-S-006: maximum candidate API keys we'll Argon2-verify per
/// request. Hoisted to module scope so clippy::items_after_statements
/// is happy.
const MAX_API_KEY_CANDIDATES: usize = 16;

impl PrincipalCache {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    fn get(&self, token_hash: &str) -> Option<Principal> {
        let mut guard = self.inner.lock().ok()?;
        let entry = guard.get(token_hash)?;
        if entry.0.elapsed() > PRINCIPAL_CACHE_TTL {
            guard.remove(token_hash);
            return None;
        }
        Some(entry.1.clone())
    }

    fn insert(&self, token_hash: String, principal: Principal) {
        let Ok(mut guard) = self.inner.lock() else {
            return;
        };
        // Lazy GC: cap unbounded growth by sweeping expired entries
        // when the cache crosses an arbitrary threshold.
        if guard.len() > 1024 {
            let now = Instant::now();
            guard.retain(|_, v| now.duration_since(v.0) <= PRINCIPAL_CACHE_TTL);
        }
        guard.insert(token_hash, (Instant::now(), principal));
    }
}

pub const SESSION_COOKIE: &str = "picloud_session";

/// Prefix on the wire that selects the API-key path. The body that
/// follows is `base32(32 random bytes)`; the first 8 chars of the body
/// index into `api_keys.prefix` for verification.
pub const API_KEY_PREFIX: &str = "pic_";

/// Length of the indexed prefix portion of an API key (the 8 chars
/// immediately after `pic_`). Schema-side index is on this slice.
pub const API_KEY_PREFIX_LEN: usize = 8;

/// Shared state for auth: the user / session / API-key repos plus the
/// configured sliding session TTL. Cheap to clone (`Arc` everywhere).
#[derive(Clone)]
pub struct AuthState {
    pub users: Arc<dyn AdminUserRepository>,
    pub sessions: Arc<dyn AdminSessionRepository>,
    pub keys: Arc<dyn ApiKeyRepository>,
    pub ttl: Duration,
    /// F-P-009 — shared cache of resolved Principals. Constructed once
    /// at startup; cloned cheaply into every router state. `None` for
    /// tests / harnesses that don't wire it.
    pub principal_cache: Arc<PrincipalCache>,
}

/// Legacy request-extension alias retained so the (only remaining)
/// handler that pulled `AuthedAdmin` out — `GET /admin/auth/me` —
/// keeps compiling during the migration. New handlers should pull
/// `Extension<Principal>` directly.
#[deprecated(note = "use Extension<Principal> directly")]
#[derive(Debug, Clone)]
pub struct AuthedAdmin {
    pub id: AdminUserId,
    pub username: String,
}

/// Middleware entry point. Wire with
/// `axum::middleware::from_fn_with_state(auth_state, require_authenticated)`.
/// Inserts `Principal` (and the legacy `AuthedAdmin`) as request
/// extensions on success; returns 401 on any failure mode.
pub async fn require_authenticated(
    State(state): State<AuthState>,
    mut req: Request<Body>,
    next: Next,
) -> Response {
    let Some(token) = extract_token(&req) else {
        return unauthorized();
    };
    let principal = match resolve_principal(&state, &token).await {
        Ok(Some(p)) => p,
        Ok(None) => return unauthorized(),
        Err(InternalError) => return internal_error(),
    };

    let username_for_legacy = username_for(&state, principal.user_id).await;
    req.extensions_mut().insert(principal.clone());
    #[allow(deprecated)]
    if let Some(username) = username_for_legacy {
        req.extensions_mut().insert(AuthedAdmin {
            id: principal.user_id,
            username,
        });
    }
    next.run(req).await
}

/// Backwards-compatible alias — the single callsite that still names
/// `require_admin` keeps working without an immediate rename. New
/// wiring should call `require_authenticated`.
#[deprecated(note = "renamed to require_authenticated")]
pub async fn require_admin(state: State<AuthState>, req: Request<Body>, next: Next) -> Response {
    require_authenticated(state, req, next).await
}

/// Opportunistic data-plane variant: always inserts an
/// `Extension<Option<Principal>>` and forwards the request. Used on
/// `/execute/{id}` and the user-route fallback, where most invocations
/// are anonymous public HTTP and the few authed ones (dashboard
/// test-runs, API keys) should still let scripts see the caller via
/// `cx.principal` once services consume it.
///
/// Failure modes — all degrade to `None` rather than rejecting:
///   * No bearer / cookie → `None`.
///   * Malformed or unknown token → `None`.
///   * DB blip while resolving → `None` (fail-open; the data plane
///     should not 500 on transient infra failures for an *optional*
///     identity check).
///
/// Admin-side routes that REQUIRE an identity keep using
/// `require_authenticated`.
pub async fn attach_principal_if_present(
    State(state): State<AuthState>,
    mut req: Request<Body>,
    next: Next,
) -> Response {
    let principal: Option<Principal> = match extract_token(&req) {
        Some(token) => resolve_principal(&state, &token).await.unwrap_or(None),
        None => None,
    };
    req.extensions_mut().insert(principal);
    next.run(req).await
}

/// Decide whether the token is an API key (pic_ prefix) or a session
/// token, then resolve the corresponding `Principal`. `Ok(None)`
/// means the token was structurally valid but didn't match any active
/// credential; `Err(InternalError)` means a DB blip.
async fn resolve_principal(
    state: &AuthState,
    token: &str,
) -> Result<Option<Principal>, InternalError> {
    // F-P-009: cache hit short-circuits Argon2 + 3 DB round-trips.
    // Keyed on the hashed token so we never store raw bearer values.
    let token_hash = hash_token(token);
    if let Some(p) = state.principal_cache.get(&token_hash) {
        return Ok(Some(p));
    }
    let resolved = if let Some(rest) = token.strip_prefix(API_KEY_PREFIX) {
        verify_api_key(state, rest).await?
    } else {
        verify_session(state, token).await?
    };
    if let Some(p) = &resolved {
        state.principal_cache.insert(token_hash, p.clone());
    }
    Ok(resolved)
}

async fn verify_session(
    state: &AuthState,
    token: &str,
) -> Result<Option<Principal>, InternalError> {
    let token_hash = hash_token(token);

    let lookup = match state.sessions.lookup(&token_hash).await {
        Ok(Some(l)) => l,
        Ok(None) => return Ok(None),
        Err(err) => {
            tracing::error!(?err, "admin_sessions lookup failed");
            return Err(InternalError);
        }
    };

    let user = match state.users.get(lookup.user_id).await {
        Ok(Some(u)) if u.is_active => u,
        Ok(_) => return Ok(None),
        Err(err) => {
            tracing::error!(?err, "admin_users lookup failed");
            return Err(InternalError);
        }
    };

    // Sliding-window bump — inline so a DB blip surfaces as 500 rather
    // than silent stale sessions. Same shape as Phase 3a.
    let new_expires_at = Utc::now() + chrono::Duration::from_std(state.ttl).unwrap_or_default();
    if let Err(err) = state.sessions.touch(&token_hash, new_expires_at).await {
        tracing::error!(?err, "admin_sessions touch failed");
        return Err(InternalError);
    }

    Ok(Some(Principal {
        user_id: user.id,
        instance_role: user.instance_role,
        scopes: None,
        app_binding: None,
    }))
}

/// API-key verification path. `rest` is the portion of the bearer
/// value *after* `pic_`. We slice off the first 8 chars as the
/// indexed lookup key, then Argon2id-verify each candidate's hash
/// against the full `rest`. At most one match is expected; multiple
/// candidates with the same prefix is statistically negligible but
/// handled correctly (verify each, take the first match).
async fn verify_api_key(state: &AuthState, rest: &str) -> Result<Option<Principal>, InternalError> {
    if rest.len() <= API_KEY_PREFIX_LEN {
        return Ok(None);
    }
    let prefix = &rest[..API_KEY_PREFIX_LEN];

    let mut candidates = match state.keys.find_active_by_prefix(prefix).await {
        Ok(v) => v,
        Err(err) => {
            tracing::error!(?err, "api_keys lookup failed");
            return Err(InternalError);
        }
    };

    // F-S-006: bound the candidate set. Prefix collisions are
    // statistically rare for a 32-byte random body, but the 8-char
    // index space is small enough that an attacker who provisions
    // many keys against one indexed prefix could amplify each public
    // bearer-tagged request into M×Argon2 verifies. Cap at MAX and
    // log the truncation so operators can spot it.
    if candidates.len() > MAX_API_KEY_CANDIDATES {
        tracing::warn!(
            prefix,
            total = candidates.len(),
            kept = MAX_API_KEY_CANDIDATES,
            "api_keys prefix-bucket overflow; truncating candidate verify set"
        );
        candidates.truncate(MAX_API_KEY_CANDIDATES);
    }

    // F-P-002: Argon2id verify is CPU-bound (~tens of ms each at OWASP
    // defaults). Move it off the Tokio async worker so a hot user with
    // N prefix-colliding keys doesn't park the worker for N×Argon2.
    let rest_owned = rest.to_string();
    let matched: Option<ApiKeyVerification> = tokio::task::spawn_blocking(move || {
        candidates
            .into_iter()
            .find(|c| verify_password(&c.hash, &rest_owned))
    })
    .await
    .map_err(|e| {
        tracing::error!(error = ?e, "verify_api_key spawn_blocking join failed");
        InternalError
    })?;
    let Some(matched) = matched else {
        return Ok(None);
    };

    // Resolve the owning user. is_active = false → reject even if the
    // key itself hasn't been expired yet (the expire_all_for_user
    // cascade on deactivation is the primary defense; this is the
    // belt-and-suspenders check at request time).
    let user = match state.users.get(matched.user_id).await {
        Ok(Some(u)) if u.is_active => u,
        Ok(_) => return Ok(None),
        Err(err) => {
            tracing::error!(?err, "admin_users lookup for api key failed");
            return Err(InternalError);
        }
    };

    if let Err(err) = state.keys.touch_last_used(matched.id).await {
        tracing::error!(?err, "api_keys touch_last_used failed");
        // Soft-fail: a timestamp blip should not invalidate the
        // request. Continue with the resolved Principal.
    }

    Ok(Some(Principal {
        user_id: user.id,
        instance_role: user.instance_role,
        scopes: Some(matched.scopes),
        app_binding: matched.app_id,
    }))
}

/// Best-effort username lookup for the legacy `AuthedAdmin` extension.
/// Returns `None` on DB error (the caller treats `None` as "skip the
/// legacy extension"). New handlers use `Principal` and don't depend
/// on this.
async fn username_for(state: &AuthState, id: AdminUserId) -> Option<String> {
    match state.users.get(id).await {
        Ok(Some(u)) => Some(u.username),
        Ok(None) => None,
        Err(err) => {
            tracing::warn!(
                ?err,
                "username lookup for AuthedAdmin failed; skipping legacy ext"
            );
            None
        }
    }
}

/// Pull the bearer token out of an `Authorization` header (preferred)
/// or the `picloud_session` cookie (fallback for browser clients).
/// Same shape as Phase 3a; the cookie only ever carries session
/// tokens — no `pic_` prefix expected there.
fn extract_token(req: &Request<Body>) -> Option<String> {
    if let Some(value) = req.headers().get(header::AUTHORIZATION) {
        if let Ok(s) = value.to_str() {
            if let Some(token) = s.strip_prefix("Bearer ") {
                let trimmed = token.trim();
                if !trimmed.is_empty() {
                    return Some(trimmed.to_string());
                }
            }
        }
    }
    if let Some(value) = req.headers().get(header::COOKIE) {
        if let Ok(s) = value.to_str() {
            for chunk in s.split(';') {
                let chunk = chunk.trim();
                if let Some(rest) = chunk.strip_prefix(&format!("{SESSION_COOKIE}=")) {
                    if !rest.is_empty() {
                        return Some(rest.to_string());
                    }
                }
            }
        }
    }
    None
}

/// Sentinel returned from the resolve functions when a DB error should
/// produce a 500 rather than a 401. Empty struct because the actual
/// error is already logged at the failure site.
struct InternalError;

fn unauthorized() -> Response {
    (
        StatusCode::UNAUTHORIZED,
        Json(json!({ "error": "authentication required" })),
    )
        .into_response()
}

fn internal_error() -> Response {
    (
        StatusCode::INTERNAL_SERVER_ERROR,
        Json(json!({ "error": "internal error" })),
    )
        .into_response()
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::http::Request;
    use picloud_shared::InstanceRole;

    fn req_with_header(name: &str, value: &str) -> Request<Body> {
        Request::builder()
            .header(name, value)
            .body(Body::empty())
            .unwrap()
    }

    #[test]
    fn extracts_bearer_token() {
        let r = req_with_header("authorization", "Bearer abc123");
        assert_eq!(extract_token(&r).as_deref(), Some("abc123"));
    }

    #[test]
    fn extracts_bearer_pic_prefixed_token() {
        let r = req_with_header("authorization", "Bearer pic_abcdefghIJKL");
        assert_eq!(extract_token(&r).as_deref(), Some("pic_abcdefghIJKL"));
    }

    #[test]
    fn ignores_bearer_with_no_token() {
        let r = req_with_header("authorization", "Bearer ");
        assert_eq!(extract_token(&r), None);
    }

    #[test]
    fn extracts_cookie_token() {
        let r = req_with_header("cookie", "foo=bar; picloud_session=xyz; baz=qux");
        assert_eq!(extract_token(&r).as_deref(), Some("xyz"));
    }

    #[test]
    fn bearer_wins_over_cookie() {
        let r = Request::builder()
            .header("authorization", "Bearer header-token")
            .header("cookie", "picloud_session=cookie-token")
            .body(Body::empty())
            .unwrap();
        assert_eq!(extract_token(&r).as_deref(), Some("header-token"));
    }

    #[test]
    fn returns_none_when_neither_present() {
        let r = Request::builder().body(Body::empty()).unwrap();
        assert_eq!(extract_token(&r), None);
    }

    // Round-trip test for the unused-variable to keep `Principal`
    // visibly tied to InstanceRole — caught a real bug during dev when
    // the field order in the struct literal had drifted.
    #[test]
    fn principal_construction_is_explicit() {
        let p = Principal {
            user_id: AdminUserId::new(),
            instance_role: InstanceRole::Owner,
            scopes: None,
            app_binding: None,
        };
        assert_eq!(p.instance_role, InstanceRole::Owner);
        assert!(p.scopes.is_none());
        assert!(p.app_binding.is_none());
    }
}