slave-framework 0.11.0

//! Sign-in through an OpenID Connect provider.
//!
//! slave apps do not handle passwords or mint their own tokens. Users sign in
//! with an OIDC provider (Google, Microsoft Entra, Auth0, Okta, Zitadel,
//! Keycloak, Authentik, Logto, …); clients send the provider's **ID token** as
//! `Authorization: Bearer <id_token>`, and [`OidcVerifier`] checks it:
//! signature (the provider's JWKS), issuer, audience (one of your client IDs)
//! and expiry.
//!
//! Configure it with `OIDC_ISSUER_URL` and `OIDC_AUDIENCES` (the client IDs of
//! your apps, comma-separated).

use std::sync::Arc;
use std::time::{Duration, Instant};

use openidconnect::core::{
    CoreIdToken, CoreIdTokenVerifier, CoreJsonWebKeySet, CoreJwsSigningAlgorithm,
    CoreProviderMetadata,
};
use openidconnect::{ClientId, IssuerUrl, JsonWebKeySetUrl, Nonce};
use tokio::sync::RwLock;

use crate::{Error, Result};

/// Environment variable holding the provider's issuer URL.
pub const ISSUER_ENV: &str = "OIDC_ISSUER_URL";
/// Environment variable holding the accepted client IDs (comma-separated).
pub const AUDIENCES_ENV: &str = "OIDC_AUDIENCES";
/// Environment variable holding the client IDs of customer-facing apps
/// (storefronts): their tokens are accepted with [`Identity::customer`] set.
pub const CUSTOMER_AUDIENCES_ENV: &str = "OIDC_CUSTOMER_AUDIENCES";

/// How to configure sign-in, shown when settings are missing.
pub const SETUP_HINT: &str = "Register your apps with an OpenID Connect provider, then set \
    OIDC_ISSUER_URL (e.g. https://accounts.google.com) and OIDC_AUDIENCES (their client IDs, \
    comma-separated) in .env.";

/// Minimum time between JWKS refreshes triggered by unknown signing keys.
const JWKS_REFRESH_INTERVAL: Duration = Duration::from_secs(30);

/// Which provider to trust, and for which client IDs.
#[derive(Debug, Clone)]
pub struct OidcConfig {
    pub issuer: String,
    /// Client IDs whose ID tokens are accepted (the token's `aud`).
    pub audiences: Vec<String>,
    /// Client IDs of customer-facing apps, also accepted; see [`Identity::customer`].
    pub customer_audiences: Vec<String>,
}

impl OidcConfig {
    /// Reads `OIDC_ISSUER_URL`, `OIDC_AUDIENCES` and `OIDC_CUSTOMER_AUDIENCES`.
    pub fn from_env() -> Result<Self> {
        let issuer = std::env::var(ISSUER_ENV)
            .ok()
            .filter(|issuer| !issuer.trim().is_empty())
            .ok_or(Error::MissingOidcSetting(ISSUER_ENV))?;
        let list = |key: &str| -> Vec<String> {
            std::env::var(key)
                .unwrap_or_default()
                .split(',')
                .map(str::trim)
                .filter(|audience| !audience.is_empty())
                .map(str::to_owned)
                .collect()
        };
        let audiences = list(AUDIENCES_ENV);
        let customer_audiences = list(CUSTOMER_AUDIENCES_ENV);
        if audiences.is_empty() && customer_audiences.is_empty() {
            return Err(Error::MissingOidcSetting(AUDIENCES_ENV));
        }
        Ok(Self {
            issuer: issuer.trim().to_owned(),
            audiences,
            customer_audiences,
        })
    }
}

/// Who signed in, from a verified ID token.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Identity {
    /// The provider (`iss`); with `subject`, the stable identity of the user.
    pub issuer: String,
    /// The user's ID at the provider (`sub`).
    pub subject: String,
    pub email: Option<String>,
    pub email_verified: Option<bool>,
    pub name: Option<String>,
    /// Whether the token was issued to a customer-facing app (its `aud` is in
    /// `OIDC_CUSTOMER_AUDIENCES`): the backend grants it limited access.
    pub customer: bool,
}

/// Why a token was not accepted.
#[derive(Debug, thiserror::Error)]
pub enum VerifyError {
    /// Malformed, wrongly signed, expired, or for another issuer or audience.
    #[error("invalid ID token: {0}")]
    Invalid(String),
    /// The provider's keys could not be fetched.
    #[error("cannot reach the OpenID Connect provider: {0}")]
    Unavailable(String),
}

/// Verifies ID tokens from one provider. Cheap to clone.
#[derive(Clone)]
pub struct OidcVerifier {
    inner: Arc<Inner>,
}

struct Inner {
    issuer: IssuerUrl,
    /// Accepted client IDs, and whether each is customer-facing.
    audiences: Vec<(ClientId, bool)>,
    jwks_url: JsonWebKeySetUrl,
    http: openidconnect::reqwest::Client,
    jwks: RwLock<(CoreJsonWebKeySet, Instant)>,
}

impl std::fmt::Debug for OidcVerifier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OidcVerifier")
            .field("issuer", &self.inner.issuer.as_str())
            .field("audiences", &self.inner.audiences)
            .finish()
    }
}

/// Asymmetric algorithms accepted for ID tokens (never `none` or HMAC).
fn allowed_algs() -> [CoreJwsSigningAlgorithm; 8] {
    use CoreJwsSigningAlgorithm::*;
    [
        RsaSsaPkcs1V15Sha256,
        RsaSsaPkcs1V15Sha384,
        RsaSsaPkcs1V15Sha512,
        RsaSsaPssSha256,
        RsaSsaPssSha384,
        EcdsaP256Sha256,
        EcdsaP384Sha384,
        EdDsa,
    ]
}

impl OidcVerifier {
    /// Discovers the provider (`/.well-known/openid-configuration`) and
    /// fetches its signing keys.
    pub async fn discover(config: &OidcConfig) -> Result<Self> {
        let issuer = IssuerUrl::new(config.issuer.clone())
            .map_err(|err| Error::Oidc(format!("invalid {ISSUER_ENV}: {err}")))?;
        let http = openidconnect::reqwest::ClientBuilder::new()
            .redirect(openidconnect::reqwest::redirect::Policy::none())
            .build()
            .map_err(|err| Error::Oidc(format!("HTTP client: {err}")))?;
        let metadata = CoreProviderMetadata::discover_async(issuer.clone(), &http)
            .await
            .map_err(|err| {
                Error::Oidc(format!("discovery at {} failed: {err}", issuer.as_str()))
            })?;
        let jwks_url = metadata.jwks_uri().clone();
        let jwks = CoreJsonWebKeySet::fetch_async(&jwks_url, &http)
            .await
            .map_err(|err| Error::Oidc(format!("fetching signing keys failed: {err}")))?;
        Ok(Self {
            inner: Arc::new(Inner {
                issuer,
                // Customer audiences first: a client ID in both lists is limited.
                audiences: config
                    .customer_audiences
                    .iter()
                    .map(|id| (ClientId::new(id.clone()), true))
                    .chain(
                        config
                            .audiences
                            .iter()
                            .map(|id| (ClientId::new(id.clone()), false)),
                    )
                    .collect(),
                jwks_url,
                http,
                jwks: RwLock::new((jwks, Instant::now())),
            }),
        })
    }

    /// The provider's issuer URL.
    pub fn issuer(&self) -> &str {
        self.inner.issuer.as_str()
    }

    /// Checks `id_token` and returns who it identifies.
    pub async fn verify(&self, id_token: &str) -> std::result::Result<Identity, VerifyError> {
        let token: CoreIdToken = id_token
            .trim()
            .parse()
            .map_err(|err| VerifyError::Invalid(format!("malformed: {err}")))?;

        match self.check(&token).await {
            Ok(identity) => Ok(identity),
            // An unknown key usually means the provider rotated its keys.
            Err(err) if self.refresh_keys_if_stale().await? => {
                self.check(&token).await.map_err(|_| err)
            }
            Err(err) => Err(err),
        }
    }

    async fn check(&self, token: &CoreIdToken) -> std::result::Result<Identity, VerifyError> {
        let jwks = self.inner.jwks.read().await.0.clone();
        let mut last_error = None;
        for (audience, customer) in &self.inner.audiences {
            let verifier = CoreIdTokenVerifier::new_public_client(
                audience.clone(),
                self.inner.issuer.clone(),
                jwks.clone(),
            )
            .set_allowed_algs(allowed_algs());
            // Clients check the nonce of their own sign-in; the backend cannot.
            match token.claims(&verifier, |_: Option<&Nonce>| Ok(())) {
                Ok(claims) => {
                    return Ok(Identity {
                        issuer: claims.issuer().to_string(),
                        subject: claims.subject().to_string(),
                        email: claims.email().map(|email| email.to_string()),
                        email_verified: claims.email_verified(),
                        name: claims
                            .name()
                            .and_then(|name| name.get(None))
                            .map(|name| name.to_string()),
                        customer: *customer,
                    });
                }
                Err(err) => last_error = Some(err.to_string()),
            }
        }
        Err(VerifyError::Invalid(
            last_error.unwrap_or_else(|| "no audience configured".into()),
        ))
    }

    /// Refetches the JWKS unless it was fetched recently. Returns whether it did.
    async fn refresh_keys_if_stale(&self) -> std::result::Result<bool, VerifyError> {
        let mut jwks = self.inner.jwks.write().await;
        if jwks.1.elapsed() < JWKS_REFRESH_INTERVAL {
            return Ok(false);
        }
        let fresh = CoreJsonWebKeySet::fetch_async(&self.inner.jwks_url, &self.inner.http)
            .await
            .map_err(|err| VerifyError::Unavailable(err.to_string()))?;
        *jwks = (fresh, Instant::now());
        Ok(true)
    }
}

/// The sign-in side for apps that render pages on a server (e.g. an Axum
/// theme): the Authorization Code flow with PKCE and a nonce.
///
/// 1. [`authorize_url`](Self::authorize_url) gives the provider URL to
///    redirect the browser to, and a [`PendingSignIn`] to keep (for example
///    in a short-lived HttpOnly cookie).
/// 2. On the redirect back, [`exchange`](Self::exchange) checks the state,
///    swaps the code for tokens and verifies the ID token (signature, issuer,
///    audience, nonce). Send that ID token to the MCP backend as the bearer.
#[derive(Clone)]
pub struct OidcClient {
    metadata: CoreProviderMetadata,
    client_id: ClientId,
    client_secret: Option<openidconnect::ClientSecret>,
    http: openidconnect::reqwest::Client,
}

impl std::fmt::Debug for OidcClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OidcClient")
            .field("issuer", &self.metadata.issuer().as_str())
            .field("client_id", &self.client_id.as_str())
            .finish()
    }
}

/// What a sign-in in progress must remember between the redirect to the
/// provider and the redirect back. Store it where only this browser can read
/// it (an HttpOnly cookie); it is useless once the sign-in completes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PendingSignIn {
    pub state: String,
    pub pkce_verifier: String,
    pub nonce: String,
}

impl PendingSignIn {
    /// A compact form for a cookie value.
    pub fn encode(&self) -> String {
        format!("{}.{}.{}", self.state, self.pkce_verifier, self.nonce)
    }

    pub fn decode(value: &str) -> Option<Self> {
        let mut parts = value.split('.');
        let pending = Self {
            state: parts.next()?.to_owned(),
            pkce_verifier: parts.next()?.to_owned(),
            nonce: parts.next()?.to_owned(),
        };
        parts.next().is_none().then_some(pending)
    }
}

impl OidcClient {
    /// Discovers the provider for the app registered as `client_id`
    /// (`client_secret` for confidential clients, `None` for public ones).
    pub async fn discover(
        issuer: &str,
        client_id: &str,
        client_secret: Option<&str>,
    ) -> Result<Self> {
        let issuer = IssuerUrl::new(issuer.to_owned())
            .map_err(|err| Error::Oidc(format!("invalid issuer URL: {err}")))?;
        let http = openidconnect::reqwest::ClientBuilder::new()
            .redirect(openidconnect::reqwest::redirect::Policy::none())
            .build()
            .map_err(|err| Error::Oidc(format!("HTTP client: {err}")))?;
        let metadata = CoreProviderMetadata::discover_async(issuer.clone(), &http)
            .await
            .map_err(|err| {
                Error::Oidc(format!("discovery at {} failed: {err}", issuer.as_str()))
            })?;
        Ok(Self {
            metadata,
            client_id: ClientId::new(client_id.to_owned()),
            client_secret: client_secret
                .filter(|secret| !secret.is_empty())
                .map(|secret| openidconnect::ClientSecret::new(secret.to_owned())),
            http,
        })
    }

    /// The provider URL to send the browser to, and what to remember until it
    /// comes back to `redirect_url` (registered at the provider).
    pub fn authorize_url(&self, redirect_url: &str) -> Result<(String, PendingSignIn)> {
        self.authorize_url_with(redirect_url, &[])
    }

    /// Like [`Self::authorize_url`], with extra query parameters for the
    /// provider, e.g. `("prompt", "create")` to open its sign-up page.
    pub fn authorize_url_with(
        &self,
        redirect_url: &str,
        extra: &[(&str, &str)],
    ) -> Result<(String, PendingSignIn)> {
        let client = openidconnect::core::CoreClient::from_provider_metadata(
            self.metadata.clone(),
            self.client_id.clone(),
            self.client_secret.clone(),
        )
        .set_redirect_uri(
            openidconnect::RedirectUrl::new(redirect_url.to_owned())
                .map_err(|err| Error::Oidc(format!("invalid redirect URL: {err}")))?,
        );
        let (challenge, verifier) = openidconnect::PkceCodeChallenge::new_random_sha256();
        let mut request = client
            .authorize_url(
                openidconnect::core::CoreAuthenticationFlow::AuthorizationCode,
                openidconnect::CsrfToken::new_random,
                Nonce::new_random,
            )
            .add_scope(openidconnect::Scope::new("email".into()))
            .add_scope(openidconnect::Scope::new("profile".into()))
            .set_pkce_challenge(challenge);
        for (name, value) in extra {
            request = request.add_extra_param(*name, *value);
        }
        let (url, state, nonce) = request.url();
        Ok((
            url.to_string(),
            PendingSignIn {
                state: state.secret().clone(),
                pkce_verifier: verifier.secret().clone(),
                nonce: nonce.secret().clone(),
            },
        ))
    }

    /// Completes the sign-in started with `pending`: checks `state`, swaps
    /// `code` for tokens and returns the verified ID token.
    pub async fn exchange(
        &self,
        redirect_url: &str,
        code: &str,
        state: &str,
        pending: &PendingSignIn,
    ) -> Result<String> {
        if state != pending.state {
            return Err(Error::Oidc(
                "sign-in state does not match; start again".into(),
            ));
        }
        let client = openidconnect::core::CoreClient::from_provider_metadata(
            self.metadata.clone(),
            self.client_id.clone(),
            self.client_secret.clone(),
        )
        .set_redirect_uri(
            openidconnect::RedirectUrl::new(redirect_url.to_owned())
                .map_err(|err| Error::Oidc(format!("invalid redirect URL: {err}")))?,
        );
        let response = client
            .exchange_code(openidconnect::AuthorizationCode::new(code.to_owned()))
            .map_err(|err| Error::Oidc(format!("token endpoint: {err}")))?
            .set_pkce_verifier(openidconnect::PkceCodeVerifier::new(
                pending.pkce_verifier.clone(),
            ))
            .request_async(&self.http)
            .await
            .map_err(|err| Error::Oidc(format!("exchanging the code failed: {err}")))?;
        let id_token = openidconnect::TokenResponse::id_token(&response)
            .ok_or_else(|| Error::Oidc("the provider returned no ID token".into()))?;
        let verifier = client.id_token_verifier().set_allowed_algs(allowed_algs());
        id_token
            .claims(&verifier, &Nonce::new(pending.nonce.clone()))
            .map_err(|err| Error::Oidc(format!("invalid ID token: {err}")))?;
        Ok(id_token.to_string())
    }
}