como/como_auth/src/introspection.rs

use std::{ops::Deref, sync::Arc};

use async_trait::async_trait;
use axum::extract::FromRef;
use oauth2::TokenIntrospectionResponse;
use openidconnect::IntrospectionUrl;
use zitadel::{
    axum::introspection::IntrospectionStateBuilderError,
    credentials::Application,
    oidc::{
        discovery::discover,
        introspection::{introspect, AuthorityAuthentication},
    },
};

use crate::AuthClap;

#[async_trait]
pub trait Introspection {
    async fn get_user(&self) -> anyhow::Result<()>;
    async fn get_id_token(&self, token: &str) -> anyhow::Result<String>;
}

pub struct IntrospectionService(Arc<dyn Introspection + Send + Sync + 'static>);
impl IntrospectionService {
    pub async fn new_zitadel(config: &AuthClap) -> anyhow::Result<Self> {
        let res = IntrospectionStateBuilder::new(&config.zitadel.authority_url.clone().unwrap())
            .with_basic_auth(
                &config.zitadel.client_id.clone().unwrap(),
                &config.zitadel.client_secret.clone().unwrap(),
            )
            .build()
            .await?;

        Ok(IntrospectionService(Arc::new(ZitadelIntrospection::new(
            res,
        ))))
    }
}

impl Deref for IntrospectionService {
    type Target = Arc<dyn Introspection + Send + Sync + 'static>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

pub struct ZitadelIntrospection {
    state: IntrospectionState,
}

impl ZitadelIntrospection {
    pub fn new(state: IntrospectionState) -> Self {
        Self { state }
    }
}

#[async_trait]
impl Introspection for ZitadelIntrospection {
    async fn get_user(&self) -> anyhow::Result<()> {
        Ok(())
    }
    async fn get_id_token(&self, token: &str) -> anyhow::Result<String> {
        let config = &self.state.config;
        let res = introspect(
            &config.introspection_uri,
            &config.authority,
            &config.authentication,
            token,
        )
        .await?;

        Ok(res
            .sub()
            .ok_or(anyhow::anyhow!("could not find a userid (sub) in token"))?
            .to_string())
    }
}

#[derive(Clone, Debug)]
pub struct IntrospectionState {
    pub(crate) config: IntrospectionConfig,
}

/// Configuration that must be inject into the axum application state. Used by the
/// [IntrospectionStateBuilder](super::IntrospectionStateBuilder). This struct is also used to create the [IntrospectionState](IntrospectionState)
#[derive(Debug, Clone)]
pub struct IntrospectionConfig {
    pub authority: String,
    pub authentication: AuthorityAuthentication,
    pub introspection_uri: IntrospectionUrl,
}

impl FromRef<IntrospectionState> for IntrospectionConfig {
    fn from_ref(input: &IntrospectionState) -> Self {
        input.config.clone()
    }
}

pub struct IntrospectionStateBuilder {
    authority: String,
    authentication: Option<AuthorityAuthentication>,
}

/// Builder for [IntrospectionConfig]
impl IntrospectionStateBuilder {
    pub fn new(authority: &str) -> Self {
        Self {
            authority: authority.to_string(),
            authentication: None,
        }
    }

    pub fn with_basic_auth(
        &mut self,
        client_id: &str,
        client_secret: &str,
    ) -> &mut IntrospectionStateBuilder {
        self.authentication = Some(AuthorityAuthentication::Basic {
            client_id: client_id.to_string(),
            client_secret: client_secret.to_string(),
        });

        self
    }

    pub fn with_jwt_profile(&mut self, application: Application) -> &mut IntrospectionStateBuilder {
        self.authentication = Some(AuthorityAuthentication::JWTProfile { application });

        self
    }

    pub async fn build(&mut self) -> Result<IntrospectionState, IntrospectionStateBuilderError> {
        let authentication = self
            .authentication
            .clone()
            .ok_or(IntrospectionStateBuilderError::NoAuthSchema)?;

        let metadata = discover(&self.authority)
            .await
            .map_err(|source| IntrospectionStateBuilderError::Discovery { source })?;

        let introspection_uri = metadata
            .additional_metadata()
            .introspection_endpoint
            .clone()
            .ok_or(IntrospectionStateBuilderError::NoIntrospectionUrl)?;

        Ok(IntrospectionState {
            config: IntrospectionConfig {
                authority: self.authority.clone(),
                introspection_uri: introspection_uri,
                authentication: authentication,
            },
        })
    }
}