nefarious-login/crates/nefarious-login/src/axum.rs

use std::fmt::Display;

use axum::extract::{FromRef, FromRequestParts, Query, State};
use axum::http::request::Parts;
use axum::http::{HeaderMap, StatusCode, Uri};
use axum::response::{ErrorResponse, IntoResponse, Redirect, Response};
use axum::routing::get;
use axum::{async_trait, Json, RequestPartsExt, Router};

use axum_extra::extract::CookieJar;
use axum_extra::headers::authorization::Basic;
use axum_extra::headers::{Authorization, Cookie};
use axum_extra::TypedHeader;
use serde::Deserialize;
use serde_json::json;

use crate::auth::{AuthService, COOKIE_APP_SESSION_NAME};
use crate::session::User;

#[derive(Debug, Deserialize)]
pub struct ZitadelAuthParams {
    return_url: Option<String>,
}

trait AnyhowExtensions<T, E>
where
    E: Display,
{
    fn into_response(self) -> Result<T, ErrorResponse>;
}
impl<T, E> AnyhowExtensions<T, E> for anyhow::Result<T, E>
where
    E: Display,
{
    fn into_response(self) -> Result<T, ErrorResponse> {
        match self {
            Ok(o) => Ok(o),
            Err(e) => {
                tracing::error!("failed with anyhow error: {}", e);
                Err(ErrorResponse::from((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    Json(json!({
                        "status": "something",
                    })),
                )))
            }
        }
    }
}

pub async fn zitadel_auth(
    State(auth_service): State<AuthService>,
) -> Result<impl IntoResponse, ErrorResponse> {
    let (headers, url) = auth_service.login(None).await.into_response()?;

    Ok((headers, Redirect::to(url.as_ref())))
}

#[derive(Debug, Deserialize)]
#[allow(dead_code)]
pub struct AuthRequest {
    code: String,
    state: String,
}

pub async fn login_authorized(
    Query(query): Query<AuthRequest>,
    State(auth_service): State<AuthService>,
    cookie_jar: CookieJar,
) -> Result<impl IntoResponse, ErrorResponse> {
    let cookie_value = cookie_jar
        .get(COOKIE_APP_SESSION_NAME)
        .map(|c| c.value().to_string());

    let (headers, url) = auth_service
        .login_authorized(&query.code, &query.state, cookie_value)
        .await
        .into_response()?;

    Ok((headers, Redirect::to(url.as_str())))
}

pub struct AuthController;

impl AuthController {
    pub async fn new_router(auth_service: AuthService) -> anyhow::Result<Router> {
        Ok(Router::new()
            .route("/zitadel", get(zitadel_auth))
            .route("/authorized", get(login_authorized))
            .with_state(auth_service))
    }
}

pub struct UserFromSession {
    pub user: User,
}

pub static COOKIE_NAME: &str = "SESSION";

pub struct AuthRedirect((HeaderMap, String));

impl IntoResponse for AuthRedirect {
    fn into_response(self) -> Response {
        (self.0 .0, Redirect::temporary(&self.0 .1.as_str())).into_response()
    }
}

#[async_trait]
impl<S> FromRequestParts<S> for UserFromSession
where
    AuthService: FromRef<S>,
    S: Send + Sync,
{
    type Rejection = AuthRedirect;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        let auth_service = AuthService::from_ref(state);

        let cookie: Option<TypedHeader<Cookie>> = parts.extract().await.unwrap();
        let session_cookie = cookie.as_ref().and_then(|cookie| cookie.get(COOKIE_NAME));
        if session_cookie.is_none() {
            let basic: Option<TypedHeader<Authorization<Basic>>> = parts.extract().await.unwrap();

            if let Some(basic) = basic {
                let token = match auth_service
                    .login_token(basic.username(), basic.password())
                    .await
                    .into_response()
                {
                    Ok(login) => login,
                    Err(e) => {
                        tracing::info!("did not find a basic login token, will trigger login");
                        let (headers, url) = auth_service
                            .login(Some(parts.uri.to_string()))
                            .await
                            .expect("to be able to request login");
                        return Err(AuthRedirect((headers, url.to_string())));
                    }
                };

                return Ok(UserFromSession {
                    user: User {
                        id: token.sub,
                        email: token.email,
                        name: token.name,
                    },
                });
            }

            tracing::info!("did not find a cookie, will trigger login");
            let (headers, url) = auth_service
                .login(Some(parts.uri.to_string()))
                .await
                .expect("to be able to request login");
            return Err(AuthRedirect((headers, url.to_string())));
        }

        let session_cookie = session_cookie.unwrap();

        // continue to decode the session cookie
        let user = match auth_service
            .get_user_from_session(session_cookie)
            .await
            .into_response()
        {
            Ok(user) => user,
            Err(_) => {
                tracing::info!("could not get user from session, will trigger login");
                let (headers, url) = auth_service
                    .login(Some(parts.uri.to_string()))
                    .await
                    .expect("to be able to request login");
                return Err(AuthRedirect((headers, url.to_string())));
            }
        };

        Ok(UserFromSession { user })
    }
}