pas-external 4.0.2

//! 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.
    ///
    /// Failure handling is **fail-CLOSED** for every step beyond the base
    /// resolver: refresh-token lookup error, `/token` failure, userinfo
    /// failure, anomalous `session_version=None` on a Human session, and
    /// `update_sv` persistence failure all surface as `Expired`.
    ///
    /// # Errors
    ///
    /// Surfaces [`SessionStore::Error`] from the underlying store
    /// lookup. All other failure modes 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, &cache_key).await,
        }
    }

    async fn refresh_and_recheck(
        &self,
        jar: &PrivateCookieJar,
        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 => {
                    // We reached this branch because `session.sv() ==
                    // Some(token_sv)` — a Human-entity session. Yet PAS
                    // just minted a fresh access_token whose userinfo
                    // omits `session_version`. That is anomalous (PAS
                    // regression, dev/prod skew, proxy mangling), not
                    // the AI-agent admit case. Fail CLOSED: do not
                    // bypass the very check this resolver exists for.
                    tracing::error!(
                        session_id = %session_id,
                        "Human session refreshed but userinfo.session_version=None — Expired",
                    );
                    return Ok(SessionResolution::Expired);
                }
            },
            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 the store rejects the update (DB outage, serialization
        //    error), we MUST NOT update the cache and admit — that would
        //    leave the in-memory state diverged from the durable store
        //    indefinitely. Fail CLOSED so the consumer redirects to
        //    login; the next attempt retries persistence.
        if let Err(e) = self.store.update_sv(&session_id, new_sv).await {
            tracing::warn!(
                session_id = %session_id,
                new_sv = new_sv,
                error = %e,
                "update_sv failed after refresh — Expired",
            );
            return Ok(SessionResolution::Expired);
        }
        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),
        }
    }
}