use anyhow::Context; use async_session::{Session, SessionStore as _}; use axum::{ extract::{FromRequestParts, OriginalUri}, http::{request::Parts, Method}, response::{IntoResponse, Redirect, Response}, RequestPartsExt, }; use axum_extra::extract::{ cookie::{Cookie, SameSite}, CookieJar, }; use diesel::{ associations::Identifiable, deserialize::Queryable, dsl::{auto_type, insert_into, AsSelect, Eq, Select}, pg::Pg, prelude::*, Selectable, }; use uuid::Uuid; use crate::{ app_error::AppError, app_state::AppState, auth::{AuthInfo, SESSION_KEY_AUTH_INFO, SESSION_KEY_AUTH_REDIRECT}, schema::{team_memberships, teams, users}, sessions::AppSession, team_memberships::TeamMembership, teams::Team, }; #[derive(Clone, Debug, Identifiable, Insertable, Queryable, Selectable)] #[diesel(table_name = users)] #[diesel(check_for_backend(Pg))] pub struct User { pub id: Uuid, pub uid: String, pub email: String, } impl User { pub fn all() -> Select> { users::table.select(User::as_select()) } #[auto_type(no_type_alias)] pub fn with_uid(uid_value: &str) -> _ { users::uid.eq(uid_value) } #[auto_type(no_type_alias)] pub fn team_memberships(&self) -> _ { let user_id_filter: Eq = TeamMembership::with_user_id(&self.id); let select: AsSelect<(TeamMembership, Team), Pg> = <(TeamMembership, Team)>::as_select(); team_memberships::table .inner_join(teams::table) .filter(user_id_filter) .select(select) } } #[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.base_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.base_path), )); }; let db_conn = app_state.db_pool.get().await?; let current_user = db_conn .interact(move |conn| { let maybe_current_user = User::all() .filter(User::with_uid(&auth_info.sub)) .first(conn) .optional() .context("failed to load maybe_current_user")?; if let Some(current_user) = maybe_current_user { return Ok(current_user); } let new_user = User { id: Uuid::now_v7(), uid: auth_info.sub.clone(), email: auth_info.email, }; match insert_into(users::table) .values(new_user) .on_conflict(users::uid) .do_nothing() .returning(User::as_returning()) .get_result(conn) { QueryResult::Err(diesel::result::Error::NotFound) => { tracing::debug!("detected race to insert current user record"); User::all() .filter(User::with_uid(&auth_info.sub)) .first(conn) .context( "failed to load record after detecting race to insert current user", ) } QueryResult::Err(err) => { Err(err).context("failed to insert current user record") } QueryResult::Ok(result) => Ok(result), } }) .await .unwrap()?; 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() } } } }