klieo-a2a 3.4.0

//! A2A-specific request-context bag.
//!
//! Authentication primitives (`Authenticator`, `AuthError`, `Identity`,
//! `BearerTokenAuthenticator`, `AllowAnonymous`, etc.) live in
//! `klieo-auth-common` — import them from there.
//!
//! [`RequestContext`] stays here because it wraps an A2A-specific
//! [`crate::envelope::A2aHeaders`] bag.

use crate::envelope::A2aHeaders;
use klieo_auth_common::Identity as AuthIdentity;

/// Context passed to every [`A2aHandler`](crate::handler::A2aHandler) method.
///
/// Holds the decoded transport headers, the verified caller identity, and
/// a cooperative cancel signal (set by the HTTP/SSE transport when a client
/// disconnects mid-stream).
#[derive(Debug, Clone)]
pub struct RequestContext {
    /// All A2A standard headers decoded from the inbound message.
    pub headers: A2aHeaders,
    /// The verified caller identity. `None` is reserved for the explicit
    /// [`klieo_auth_common::AllowAnonymous`] path; production authenticators
    /// always set this to `Some(...)` on success.
    pub caller: Option<AuthIdentity>,
    /// Cooperative cancellation token. HTTP/SSE transports pass a
    /// request-scoped child token; handlers should `select!` against
    /// `cancel.cancelled()` to honor client disconnects.
    pub cancel: tokio_util::sync::CancellationToken,
}

impl RequestContext {
    /// Construct a request context with default `cancel` (fresh, uncancelled).
    pub fn new(headers: A2aHeaders, caller: Option<AuthIdentity>) -> Self {
        Self {
            headers,
            caller,
            cancel: tokio_util::sync::CancellationToken::new(),
        }
    }

    /// Override the cancellation token (HTTP/SSE callers pass a
    /// request-scoped child token).
    #[must_use]
    pub fn with_cancel(mut self, cancel: tokio_util::sync::CancellationToken) -> Self {
        self.cancel = cancel;
        self
    }

    /// Returns `true` iff a caller identity is present.
    pub fn is_authenticated(&self) -> bool {
        self.caller.is_some()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::envelope::A2aHeaders;
    use async_trait::async_trait;
    #[allow(deprecated)]
    use klieo_auth_common::{AllowAnonymous, AuthError, Authenticator, BearerTokenAuthenticator};
    use klieo_auth_common::{Headers as AuthHeaders, Identity};
    use tokio_util::sync::CancellationToken;

    fn headers_with(auth: Option<&str>) -> A2aHeaders {
        let mut h = klieo_core::Headers::new();
        if let Some(value) = auth {
            h.insert("Authorization".into(), value.into());
        }
        A2aHeaders::decode_from(&h)
    }

    #[tokio::test]
    async fn identity_anonymous_round_trips() {
        let id = Identity::anonymous();
        assert!(id.is_anonymous());
        assert_eq!(id.as_str(), "anonymous");
    }

    #[tokio::test]
    async fn allow_anonymous_returns_anonymous_identity_for_any_request() {
        let authn = AllowAnonymous;
        let h = headers_with(None);
        let id = authn
            .authenticate(&h as &dyn AuthHeaders, b"{}")
            .await
            .expect("AllowAnonymous never errors");
        assert!(id.is_anonymous());
    }

    #[tokio::test]
    async fn bearer_rejects_missing_authorization_header() {
        let authn = BearerTokenAuthenticator::new(|_| Ok(Identity::new("never-called")));
        let h = headers_with(None);
        let err = authn
            .authenticate(&h as &dyn AuthHeaders, b"{}")
            .await
            .unwrap_err();
        assert!(matches!(err, AuthError::Missing));
    }

    #[tokio::test]
    async fn bearer_rejects_wrong_scheme() {
        let authn = BearerTokenAuthenticator::new(|_| Ok(Identity::new("never-called")));
        let h = headers_with(Some("Basic abc"));
        let err = authn
            .authenticate(&h as &dyn AuthHeaders, b"{}")
            .await
            .unwrap_err();
        assert!(matches!(err, AuthError::Malformed));
    }

    /// W2.A14 / RFC 7235 §2.1: scheme matching must be case-insensitive.
    #[tokio::test]
    async fn bearer_accepts_case_insensitive_scheme() {
        let authn = BearerTokenAuthenticator::new(|tok| {
            assert_eq!(tok, "abc");
            Ok(Identity::new("alice"))
        });
        for header in ["bearer abc", "BEARER abc", "BeArEr abc", "Bearer\tabc"] {
            let h = headers_with(Some(header));
            authn
                .authenticate(&h as &dyn AuthHeaders, b"{}")
                .await
                .unwrap_or_else(|e| panic!("header {header:?} rejected: {e:?}"));
        }
    }

    #[tokio::test]
    async fn bearer_rejects_empty_token() {
        let authn = BearerTokenAuthenticator::new(|_| Ok(Identity::new("never-called")));
        let h = headers_with(Some("Bearer "));
        let err = authn
            .authenticate(&h as &dyn AuthHeaders, b"{}")
            .await
            .unwrap_err();
        assert!(matches!(err, AuthError::Malformed));
    }

    #[tokio::test]
    async fn bearer_delegates_token_to_verifier_and_returns_identity() {
        let authn = BearerTokenAuthenticator::new(|tok| {
            if tok == "good" {
                Ok(Identity::new("alice"))
            } else {
                Err(AuthError::Rejected("bad".into()))
            }
        });
        let ok = headers_with(Some("Bearer good"));
        let ok_id = authn
            .authenticate(&ok as &dyn AuthHeaders, b"{}")
            .await
            .unwrap();
        assert_eq!(ok_id.as_str(), "alice");
        let bad = headers_with(Some("Bearer bad"));
        let err = authn
            .authenticate(&bad as &dyn AuthHeaders, b"{}")
            .await
            .unwrap_err();
        assert!(matches!(err, AuthError::Rejected(_)));
    }

    #[tokio::test]
    async fn request_context_authenticated_flag_reflects_caller() {
        let ctx = RequestContext::new(headers_with(None), Some(Identity::new("alice")));
        assert!(ctx.is_authenticated());
        let ctx = RequestContext::new(headers_with(None), None);
        assert!(!ctx.is_authenticated());
    }

    #[test]
    fn request_context_with_cancel_overrides_default() {
        let token = CancellationToken::new();
        let ctx = RequestContext::new(headers_with(None), None).with_cancel(token.clone());
        token.cancel();
        assert!(ctx.cancel.is_cancelled());
    }

    #[test]
    fn request_context_new_has_fresh_uncancelled_token() {
        let ctx = RequestContext::new(headers_with(None), None);
        assert!(!ctx.cancel.is_cancelled());
    }

    #[tokio::test]
    async fn allow_anonymous_authorize_method_returns_ok_for_any_method() {
        let authn = AllowAnonymous;
        let id = Identity::anonymous();
        for method in ["SendMessage", "GetTask", "CancelTask"] {
            authn
                .authorize_method(&id, method)
                .await
                .unwrap_or_else(|e| panic!("AllowAnonymous rejected {method:?}: {e:?}"));
        }
    }

    #[tokio::test]
    async fn bearer_authorize_method_default_impl_returns_ok() {
        let authn = BearerTokenAuthenticator::new(|_| Ok(Identity::new("never-called")));
        let id = Identity::new("alice");
        authn
            .authorize_method(&id, "SendMessage")
            .await
            .expect("default authorize_method impl returns Ok");
    }

    /// Verifies an authenticator can reject a method-scope mismatch
    /// without changing the dispatch shape — the T3 OAuth impl will
    /// follow this pattern.
    #[tokio::test]
    async fn custom_authenticator_can_reject_method_via_authorize_method() {
        struct ScopeGatedAuthn;
        #[async_trait]
        impl Authenticator for ScopeGatedAuthn {
            async fn authenticate(
                &self,
                _headers: &dyn AuthHeaders,
                _payload: &[u8],
            ) -> Result<Identity, AuthError> {
                Ok(Identity::new("scoped-caller"))
            }

            async fn authorize_method(
                &self,
                _identity: &Identity,
                method: &str,
            ) -> Result<(), AuthError> {
                if method == "CancelTask" {
                    Err(AuthError::Rejected("scope mismatch".into()))
                } else {
                    Ok(())
                }
            }
        }

        let authn = ScopeGatedAuthn;
        let id = Identity::new("scoped-caller");
        authn
            .authorize_method(&id, "SendMessage")
            .await
            .expect("SendMessage permitted by scope");
        let err = authn
            .authorize_method(&id, "CancelTask")
            .await
            .expect_err("CancelTask must be rejected");
        assert!(matches!(err, AuthError::Rejected(msg) if msg == "scope mismatch"));
    }
}