a2a-rs 0.8.2

//! Authentication adapter implementations

use std::collections::HashMap;
#[cfg(feature = "http-server")]
use std::sync::Arc;

use async_trait::async_trait;
#[cfg(feature = "http-server")]
use axum::{
    extract::State,
    http::{HeaderMap, Request, StatusCode},
    middleware::Next,
    response::Response,
};

#[cfg(not(feature = "http-server"))]
type HeaderMap = std::collections::HashMap<String, String>;

use crate::{
    domain::{A2AError, core::agent::SecurityScheme},
    port::authenticator::{AuthContext, AuthContextExtractor, AuthPrincipal, Authenticator},
};

/// HTTP Bearer token authenticator
#[derive(Clone)]
pub struct BearerTokenAuthenticator {
    /// The valid tokens
    tokens: Vec<String>,
    /// The security scheme configuration
    scheme: SecurityScheme,
}

impl BearerTokenAuthenticator {
    /// Create a new bearer token authenticator with the given tokens
    pub fn new(tokens: Vec<String>) -> Self {
        Self {
            tokens,
            scheme: SecurityScheme::http(
                "bearer".to_string(),
                None,
                Some("Bearer token authentication".to_string()),
            ),
        }
    }

    /// Create with a specific bearer format
    pub fn with_format(tokens: Vec<String>, format: String) -> Self {
        Self {
            tokens,
            scheme: SecurityScheme::http(
                "bearer".to_string(),
                Some(format),
                Some("Bearer token authentication".to_string()),
            ),
        }
    }
}

#[async_trait]
impl Authenticator for BearerTokenAuthenticator {
    async fn authenticate(&self, context: &AuthContext) -> Result<AuthPrincipal, A2AError> {
        self.validate_context(context)?;

        if self.tokens.contains(&context.credential) {
            Ok(AuthPrincipal::new(
                context.credential.clone(),
                "bearer".to_string(),
            ))
        } else {
            Err(A2AError::Internal(
                "Invalid authentication token".to_string(),
            ))
        }
    }

    fn security_scheme(&self) -> &SecurityScheme {
        &self.scheme
    }

    fn validate_context(&self, context: &AuthContext) -> Result<(), A2AError> {
        if context.scheme_type != "bearer" {
            return Err(A2AError::Internal(format!(
                "Invalid authentication scheme: expected 'bearer', got '{}'",
                context.scheme_type
            )));
        }
        Ok(())
    }
}

/// HTTP context extractor for Bearer tokens
#[derive(Clone)]
pub struct BearerTokenExtractor;

#[async_trait]
impl AuthContextExtractor for BearerTokenExtractor {
    #[cfg(feature = "http-server")]
    async fn extract_from_headers(&self, headers: &HeaderMap) -> Option<AuthContext> {
        headers
            .get(axum::http::header::AUTHORIZATION)
            .and_then(|h| h.to_str().ok())
            .and_then(|auth| {
                let parts: Vec<&str> = auth.splitn(2, ' ').collect();
                if parts.len() == 2 && parts[0].to_lowercase() == "bearer" {
                    Some(AuthContext::new("bearer".to_string(), parts[1].to_string()))
                } else {
                    None
                }
            })
    }

    #[cfg(not(feature = "http-server"))]
    async fn extract_from_headers(&self, headers: &HeaderMap) -> Option<AuthContext> {
        headers
            .get("authorization")
            .or_else(|| headers.get("Authorization"))
            .and_then(|auth| {
                let parts: Vec<&str> = auth.splitn(2, ' ').collect();
                if parts.len() == 2 && parts[0].to_lowercase() == "bearer" {
                    Some(AuthContext::new("bearer".to_string(), parts[1].to_string()))
                } else {
                    None
                }
            })
    }

    async fn extract_from_query(&self, _params: &HashMap<String, String>) -> Option<AuthContext> {
        // Bearer tokens are not typically passed in query parameters
        None
    }

    async fn extract_from_cookies(&self, _cookies: &str) -> Option<AuthContext> {
        // Bearer tokens are not typically passed in cookies
        None
    }
}

/// API Key authenticator
#[derive(Clone)]
pub struct ApiKeyAuthenticator {
    /// Valid API keys
    api_keys: Vec<String>,
    /// The security scheme configuration
    scheme: SecurityScheme,
}

impl ApiKeyAuthenticator {
    /// Create a new API key authenticator
    pub fn new(api_keys: Vec<String>, location: String, name: String) -> Self {
        Self {
            api_keys,
            scheme: SecurityScheme::api_key(
                name,
                location,
                Some("API key authentication".to_string()),
            ),
        }
    }

    /// Create for header-based API key
    pub fn header(api_keys: Vec<String>, header_name: String) -> Self {
        Self::new(api_keys, "header".to_string(), header_name)
    }

    /// Create for query parameter-based API key
    pub fn query(api_keys: Vec<String>, param_name: String) -> Self {
        Self::new(api_keys, "query".to_string(), param_name)
    }

    /// Create for cookie-based API key
    pub fn cookie(api_keys: Vec<String>, cookie_name: String) -> Self {
        Self::new(api_keys, "cookie".to_string(), cookie_name)
    }
}

#[async_trait]
impl Authenticator for ApiKeyAuthenticator {
    async fn authenticate(&self, context: &AuthContext) -> Result<AuthPrincipal, A2AError> {
        self.validate_context(context)?;

        if self.api_keys.contains(&context.credential) {
            Ok(
                AuthPrincipal::new(context.credential.clone(), "apikey".to_string())
                    .with_attribute(
                        "location".to_string(),
                        context
                            .metadata
                            .get("location")
                            .unwrap_or(&String::new())
                            .clone(),
                    ),
            )
        } else {
            Err(A2AError::Internal("Invalid API key".to_string()))
        }
    }

    fn security_scheme(&self) -> &SecurityScheme {
        &self.scheme
    }

    fn validate_context(&self, context: &AuthContext) -> Result<(), A2AError> {
        if context.scheme_type != "apikey" {
            return Err(A2AError::Internal(format!(
                "Invalid authentication scheme: expected 'apikey', got '{}'",
                context.scheme_type
            )));
        }
        Ok(())
    }
}

/// API Key context extractor
#[derive(Clone)]
pub struct ApiKeyExtractor {
    location: String,
    name: String,
}

impl ApiKeyExtractor {
    pub fn new(location: String, name: String) -> Self {
        Self { location, name }
    }
}

#[async_trait]
impl AuthContextExtractor for ApiKeyExtractor {
    #[cfg(feature = "http-server")]
    async fn extract_from_headers(&self, headers: &HeaderMap) -> Option<AuthContext> {
        if self.location != "header" {
            return None;
        }

        headers
            .get(axum::http::HeaderName::from_bytes(self.name.as_bytes()).ok()?)
            .and_then(|h| h.to_str().ok())
            .map(|value| {
                AuthContext::new("apikey".to_string(), value.to_string())
                    .with_metadata("location".to_string(), "header".to_string())
                    .with_metadata("name".to_string(), self.name.clone())
            })
    }

    #[cfg(not(feature = "http-server"))]
    async fn extract_from_headers(&self, headers: &HeaderMap) -> Option<AuthContext> {
        if self.location != "header" {
            return None;
        }

        headers.get(&self.name).map(|value| {
            AuthContext::new("apikey".to_string(), value.clone())
                .with_metadata("location".to_string(), "header".to_string())
                .with_metadata("name".to_string(), self.name.clone())
        })
    }

    async fn extract_from_query(&self, params: &HashMap<String, String>) -> Option<AuthContext> {
        if self.location != "query" {
            return None;
        }

        params.get(&self.name).map(|value| {
            AuthContext::new("apikey".to_string(), value.clone())
                .with_metadata("location".to_string(), "query".to_string())
                .with_metadata("name".to_string(), self.name.clone())
        })
    }

    async fn extract_from_cookies(&self, cookies: &str) -> Option<AuthContext> {
        if self.location != "cookie" {
            return None;
        }

        // Simple cookie parsing - in production, use a proper cookie parser
        cookies
            .split(';')
            .map(|cookie| cookie.trim())
            .find_map(|cookie| {
                let parts: Vec<&str> = cookie.splitn(2, '=').collect();
                if parts.len() == 2 && parts[0] == self.name {
                    Some(
                        AuthContext::new("apikey".to_string(), parts[1].to_string())
                            .with_metadata("location".to_string(), "cookie".to_string())
                            .with_metadata("name".to_string(), self.name.clone()),
                    )
                } else {
                    None
                }
            })
    }
}

/// No-op authenticator that allows all requests
#[derive(Clone)]
pub struct NoopAuthenticator {
    scheme: SecurityScheme,
}

impl NoopAuthenticator {
    /// Create a new no-op authenticator
    pub fn new() -> Self {
        Self {
            scheme: SecurityScheme::http(
                "none".to_string(),
                None,
                Some("No authentication required".to_string()),
            ),
        }
    }
}

impl Default for NoopAuthenticator {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Authenticator for NoopAuthenticator {
    async fn authenticate(&self, _context: &AuthContext) -> Result<AuthPrincipal, A2AError> {
        Ok(AuthPrincipal::new(
            "anonymous".to_string(),
            "none".to_string(),
        ))
    }

    fn security_scheme(&self) -> &SecurityScheme {
        &self.scheme
    }

    fn validate_context(&self, _context: &AuthContext) -> Result<(), A2AError> {
        // No validation needed for no-op
        Ok(())
    }
}

#[cfg(feature = "http-server")]
mod http_auth {
    use super::*;

    /// Authentication middleware state
    #[derive(Clone)]
    pub struct AuthState {
        /// The authenticator to use
        authenticator: Arc<dyn Authenticator>,
        /// Context extractors
        extractors: Vec<Arc<dyn AuthContextExtractor>>,
    }

    impl AuthState {
        /// Create a new authentication state, reading credentials the way the
        /// authenticator's own security scheme says they arrive.
        pub fn new(authenticator: impl Authenticator + 'static) -> Self {
            let extractors = extractors_for(authenticator.security_scheme());
            Self {
                authenticator: Arc::new(authenticator),
                extractors,
            }
        }

        /// Create with custom extractors, for a credential
        /// [`new`](Self::new) cannot derive.
        #[allow(dead_code)]
        pub fn with_extractors(
            authenticator: impl Authenticator + 'static,
            extractors: Vec<Arc<dyn AuthContextExtractor>>,
        ) -> Self {
            Self {
                authenticator: Arc::new(authenticator),
                extractors,
            }
        }
    }

    /// The extractor that reads the credential `scheme` describes.
    ///
    /// Derived rather than assumed, and that is the whole of it: every
    /// [`Authenticator`] declares its scheme — `security_scheme()` is required
    /// by the port, and the same value goes on the agent card — and every one
    /// of them refuses a context labelled with any other scheme. Hard-coding
    /// [`BearerTokenExtractor`] here therefore made three of the five
    /// authenticators unreachable: an API key, an OAuth2 access token and an
    /// OIDC ID token were each extracted as `bearer` and then refused by the
    /// authenticator that had just been wired up to accept them, so a server
    /// configured with one answered 401 to every request — the valid
    /// credentials included. Only `bearer` and JWT worked, because those are
    /// the two that expect what the bearer extractor produces.
    ///
    /// What the card advertises is now what the server reads.
    ///
    /// One limit worth stating: this middleware reads *headers*, so an API key
    /// configured with `location = "query"` or `"cookie"` extracts nothing and
    /// the request is refused. That is deliberate on the query side — a
    /// credential in a URL is a credential in every access log and referrer —
    /// and a caller who needs either can pass its own extractors to
    /// [`AuthState::with_extractors`].
    fn extractors_for(scheme: &SecurityScheme) -> Vec<Arc<dyn AuthContextExtractor>> {
        use crate::domain::core::agent::security_scheme::Scheme;

        match &scheme.scheme {
            Some(Scheme::ApiKeySecurityScheme(api_key)) => vec![Arc::new(ApiKeyExtractor::new(
                api_key.location.clone(),
                api_key.name.clone(),
            ))],
            // An OAuth2 access token and an OIDC ID token both arrive in an
            // `Authorization: Bearer` header and are both labelled `oauth2` —
            // which is what `OpenIdConnectAuthenticator::validate_context`
            // already accepts alongside its own name.
            #[cfg(feature = "auth")]
            Some(Scheme::Oauth2SecurityScheme(_) | Scheme::OpenIdConnectSecurityScheme(_)) => {
                vec![Arc::new(crate::adapter::auth::OAuth2Extractor)]
            }
            // `http` (bearer, and JWT in a bearer header), mTLS — where the
            // credential is not in a header this middleware reads at all — and
            // any scheme added later. The bearer extractor is what
            // `BearerTokenAuthenticator`, `JwtAuthenticator` and
            // `NoopAuthenticator` expect, so this is also what every working
            // deployment already had.
            _ => vec![Arc::new(BearerTokenExtractor)],
        }
    }

    /// Authentication middleware for Axum.
    ///
    /// On success the [`AuthPrincipal`] is inserted into the request extensions,
    /// which is how it reaches the transport adapter and from there the message
    /// handler. Dropping it here is what made every caller look identical to a
    /// handler keeping per-caller state.
    pub async fn http_auth_middleware(
        State(state): State<AuthState>,
        mut req: Request<axum::body::Body>,
        next: Next,
    ) -> Result<Response, StatusCode> {
        // The first extractor that finds credentials decides the request; the
        // rest are not consulted. Resolved before touching the extensions so the
        // borrow of the headers is done with by then.
        let mut outcome = None;
        for extractor in &state.extractors {
            if let Some(context) = extractor.extract_from_headers(req.headers()).await {
                outcome = Some(state.authenticator.authenticate(&context).await);
                break;
            }
        }

        match outcome {
            Some(Ok(principal)) => {
                req.extensions_mut().insert(principal);
                Ok(next.run(req).await)
            }
            // Credentials were presented and rejected, or none were presented at
            // all. Both are 401.
            Some(Err(_)) | None => Err(StatusCode::UNAUTHORIZED),
        }
    }

    /// Helper function to apply authentication middleware to a router
    pub fn with_auth<R>(router: R, authenticator: impl Authenticator + 'static) -> axum::Router
    where
        R: Into<axum::Router>,
    {
        let auth_state = AuthState::new(authenticator);
        let router = router.into();

        router.layer(axum::middleware::from_fn_with_state(
            auth_state,
            http_auth_middleware,
        ))
    }
}

#[cfg(feature = "http-server")]
pub use http_auth::with_auth;