plexus_auth_core/

auth.rs

//! AuthContext and SessionValidator — relocated from plexus-core.
//!
//! Per AUTHZ-CORE-CRATE-1, the public type surface (fields, methods,
//! constructors) is preserved verbatim from plexus-core's previous home at
//! `plexus_core::plexus::auth`. The migration is mechanical; behavior is
//! unchanged. plexus-core re-exports these from here with a `#[deprecated]`
//! pointer at the new path so existing callers compile during the
//! deprecation window.
//!
//! Future work tightens the seal on `AuthContext::new` and the field
//! visibility (per AUTHZ-0 §"Crate-level isolation amplifies the seal").
//! That tightening is intentionally NOT in this ticket because it would
//! break ~30 direct construction sites across plexus-trak, plexus-transport
//! tests, and others — a workspace-wide blast radius outside this ticket's
//! scope. See `plans/AUTHZ/AUTHZ-CORE-CRATE-1-RUN-NOTES.md`.

use async_trait::async_trait;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::Value;

/// Per-connection authentication context, populated during WS upgrade.
///
/// This context is extracted from HTTP cookies (or other auth mechanisms) during
/// the WebSocket handshake and attached to the connection. Every RPC call on that
/// connection has access to this context.
///
/// # Multi-tenancy with Keycloak
///
/// When using Keycloak for multi-tenancy, the `AuthContext` typically contains:
/// - `user_id`: Keycloak user ID (sub claim from JWT)
/// - `session_id`: Keycloak session ID
/// - `roles`: User roles within the tenant/realm
/// - `metadata`: Additional JWT claims (realm, tenant ID, custom attributes)
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct AuthContext {
    /// User identifier (e.g., Keycloak sub claim, user UUID)
    pub user_id: String,

    /// Session identifier (e.g., Keycloak session ID)
    pub session_id: String,

    /// User roles (e.g., ["user", "admin"], Keycloak realm roles)
    pub roles: Vec<String>,

    /// Additional metadata (e.g., JWT claims, tenant/realm info, custom attributes)
    /// For Keycloak multi-tenancy, this typically includes:
    /// - `realm`: Keycloak realm name
    /// - `tenant_id`: Organization/tenant identifier
    /// - Any custom claims from the JWT token
    pub metadata: Value,
}

impl AuthContext {
    /// Create a new AuthContext.
    ///
    /// # Note
    ///
    /// This constructor is `pub` to preserve the existing public API across
    /// the workspace. AUTHZ-0's structural-defense vision calls for this
    /// constructor to be `pub(crate)`; tightening that seal lands in a
    /// follow-up ticket because of the workspace-wide blast radius. See
    /// `plans/AUTHZ/AUTHZ-CORE-CRATE-1-RUN-NOTES.md`.
    pub fn new(user_id: String, session_id: String, roles: Vec<String>, metadata: Value) -> Self {
        Self {
            user_id,
            session_id,
            roles,
            metadata,
        }
    }

    /// Create an anonymous/unauthenticated context.
    ///
    /// This can be used as a fallback when methods accept `Option<&AuthContext>`
    /// and no authentication was provided.
    pub fn anonymous() -> Self {
        Self {
            user_id: "anonymous".to_string(),
            session_id: String::new(),
            roles: vec![],
            metadata: Value::Null,
        }
    }

    /// Check if this context represents an authenticated user.
    pub fn is_authenticated(&self) -> bool {
        self.user_id != "anonymous" && !self.session_id.is_empty()
    }

    /// Check if the user has a specific role.
    pub fn has_role(&self, role: &str) -> bool {
        self.roles.iter().any(|r| r == role)
    }

    /// Get a metadata field as a string.
    pub fn get_metadata_string(&self, key: &str) -> Option<String> {
        self.metadata
            .get(key)
            .and_then(|v| v.as_str())
            .map(String::from)
    }

    /// Get the tenant/realm from metadata (Keycloak multi-tenancy).
    pub fn tenant(&self) -> Option<String> {
        self.get_metadata_string("tenant_id")
            .or_else(|| self.get_metadata_string("realm"))
    }

    /// Framework-only constructor: derive a callee `AuthContext` from a
    /// caller context and a [`ForwardDerivation`].
    ///
    /// This is the **only** path that mints a callee context from a
    /// caller's. It is `pub(crate)` to `plexus-auth-core` — no downstream
    /// crate can reach it. A [`ForwardPolicy`] impl returns a
    /// `ForwardDerivation` (parameters); the framework dispatch path
    /// (AUTHLANG-3 wires this into `plexus_core::route_to_child`) is the
    /// only caller of this constructor. Per AUTHZ-0 §"The sealed-type
    /// pattern": the policy proposes; the framework disposes.
    ///
    /// The `_immediate_caller_stamp` parameter is reserved for the
    /// principal-chain extension that AUTHLANG-3 lands (it will append the
    /// stamp to the callee's invocation chain when `AuthContext` grows the
    /// chain field — today's `AuthContext` does not carry one, so the
    /// parameter is bound for forward compatibility and accepted but
    /// otherwise ignored). Surfacing it now keeps the constructor's
    /// signature stable across the AUTHZ-0 sealed-context migration.
    ///
    /// # Derivation semantics
    ///
    /// Each flag on the derivation maps to a logical field group on the
    /// current `AuthContext`:
    ///
    /// - `keep_verified_user` — retains `user_id` and `session_id` (the
    ///   identity of the originator). When `false`, both are reset to
    ///   `AuthContext::anonymous`'s values (`"anonymous"` / empty).
    /// - `keep_roles` — retains the role vector. When `false`, the
    ///   callee's `roles` is empty.
    /// - `keep_capabilities` — reserved for the AUTHZ-DATA / AUTHZ-CRED
    ///   migration; today's `AuthContext` carries no capabilities field, so
    ///   this flag is a no-op. Surfacing it keeps the v1 policy shape
    ///   forward-compatible.
    /// - `keep_metadata` — retains the metadata bag. When `false`, the
    ///   callee's `metadata` is `Value::Null`.
    ///
    /// The constructor never grows the context: every field in the
    /// returned `AuthContext` is either copied from `caller_ctx` or reset
    /// to the corresponding empty value. There is no path through this
    /// function that adds an unrelated role or fabricates a user_id.
    ///
    /// [`ForwardDerivation`]: crate::forward::ForwardDerivation
    /// [`ForwardPolicy`]: crate::forward::ForwardPolicy
    ///
    pub(crate) fn derive_callee_context(
        caller_ctx: &AuthContext,
        derivation: &crate::forward::ForwardDerivation,
        _immediate_caller_stamp: &crate::principal::Principal,
    ) -> AuthContext {
        let (user_id, session_id) = if derivation.keep_verified_user {
            (caller_ctx.user_id.clone(), caller_ctx.session_id.clone())
        } else {
            ("anonymous".to_string(), String::new())
        };
        let roles = if derivation.keep_roles {
            caller_ctx.roles.clone()
        } else {
            Vec::new()
        };
        // `keep_capabilities` is a no-op today: the current AuthContext has no
        // capabilities field. The flag is preserved on `ForwardDerivation` for
        // forward compatibility with the AUTHZ-DATA / AUTHZ-CRED migration.
        let metadata = if derivation.keep_metadata {
            caller_ctx.metadata.clone()
        } else {
            Value::Null
        };
        AuthContext {
            user_id,
            session_id,
            roles,
            metadata,
        }
    }

    /// Scoped-callback API for deriving a callee context.
    ///
    /// The framework's dispatch path (plexus-core `route_to_child`) calls
    /// this with a `ForwardDerivation` and a caller-principal stamp; the
    /// closure receives the derived callee `AuthContext` by value and
    /// returns whatever the dispatch yields (typically a `Future`).
    /// Passing by value rather than reference is intentional: it allows
    /// the closure to move the callee into an async block so dispatch can
    /// await the child call while the callee lives inside the future's
    /// state machine.
    ///
    /// This is the public entry point for AUTHLANG-3. The underlying
    /// constructor [`derive_callee_context`] remains `pub(crate)` so the
    /// raw "mint a callee from a caller" symbol is not callable from
    /// outside `plexus-auth-core`. Anyone can still call `AuthContext::new`
    /// and craft their own context from scratch — what they cannot do is
    /// obtain one through the framework-blessed derivation path except
    /// inside this callback, where the lifetime is scoped to the dispatch
    /// invocation.
    ///
    /// Per AUTHZ-0 §"The sealed-type pattern": the policy proposes (via
    /// `ForwardDerivation`); the framework disposes (via this callback).
    pub fn with_callee_context<F, R>(
        &self,
        derivation: &crate::forward::ForwardDerivation,
        immediate_caller_stamp: &crate::principal::Principal,
        f: F,
    ) -> R
    where
        F: FnOnce(AuthContext) -> R,
    {
        f(Self::derive_callee_context(self, derivation, immediate_caller_stamp))
    }
}

/// Backends implement this trait to validate cookies/tokens during WS upgrade.
///
/// This trait is designed to be object-safe and work with async/await, allowing
/// backends to use any authentication mechanism:
/// - JWT validation (e.g., Keycloak tokens)
/// - Database session lookups
/// - Redis session stores
/// - OAuth token introspection
///
/// # Example: Keycloak JWT Validation
///
/// ```rust,ignore
/// use plexus_auth_core::{AuthContext, SessionValidator};
/// use async_trait::async_trait;
///
/// struct KeycloakValidator {
///     jwks_client: JwksClient,
///     realm: String,
/// }
///
/// #[async_trait]
/// impl SessionValidator for KeycloakValidator {
///     async fn validate(&self, cookie_value: &str) -> Option<AuthContext> {
///         // Parse JWT from cookie
///         let token = parse_jwt_from_cookie(cookie_value)?;
///
///         // Validate signature and claims
///         let claims = self.jwks_client.verify(&token).await.ok()?;
///
///         // Extract user info and tenant from JWT claims
///         Some(AuthContext::new(
///             claims.sub,
///             claims.sid.unwrap_or_default(),
///             claims.realm_access.roles,
///             serde_json::json!({
///                 "realm": self.realm,
///                 "tenant_id": claims.get("tenant_id"),
///                 "email": claims.email,
///             }),
///         ))
///     }
/// }
/// ```
#[async_trait]
pub trait SessionValidator: Send + Sync + 'static {
    /// Validate a cookie header value and return an AuthContext if valid.
    ///
    /// # Arguments
    ///
    /// * `cookie_value` - The raw Cookie header value (e.g., "session=abc123; path=/")
    ///
    /// # Returns
    ///
    /// - `Some(AuthContext)` if the cookie is valid and represents an authenticated session
    /// - `None` if the cookie is invalid, expired, or represents an anonymous session
    ///
    /// # Implementation Notes
    ///
    /// - This is called during the WebSocket handshake (HTTP upgrade)
    /// - Validation should be fast to avoid blocking the connection
    /// - For JWT: verify signature, check expiration, extract claims
    /// - For session-based auth: lookup session in DB/Redis
    /// - Return None for invalid/expired credentials (connection proceeds as anonymous)
    async fn validate(&self, cookie_value: &str) -> Option<AuthContext>;
}

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

    #[test]
    fn test_auth_context_creation() {
        let ctx = AuthContext::new(
            "user-123".to_string(),
            "sess-456".to_string(),
            vec!["admin".to_string()],
            serde_json::json!({"tenant_id": "acme"}),
        );

        assert_eq!(ctx.user_id, "user-123");
        assert_eq!(ctx.session_id, "sess-456");
        assert!(ctx.has_role("admin"));
        assert!(!ctx.has_role("user"));
        assert_eq!(ctx.tenant(), Some("acme".to_string()));
        assert!(ctx.is_authenticated());
    }

    #[test]
    fn test_auth_context_clone() {
        let ctx = AuthContext::new(
            "alice".to_string(),
            "sess-1".to_string(),
            vec!["admin".to_string()],
            serde_json::json!({"org": "acme"}),
        );

        let cloned = ctx.clone();
        assert_eq!(ctx.user_id, cloned.user_id);
        assert_eq!(ctx.session_id, cloned.session_id);
        assert_eq!(ctx.roles, cloned.roles);
    }

    #[test]
    fn test_anonymous_context() {
        let ctx = AuthContext::anonymous();
        assert_eq!(ctx.user_id, "anonymous");
        assert!(!ctx.is_authenticated());
        assert!(ctx.roles.is_empty());
    }

    #[test]
    fn test_role_checking() {
        let ctx = AuthContext::new(
            "user-1".to_string(),
            "sess-1".to_string(),
            vec!["user".to_string(), "editor".to_string()],
            Value::Null,
        );

        assert!(ctx.has_role("user"));
        assert!(ctx.has_role("editor"));
        assert!(!ctx.has_role("admin"));
    }

    #[test]
    fn test_metadata_access() {
        let ctx = AuthContext::new(
            "user-1".to_string(),
            "sess-1".to_string(),
            vec![],
            serde_json::json!({
                "tenant_id": "org-123",
                "realm": "production",
                "email": "user@example.com"
            }),
        );

        assert_eq!(
            ctx.get_metadata_string("tenant_id"),
            Some("org-123".to_string())
        );
        assert_eq!(
            ctx.get_metadata_string("realm"),
            Some("production".to_string())
        );
        assert_eq!(
            ctx.get_metadata_string("email"),
            Some("user@example.com".to_string())
        );
        assert_eq!(ctx.get_metadata_string("nonexistent"), None);
    }


    #[test]
    fn derive_callee_context_identity_only_strips_roles_and_metadata() {
        use crate::forward::ForwardDerivation;
        use crate::principal::Principal;

        let caller = AuthContext::new(
            "alice".to_string(),
            "sess-1".to_string(),
            vec!["admin".to_string(), "editor".to_string()],
            serde_json::json!({"tenant_id": "acme"}),
        );
        let stamp = Principal::anonymous_sealed();
        let callee = AuthContext::derive_callee_context(
            &caller,
            &ForwardDerivation::IDENTITY_ONLY,
            &stamp,
        );
        assert_eq!(callee.user_id, "alice");
        assert_eq!(callee.session_id, "sess-1");
        assert!(callee.roles.is_empty());
        assert_eq!(callee.metadata, Value::Null);
    }

    #[test]
    fn derive_callee_context_pass_through_retains_all_fields() {
        use crate::forward::ForwardDerivation;
        use crate::principal::Principal;

        let caller = AuthContext::new(
            "alice".to_string(),
            "sess-1".to_string(),
            vec!["admin".to_string()],
            serde_json::json!({"tenant_id": "acme", "k": "v"}),
        );
        let stamp = Principal::anonymous_sealed();
        let callee = AuthContext::derive_callee_context(
            &caller,
            &ForwardDerivation::PASS_THROUGH,
            &stamp,
        );
        assert_eq!(callee.user_id, caller.user_id);
        assert_eq!(callee.session_id, caller.session_id);
        assert_eq!(callee.roles, caller.roles);
        assert_eq!(callee.metadata, caller.metadata);
    }

    #[test]
    fn derive_callee_context_anonymous_drops_everything() {
        use crate::forward::ForwardDerivation;
        use crate::principal::Principal;

        let caller = AuthContext::new(
            "alice".to_string(),
            "sess-1".to_string(),
            vec!["admin".to_string()],
            serde_json::json!({"tenant_id": "acme"}),
        );
        let stamp = Principal::anonymous_sealed();
        let callee = AuthContext::derive_callee_context(
            &caller,
            &ForwardDerivation::ANONYMOUS,
            &stamp,
        );
        assert_eq!(callee.user_id, "anonymous");
        assert_eq!(callee.session_id, "");
        assert!(callee.roles.is_empty());
        assert_eq!(callee.metadata, Value::Null);
        assert!(!callee.is_authenticated());
    }

    #[test]
    fn with_callee_context_invokes_closure_with_derived_callee() {
        use crate::forward::ForwardDerivation;
        use crate::principal::Principal;

        let caller = AuthContext::new(
            "alice".to_string(),
            "sess-1".to_string(),
            vec!["admin".to_string()],
            serde_json::json!({"tenant_id": "org-1"}),
        );
        let stamp = Principal::anonymous_sealed();

        let observed_user_id =
            caller.with_callee_context(&ForwardDerivation::IDENTITY_ONLY, &stamp, |callee| {
                assert!(callee.roles.is_empty());
                assert_eq!(callee.metadata, Value::Null);
                callee.user_id
            });

        assert_eq!(observed_user_id, "alice");
    }

    #[test]
    fn with_callee_context_returns_closure_value() {
        use crate::forward::ForwardDerivation;
        use crate::principal::Principal;

        let caller = AuthContext::anonymous();
        let stamp = Principal::anonymous_sealed();
        let answer =
            caller.with_callee_context(&ForwardDerivation::PASS_THROUGH, &stamp, |_| 42_u32);
        assert_eq!(answer, 42);
    }

    #[test]
    fn derive_callee_context_never_grows_context() {
        // Sanity: the constructor cannot fabricate fields that did not
        // exist on the caller. Starting from anonymous, even pass_through
        // produces anonymous — the derivation can keep what the caller
        // had, but never add what the caller lacked.
        use crate::forward::ForwardDerivation;
        use crate::principal::Principal;

        let caller = AuthContext::anonymous();
        let stamp = Principal::anonymous_sealed();
        let callee = AuthContext::derive_callee_context(
            &caller,
            &ForwardDerivation::PASS_THROUGH,
            &stamp,
        );
        assert_eq!(callee.user_id, "anonymous");
        assert!(callee.roles.is_empty());
        assert_eq!(callee.metadata, Value::Null);
        assert!(!callee.is_authenticated());
    }

    #[test]
    fn test_tenant_from_metadata() {
        // tenant_id takes precedence
        let ctx1 = AuthContext::new(
            "user-1".to_string(),
            "sess-1".to_string(),
            vec![],
            serde_json::json!({"tenant_id": "org-123", "realm": "prod"}),
        );
        assert_eq!(ctx1.tenant(), Some("org-123".to_string()));

        // Falls back to realm if no tenant_id
        let ctx2 = AuthContext::new(
            "user-1".to_string(),
            "sess-1".to_string(),
            vec![],
            serde_json::json!({"realm": "prod"}),
        );
        assert_eq!(ctx2.tenant(), Some("prod".to_string()));

        // None if neither present
        let ctx3 = AuthContext::new(
            "user-1".to_string(),
            "sess-1".to_string(),
            vec![],
            Value::Null,
        );
        assert_eq!(ctx3.tenant(), None);
    }
}