pas-external 4.0.1

//! SvAware session resolver — STANDARDS_AUTH_INVALIDATION §5.3.
//!
//! Wraps the base [`SessionResolver`] with an `sv`-version check against
//! a consumer-local cache. On miss or stale, refreshes the PAS access
//! token to pick up the current `sv`, updates both the session store
//! and the cache, and re-checks.
//!
//! The cache TTL (60 s) bounds the cross-cluster propagation latency
//! for SDK consumers: a PAS-side break-glass / lifecycle invalidation
//! becomes effective for the consumer within `~1 PAS round-trip + cache
//! TTL` per spec §4.3. PAS-internal and PCS chat-auth use different
//! substrates (shared KVRocks); this module covers the cross-cluster
//! case.

use std::sync::Arc;

use axum_extra::extract::PrivateCookieJar;

use super::session::{SessionResolution, SessionResolver};
use super::traits::{RefreshTokenResolver, SessionStore, SvAware};
use crate::oauth::AuthClient;
use crate::session_version::{SV_CACHE_KEY_PREFIX, SV_CACHE_TTL, SessionVersionCache};
use crate::types::SessionId;

/// Wrap a base [`SessionResolver`] with sv enforcement + refresh fallback.
///
/// Construct via [`super::PasAuth::resolver`]
/// (returns `SvAwareSessionResolver<S, MemorySessionVersionCache>`) or
/// [`super::PasAuth::resolver_with_cache`] for an external cache (e.g.
/// shared Redis across consumer pods).
pub struct SvAwareSessionResolver<S: SessionStore, R: RefreshTokenResolver, C: SessionVersionCache> {
    base: SessionResolver<S>,
    store: Arc<S>,
    refresh_resolver: Arc<R>,
    auth_client: Arc<AuthClient>,
    cache: Arc<C>,
}

// Manual Clone — derive would add `S: Clone, R: Clone, C: Clone` bounds.
impl<S, R, C> Clone for SvAwareSessionResolver<S, R, C>
where
    S: SessionStore,
    R: RefreshTokenResolver,
    C: SessionVersionCache,
{
    fn clone(&self) -> Self {
        Self {
            base: self.base.clone(),
            store: Arc::clone(&self.store),
            refresh_resolver: Arc::clone(&self.refresh_resolver),
            auth_client: Arc::clone(&self.auth_client),
            cache: Arc::clone(&self.cache),
        }
    }
}

impl<S, R, C> SvAwareSessionResolver<S, R, C>
where
    S: SessionStore,
    R: RefreshTokenResolver,
    C: SessionVersionCache,
{
    pub(super) fn new(
        base: SessionResolver<S>,
        store: Arc<S>,
        refresh_resolver: Arc<R>,
        auth_client: Arc<AuthClient>,
        cache: Arc<C>,
    ) -> Self {
        Self {
            base,
            store,
            refresh_resolver,
            auth_client,
            cache,
        }
    }

    /// Resolve a session cookie into an authentication context, with
    /// `sv` enforcement.
    ///
    /// Flow:
    /// 1. Run the base resolver. Propagate `NoCookie` / `Expired` unchanged.
    /// 2. If the session has no `sv` claim (AI agent, DEV_AUTH), admit.
    /// 3. Cache hit + `token_sv >= cached_sv` → admit.
    /// 4. Otherwise (miss or stale) → call `/token` refresh, fetch fresh
    ///    sv via userinfo, update store + cache, re-check.
    /// 5. Refresh failure → `Expired`.
    ///
    /// # Errors
    ///
    /// Surfaces [`SessionStore::Error`] from the underlying store
    /// lookup. All other failure modes (refresh failure, cache backend
    /// error, missing refresh-token) are flattened into `Expired` so
    /// the consumer redirects to login rather than 500ing.
    pub async fn resolve(
        &self,
        jar: &PrivateCookieJar,
    ) -> Result<SessionResolution<S::AuthContext>, S::Error> {
        let resolution = self.base.resolve(jar).await?;
        let session = match resolution {
            SessionResolution::Authenticated(s) => s,
            other => return Ok(other),
        };

        // sv claim absent → admit (spec §4.2.1: AI agent / DEV_AUTH).
        let Some(token_sv) = session.sv() else {
            return Ok(SessionResolution::Authenticated(session));
        };

        let cache_key = format!("{SV_CACHE_KEY_PREFIX}{}", session.ppnum_id());

        match self.cache.get(&cache_key).await {
            Some(cached_sv) if token_sv >= cached_sv => {
                Ok(SessionResolution::Authenticated(session))
            }
            _ => self.refresh_and_recheck(jar, session, &cache_key).await,
        }
    }

    async fn refresh_and_recheck(
        &self,
        jar: &PrivateCookieJar,
        stale_session: S::AuthContext,
        cache_key: &str,
    ) -> Result<SessionResolution<S::AuthContext>, S::Error> {
        // 1. Extract session_id from the cookie. The base resolver
        //    already proved the cookie is present + the session exists,
        //    so a missing cookie here is anomalous — fall through to
        //    Expired so the consumer redirects to login.
        let Some(cookie) = jar.get(self.base.cookie_name()) else {
            return Ok(SessionResolution::Expired);
        };
        let session_id = SessionId(cookie.value().to_string());

        // 2. Fetch decrypted refresh_token from the consumer.
        let refresh_token = match self
            .refresh_resolver
            .resolve_refresh_token(&session_id)
            .await
        {
            Ok(Some(rt)) => rt,
            Ok(None) => {
                tracing::debug!(
                    session_id = %session_id,
                    "no refresh_token available — surface as Expired"
                );
                return Ok(SessionResolution::Expired);
            }
            Err(e) => {
                tracing::warn!(
                    session_id = %session_id,
                    error = %e,
                    "refresh_token lookup failed — surface as Expired"
                );
                return Ok(SessionResolution::Expired);
            }
        };

        // 3. Call /token to mint a fresh access_token. The new token
        //    necessarily carries the current sv claim — PAS always
        //    embeds the post-bump value.
        let token_response = match self.auth_client.refresh_token(&refresh_token).await {
            Ok(t) => t,
            Err(e) => {
                tracing::info!(
                    session_id = %session_id,
                    error = %e,
                    "token refresh failed (likely revoked refresh_token) — Expired"
                );
                return Ok(SessionResolution::Expired);
            }
        };

        // 4. Read the new sv from userinfo (matches the OAuth callback
        //    pattern — PAS exposes session_version via /userinfo).
        let new_sv = match self
            .auth_client
            .get_user_info(&token_response.access_token)
            .await
        {
            Ok(info) => match info.session_version {
                Some(sv) => sv,
                None => {
                    // Token was minted but carries no sv — should only
                    // happen for AI-agent paths which the cookie-session
                    // middleware never touches. Treat as admit (consistent
                    // with the early-return above).
                    return Ok(SessionResolution::Authenticated(stale_session));
                }
            },
            Err(e) => {
                tracing::warn!(
                    session_id = %session_id,
                    error = %e,
                    "post-refresh userinfo failed — Expired"
                );
                return Ok(SessionResolution::Expired);
            }
        };

        // 5. Persist new sv in the session store + consumer-local cache.
        if let Err(e) = self.store.update_sv(&session_id, new_sv).await {
            // Update_sv failures are unusual but not catastrophic —
            // the next resolve will hit the cache entry below and
            // re-fetch only if cache eviction occurs first.
            tracing::warn!(
                session_id = %session_id,
                new_sv = new_sv,
                error = %e,
                "update_sv failed after refresh — cache will short-circuit",
            );
        }
        self.cache.set(cache_key, new_sv, SV_CACHE_TTL).await;

        // 6. Re-fetch the session so the returned AuthContext carries
        //    the updated sv. find() returning None means the row was
        //    deleted concurrently — surface as Expired.
        match self.store.find(&session_id).await? {
            Some(refreshed) => Ok(SessionResolution::Authenticated(refreshed)),
            None => Ok(SessionResolution::Expired),
        }
    }
}