como/crates/como_auth/src/auth.rs

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

use anyhow::Context;
use async_trait::async_trait;
use axum::http::{header::SET_COOKIE, HeaderMap};
use oauth2::url::Url;

use crate::{
    introspection::IntrospectionService,
    oauth::{OAuth, ZitadelConfig},
    session::{SessionService, User},
    AuthClap, AuthEngine,
};

#[async_trait]
pub trait Auth {
    async fn login(&self) -> anyhow::Result<Url>;
    async fn login_token(&self, user: &str, password: &str) -> anyhow::Result<String>;
    async fn login_authorized(&self, code: &str, state: &str) -> anyhow::Result<(HeaderMap, Url)>;
    async fn get_user_from_session(&self, cookie: &str) -> anyhow::Result<User>;
}

#[derive(Clone)]
pub struct AuthService(Arc<dyn Auth + Send + Sync + 'static>);

impl AuthService {
    pub async fn new(config: &AuthClap, session: SessionService) -> anyhow::Result<Self> {
        match config.engine {
            AuthEngine::Noop => Ok(Self::new_noop()),
            AuthEngine::Zitadel => {
                let oauth: OAuth = ZitadelConfig::try_from(config.zitadel.clone())?.into();
                let introspection: IntrospectionService =
                    IntrospectionService::new_zitadel(config).await?;

                Ok(Self::new_zitadel(oauth, introspection, session))
            }
        }
    }

    pub fn new_zitadel(
        oauth: OAuth,
        introspection: IntrospectionService,
        session: SessionService,
    ) -> Self {
        Self(Arc::new(ZitadelAuthService {
            oauth,
            introspection,
            session,
        }))
    }

    pub fn new_noop() -> Self {
        Self(Arc::new(NoopAuthService {}))
    }
}

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

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

pub struct ZitadelAuthService {
    oauth: OAuth,
    introspection: IntrospectionService,
    session: SessionService,
}
pub static COOKIE_NAME: &str = "SESSION";

#[async_trait]
impl Auth for ZitadelAuthService {
    async fn login(&self) -> anyhow::Result<Url> {
        let authorize_url = self.oauth.authorize_url().await?;

        Ok(authorize_url)
    }
    async fn login_authorized(&self, code: &str, _state: &str) -> anyhow::Result<(HeaderMap, Url)> {
        let token = self.oauth.exchange(code).await?;
        let user_id = self.introspection.get_id_token(token.as_str()).await?;
        let cookie_value = self.session.insert_user("user", user_id.as_str()).await?;

        let cookie = format!("{}={}; SameSite=Lax; Path=/", COOKIE_NAME, cookie_value);

        let mut headers = HeaderMap::new();
        headers.insert(SET_COOKIE, cookie.parse().unwrap());

        Ok((
            headers,
            Url::parse("http://localhost:3000/dash/home")
                .context("failed to parse login_authorized zitadel return url")?,
        ))
    }
    async fn login_token(&self, _user: &str, password: &str) -> anyhow::Result<String> {
        self.introspection.get_id_token(password).await
    }
    async fn get_user_from_session(&self, cookie: &str) -> anyhow::Result<User> {
        match self.session.get_user(cookie).await? {
            Some(u) => Ok(User { id: u }),
            None => Err(anyhow::anyhow!("failed to find user")),
        }
    }
}

pub struct NoopAuthService {}

#[async_trait]
impl Auth for NoopAuthService {
    async fn login(&self) -> anyhow::Result<Url> {
        todo!()
    }
    async fn login_authorized(
        &self,
        _code: &str,
        _state: &str,
    ) -> anyhow::Result<(HeaderMap, Url)> {
        todo!()
    }

    async fn login_token(&self, _user: &str, _password: &str) -> anyhow::Result<String> {
        todo!()
    }

    async fn get_user_from_session(&self, _cookie: &str) -> anyhow::Result<User> {
        todo!()
    }
}