phonograph/interim-server/src/users.rs

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 sqlx::{query_as, PgExecutor};
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 User {
    pub id: Uuid,
    pub uid: String,
    pub email: String,
}

impl User {
    pub async fn fetch_by_ids_any<'a, I: IntoIterator<Item = Uuid>, E: PgExecutor<'a>>(
        ids: I,
        app_db: E,
    ) -> Result<Vec<Self>, sqlx::Error> {
        let ids: Vec<Uuid> = ids.into_iter().collect();
        query_as!(
            Self,
            "
select * from users
where id = any($1)
",
            ids.as_slice()
        )
        .fetch_all(app_db)
        .await
        .map_err(Into::into)
    }
}

#[derive(Clone, Debug)]
pub struct CurrentUser(pub User);

impl<S> FromRequestParts<S> for CurrentUser
where
    S: Into<AppState> + Clone + Sync,
{
    type Rejection = CurrentUserRejection;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        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::<AuthInfo>(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<E> From<E> for CurrentUserRejection
where
    E: Into<AppError>,
{
    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()
            }
        }
    }
}