klieo-a2a 0.4.0

//! Authentication port for the A2A transport.
//!
//! [`A2aServer`](crate::server::A2aServer) calls an [`Authenticator`] before
//! dispatching every inbound request. The verified [`Identity`] is then
//! threaded into the handler via [`RequestContext`] so handler authors can
//! make per-method authorisation decisions.
//!
//! # Why this lives at the server layer
//!
//! The A2A protocol carries credentials in the `Authorization` header
//! (extracted by [`crate::envelope::A2aHeaders`]). Before this module
//! existed, the header was decoded but never propagated to handlers — any
//! caller publishable on `<app_prefix>.a2a.<agent_id>.rpc` could invoke any
//! method. Pre-dispatch authentication closes that gap and gives handlers
//! a uniform identity surface.
//!
//! # Two implementations
//!
//! - [`BearerTokenAuthenticator`] — validates `Authorization: Bearer <token>`
//!   against a caller-provided verifier closure. Wrap your JWT/opaque-token
//!   verification logic in the closure.
//! - [`AllowAnonymous`] — explicit opt-in for tests, demos, and the
//!   conformance suite. Production code must not use this.

use crate::envelope::A2aHeaders;
use async_trait::async_trait;
use thiserror::Error;

/// Verified caller identity returned by an [`Authenticator`].
///
/// Wraps an opaque string (typically a JWT subject claim, service-account
/// name, or peer agent id depending on the authenticator). Handler authors
/// can pattern-match on `as_str()` to make per-method authorisation
/// decisions.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Identity(String);

impl Identity {
    /// Wrap a verified identity string.
    pub fn new(value: impl Into<String>) -> Self {
        Self(value.into())
    }

    /// Sentinel identity returned by [`AllowAnonymous`]. Handlers can
    /// detect this and refuse to act for mutating methods.
    pub fn anonymous() -> Self {
        Self("anonymous".into())
    }

    /// Borrow the underlying identity string.
    pub fn as_str(&self) -> &str {
        &self.0
    }

    /// Returns `true` iff this identity was produced by [`AllowAnonymous`].
    pub fn is_anonymous(&self) -> bool {
        self.0 == "anonymous"
    }
}

/// Context passed to every [`A2aHandler`](crate::handler::A2aHandler) method.
///
/// Holds the decoded transport headers plus the verified caller identity.
#[derive(Debug, Clone)]
pub struct RequestContext {
    /// All A2A standard headers (Authorization, Content-Type, A2A-Version,
    /// A2A-Extensions) decoded from the inbound message.
    pub headers: A2aHeaders,
    /// The verified caller identity. `None` is reserved for the explicit
    /// [`AllowAnonymous`] path; production authenticators always set this
    /// to `Some(...)` on success.
    pub caller: Option<Identity>,
}

impl RequestContext {
    /// Construct a request context with the supplied headers + caller.
    pub fn new(headers: A2aHeaders, caller: Option<Identity>) -> Self {
        Self { headers, caller }
    }

    /// Convenience: returns `true` iff a caller identity is present (any
    /// authenticator succeeded). Use this for handler-level gating on
    /// "authenticated only" methods.
    pub fn is_authenticated(&self) -> bool {
        self.caller.is_some()
    }
}

/// Failure cases for authentication.
#[derive(Debug, Error)]
pub enum AuthError {
    /// The `Authorization` header was not present on the request.
    #[error("missing Authorization header")]
    Missing,
    /// The header is present but does not match the expected format
    /// (e.g. wrong scheme prefix).
    #[error("malformed Authorization header")]
    Malformed,
    /// The verifier rejected the credential (bad signature, expired
    /// token, revoked, etc.). The string is included for log diagnostics
    /// and never returned to the caller.
    #[error("verification failed: {0}")]
    Rejected(String),
}

/// Server-side authentication port.
///
/// Implementations validate the inbound credential and return the verified
/// [`Identity`] on success. Failures map to JSON-RPC `-32001
/// Unauthenticated` at the server layer.
#[async_trait]
pub trait Authenticator: Send + Sync {
    /// Validate the request's credentials and return the verified caller
    /// identity. `headers` contains the parsed A2A header bag; `payload`
    /// is the raw JSON-RPC request body (some authenticators sign over
    /// the body — most do not).
    async fn authenticate(
        &self,
        headers: &A2aHeaders,
        payload: &[u8],
    ) -> Result<Identity, AuthError>;
}

/// **TEST FIXTURE / DEMO ONLY.** Allows every request through with
/// [`Identity::anonymous`].
///
/// Wiring this into a production [`A2aServer`](crate::server::A2aServer)
/// is a security bug — anyone publishable on the agent's subject can
/// invoke any handler method. Use [`BearerTokenAuthenticator`] (or your
/// own implementation) in production.
///
/// Gated behind the `test-fixtures` feature (W2.A14 / CWE-1188): default
/// builds cannot reach the type, so production wiring fails to compile
/// rather than silently accepting any caller. Conformance harnesses,
/// downstream test suites, and the demo `EchoHandler` enable the feature
/// explicitly.
#[cfg(any(feature = "test-fixtures", test))]
pub struct AllowAnonymous;

#[cfg(any(feature = "test-fixtures", test))]
#[async_trait]
impl Authenticator for AllowAnonymous {
    async fn authenticate(
        &self,
        _headers: &A2aHeaders,
        _payload: &[u8],
    ) -> Result<Identity, AuthError> {
        Ok(Identity::anonymous())
    }
}

/// Closure type for verifying a raw bearer token. Returns the verified
/// caller [`Identity`] on success.
pub type BearerVerifier = Box<dyn Fn(&str) -> Result<Identity, AuthError> + Send + Sync>;

/// Bearer-token authenticator: validates `Authorization: Bearer <token>`
/// against a caller-supplied verifier closure.
///
/// The closure receives the raw token string and returns a verified
/// [`Identity`] on success or [`AuthError::Rejected`] on failure. Wrap
/// your JWT, opaque-token, or Vault verification logic inside it — the
/// closure runs on every request, so prefer in-process caches over
/// round-tripping a remote service.
///
/// ```ignore
/// use klieo_a2a::auth::{BearerTokenAuthenticator, Identity, AuthError};
///
/// let authn = BearerTokenAuthenticator::new(|token: &str| {
///     // Verify the JWT signature, exp, aud — return Identity on success.
///     match my_jwt_verify(token) {
///         Ok(claims) => Ok(Identity::new(claims.sub)),
///         Err(e) => Err(AuthError::Rejected(e.to_string())),
///     }
/// });
/// ```
pub struct BearerTokenAuthenticator {
    verifier: BearerVerifier,
}

impl BearerTokenAuthenticator {
    /// Build a new bearer-token authenticator with the supplied verifier
    /// closure.
    pub fn new<F>(verifier: F) -> Self
    where
        F: Fn(&str) -> Result<Identity, AuthError> + Send + Sync + 'static,
    {
        Self {
            verifier: Box::new(verifier),
        }
    }
}

#[async_trait]
impl Authenticator for BearerTokenAuthenticator {
    async fn authenticate(
        &self,
        headers: &A2aHeaders,
        _payload: &[u8],
    ) -> Result<Identity, AuthError> {
        let header = headers.authorization.as_deref().ok_or(AuthError::Missing)?;
        let token = strip_bearer_prefix(header).ok_or(AuthError::Malformed)?;
        (self.verifier)(token)
    }
}

/// Strip the `Bearer ` scheme case-insensitively per RFC 7235 §2.1.
/// Returns the trimmed token or `None` if the header doesn't begin with
/// `bearer` (any casing) followed by whitespace + a non-empty value.
/// W2.A14 — old impl rejected `bearer ` / `BEARER ` / `Bearer\t`, which
/// is an interop bug rather than a security risk but masks legitimate
/// requests and pushes operators toward looser proxies.
fn strip_bearer_prefix(header: &str) -> Option<&str> {
    let (scheme, rest) = header.split_once(|c: char| c.is_ascii_whitespace())?;
    if !scheme.eq_ignore_ascii_case("bearer") {
        return None;
    }
    let token = rest.trim_start();
    if token.is_empty() {
        None
    } else {
        Some(token)
    }
}

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

    fn headers_with(auth: Option<&str>) -> A2aHeaders {
        let mut h = 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 id = authn
            .authenticate(&headers_with(None), 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 err = authn
            .authenticate(&headers_with(None), 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 err = authn
            .authenticate(&headers_with(Some("Basic abc")), 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"] {
            authn
                .authenticate(&headers_with(Some(header)), 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 err = authn
            .authenticate(&headers_with(Some("Bearer ")), 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_id = authn
            .authenticate(&headers_with(Some("Bearer good")), b"{}")
            .await
            .unwrap();
        assert_eq!(ok_id.as_str(), "alice");
        let err = authn
            .authenticate(&headers_with(Some("Bearer bad")), 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());
    }
}