magnetar/src/web/auth.rs

use crate::service::local_user_cache::UserCacheError;
use crate::service::MagnetarService;
use crate::web::{ApiError, IntoErrorCode};
use axum::async_trait;
use axum::extract::rejection::ExtensionRejection;
use axum::extract::{FromRequestParts, Request, State};
use axum::http::request::Parts;
use axum::http::{HeaderMap, StatusCode};
use axum::middleware::Next;
use axum::response::{IntoResponse, Response};
use headers::authorization::Bearer;
use headers::{Authorization, HeaderMapExt};
use magnetar_model::{ck, CalckeyDbError};
use std::convert::Infallible;
use std::sync::Arc;
use strum::IntoStaticStr;
use thiserror::Error;
use tracing::error;

#[derive(Clone, Debug)]
pub enum AuthMode {
    User {
        user: Arc<ck::user::Model>,
    },
    AccessToken {
        user: Arc<ck::user::Model>,
        access_token: Arc<ck::access_token::Model>,
    },
    Anonymous,
}

impl AuthMode {
    fn get_user(&self) -> Option<&Arc<ck::user::Model>> {
        match self {
            AuthMode::User { user } | AuthMode::AccessToken { user, .. } => Some(user),
            AuthMode::Anonymous => None,
        }
    }
}

pub struct AuthUserRejection(ApiError);

impl From<ExtensionRejection> for AuthUserRejection {
    fn from(rejection: ExtensionRejection) -> Self {
        AuthUserRejection(ApiError {
            status: StatusCode::UNAUTHORIZED,
            code: "Unauthorized".error_code(),
            message: if cfg!(debug_assertions) {
                format!("Missing auth extension: {}", rejection)
            } else {
                "Unauthorized".to_string()
            },
        })
    }
}

impl IntoResponse for AuthUserRejection {
    fn into_response(self) -> Response {
        self.0.into_response()
    }
}

#[derive(Clone, FromRequestParts)]
#[from_request(via(axum::Extension), rejection(AuthUserRejection))]
pub struct AuthenticatedUser(pub Arc<ck::user::Model>);

#[derive(Clone)]
pub struct MaybeUser(pub Option<Arc<ck::user::Model>>);

#[async_trait]
impl<S> FromRequestParts<S> for MaybeUser {
    type Rejection = Infallible;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        Ok(MaybeUser(
            parts
                .extensions
                .get::<AuthenticatedUser>()
                .map(|part| part.0.clone()),
        ))
    }
}

#[derive(Clone)]
pub struct AuthState {
    service: Arc<MagnetarService>,
}

#[derive(Debug, Error, IntoStaticStr)]
enum AuthError {
    #[error("Unsupported authorization scheme")]
    UnsupportedScheme,
    #[error("Cache error: {0}")]
    CacheError(#[from] UserCacheError),
    #[error("Database error: {0}")]
    DbError(#[from] CalckeyDbError),
    #[error("Invalid token")]
    InvalidToken,
    #[error("Invalid token \"{token}\" referencing user \"{user}\"")]
    InvalidTokenUser { token: String, user: String },
    #[error("Invalid access token \"{access_token}\" referencing app \"{app}\"")]
    InvalidAccessTokenApp { access_token: String, app: String },
}

impl From<AuthError> for ApiError {
    fn from(err: AuthError) -> Self {
        match err {
            AuthError::UnsupportedScheme => ApiError {
                status: StatusCode::UNAUTHORIZED,
                code: err.error_code(),
                message: "Unsupported authorization scheme".to_string(),
            },
            AuthError::CacheError(err) => err.into(),
            AuthError::DbError(err) => err.into(),
            AuthError::InvalidTokenUser {
                ref token,
                ref user,
            } => {
                error!("Invalid token \"{}\" referencing user \"{}\"", token, user);

                ApiError {
                    status: StatusCode::INTERNAL_SERVER_ERROR,
                    code: err.error_code(),
                    message: if cfg!(debug_assertions) {
                        format!("Invalid token \"{}\" referencing user \"{}\"", token, user)
                    } else {
                        "Invalid token-user link".to_string()
                    },
                }
            }
            AuthError::InvalidAccessTokenApp {
                ref access_token,
                ref app,
            } => {
                error!(
                    "Invalid access token \"{}\" referencing app \"{}\"",
                    access_token, app
                );

                ApiError {
                    status: StatusCode::INTERNAL_SERVER_ERROR,
                    code: err.error_code(),
                    message: if cfg!(debug_assertions) {
                        format!(
                            "Invalid access token \"{}\" referencing app \"{}\"",
                            access_token, app
                        )
                    } else {
                        "Invalid access token-app link".to_string()
                    },
                }
            }
            AuthError::InvalidToken => ApiError {
                status: StatusCode::UNAUTHORIZED,
                code: err.error_code(),
                message: "Invalid token".to_string(),
            },
        }
    }
}

pub fn is_user_token(token: &str) -> bool {
    token.chars().count() == 16
}

impl AuthState {
    pub fn new(magnetar: Arc<MagnetarService>) -> Self {
        Self { service: magnetar }
    }

    async fn authorize_token(
        &self,
        Authorization(token): &Authorization<Bearer>,
    ) -> Result<AuthMode, AuthError> {
        let token = token.token();

        if is_user_token(token) {
            let user_cache = &self.service.local_user_cache;
            let user = user_cache.get_by_token(token).await?;

            if let Some(user) = user {
                return Ok(AuthMode::User { user });
            }

            Err(AuthError::InvalidToken)
        } else {
            let access_token = self.service.db.get_access_token(token).await?;

            if access_token.is_none() {
                return Err(AuthError::InvalidToken);
            }

            let access_token = access_token.unwrap();

            let user = self
                .service
                .local_user_cache
                .get_by_id(&access_token.user_id)
                .await?;

            if user.is_none() {
                return Err(AuthError::InvalidTokenUser {
                    token: access_token.id,
                    user: access_token.user_id,
                });
            }

            let user = user.unwrap();

            if let Some(app_id) = &access_token.app_id {
                return match self.service.db.get_app_by_id(app_id).await? {
                    Some(app) => Ok(AuthMode::AccessToken {
                        user,
                        access_token: Arc::new(ck::access_token::Model {
                            permission: app.permission,
                            ..access_token
                        }),
                    }),
                    None => Err(AuthError::InvalidAccessTokenApp {
                        access_token: access_token.id,
                        app: access_token.user_id,
                    }),
                };
            }

            let access_token = Arc::new(access_token);

            Ok(AuthMode::AccessToken { user, access_token })
        }
    }
}

pub async fn auth(
    State(state): State<AuthState>,
    header_map: HeaderMap,
    mut req: Request,
    next: Next,
) -> Result<Response, ApiError> {
    let auth_bearer = match header_map.typed_try_get::<Authorization<Bearer>>() {
        Ok(Some(auth)) => auth,
        Ok(None) => {
            req.extensions_mut().insert(AuthMode::Anonymous);
            return Ok(next.run(req).await);
        }
        Err(_) => {
            return Err(AuthError::UnsupportedScheme.into());
        }
    };

    match state.authorize_token(&auth_bearer).await {
        Ok(auth) => {
            if let Some(user) = auth.get_user() {
                let user = AuthenticatedUser(user.clone());
                req.extensions_mut().insert(user);
            }

            req.extensions_mut().insert(auth);

            Ok(next.run(req).await)
        }
        Err(e) => Err(e.into()),
    }
}