use async_session::{Session, SessionStore as _}; use axum::{ RequestPartsExt, extract::{FromRequestParts, OriginalUri}, http::{Method, request::Parts}, response::{IntoResponse, Redirect, Response}, }; use axum_extra::extract::{ CookieJar, cookie::{Cookie, SameSite}, }; use interim_models::user::User; use sqlx::query_as; use uuid::Uuid; use crate::{ app_error::AppError, app_state::AppState, auth::{AuthInfo, SESSION_KEY_AUTH_INFO, SESSION_KEY_AUTH_REDIRECT}, sessions::AppSession, }; #[derive(Clone, Debug)] pub struct CurrentUser(pub User); impl FromRequestParts for CurrentUser where S: Into + Clone + Sync, { type Rejection = CurrentUserRejection; async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { let app_state: AppState = state.clone().into(); let mut session = if let AppSession(Some(value)) = parts.extract_with_state(&app_state).await? { value } else { Session::new() }; let auth_info = if let Some(value) = session.get::(SESSION_KEY_AUTH_INFO) { value } else { let jar: CookieJar = parts.extract().await?; let method: Method = parts.extract().await?; let jar = if method == Method::GET { let OriginalUri(uri) = parts.extract().await?; session.insert( SESSION_KEY_AUTH_REDIRECT, uri.path_and_query() .map(|value| value.to_string()) .unwrap_or(format!("{}/", app_state.settings.root_path)), )?; if let Some(cookie_value) = app_state.session_store.store_session(session).await? { tracing::debug!("adding session cookie to jar"); jar.add( Cookie::build((app_state.settings.auth.cookie_name.clone(), cookie_value)) .same_site(SameSite::Lax) .http_only(true) .path("/"), ) } else { tracing::debug!("inferred that session cookie already in jar"); jar } } else { // If request method is not GET then do not attempt to infer the // redirect target, as there may be no GET handler defined for // it. jar }; return Err(Self::Rejection::SetCookiesAndRedirect( jar, format!("{}/auth/login", app_state.settings.root_path), )); }; let current_user = if let Some(value) = query_as!(User, "select * from users where uid = $1", &auth_info.sub) .fetch_optional(&app_state.app_db) .await? { value } else if let Some(value) = query_as!( User, " insert into users (id, uid, email) values ($1, $2, $3) on conflict (uid) do nothing returning * ", Uuid::now_v7(), &auth_info.sub, &auth_info.email ) .fetch_optional(&app_state.app_db) .await? { value } else { tracing::debug!("detected race to insert current user record"); query_as!(User, "select * from users where uid = $1", &auth_info.sub) .fetch_one(&app_state.app_db) .await? }; Ok(CurrentUser(current_user)) } } pub enum CurrentUserRejection { AppError(AppError), SetCookiesAndRedirect(CookieJar, String), } // Easily convert semi-arbitrary errors to InternalServerError impl From for CurrentUserRejection where E: Into, { fn from(err: E) -> Self { Self::AppError(err.into()) } } impl IntoResponse for CurrentUserRejection { fn into_response(self) -> Response { match self { Self::AppError(err) => err.into_response(), Self::SetCookiesAndRedirect(jar, redirect_to) => { (jar, Redirect::to(&redirect_to)).into_response() } } } }