zero-to-axum/src/server/routes/auth/mod.rs

use std::fmt;

use anyhow::Context as _;
use argon2::Argon2;
use askama::Template;
use askama_web::WebTemplate;
use axum::{
    extract::{Query, State},
    http::StatusCode,
    response::{IntoResponse, Redirect},
    routing::get,
    Form, Router,
};
use password_hash::{rand_core::OsRng, PasswordHash, PasswordHasher, PasswordVerifier, SaltString};
use rand::{distr::Alphanumeric, rng, Rng as _};
use serde::Deserialize;
use tracing::{debug, error, info};
use uuid::Uuid;

use crate::server::{
    session::{Auth, CsrfCookie, User},
    AppState,
};

use super::ErrorPage;

pub fn build() -> Router<AppState> {
    Router::new()
        .route("/signup", get(signup_page).post(signup))
        .route("/login", get(login_page).post(login))
        .route("/logout", get(logout_page).post(logout))
        .route("/confirm", get(signup_confirm_page).post(signup_confirm))
}

#[derive(Template, WebTemplate, Default)]
#[template(path = "signup.html")]
struct SignupPage {
    error: Option<String>,
    csrf_token: String,
}

#[tracing::instrument]
pub async fn signup_page(csrf: CsrfCookie) -> impl IntoResponse {
    info!("get signup page");

    let csrf_token = csrf.token().to_string();

    (
        csrf,
        SignupPage {
            error: None,
            csrf_token,
        },
    )
}

#[derive(Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct SignupForm {
    email: String,
    password: String,
    csrf_token: String,
}

impl fmt::Debug for SignupForm {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SignupForm")
            .field("email", &self.email)
            .field("password", &"REDACTED")
            .field("csrf-token", &self.csrf_token)
            .finish()
    }
}

#[derive(Template, WebTemplate, Default)]
#[template(path = "signup-confirmation.html")]
struct SignupConfirmation {
    email: String,
}

#[tracing::instrument(skip(db, email_client))]
pub async fn signup(
    State(AppState {
        db,
        email_client,
        conf,
        ..
    }): State<AppState>,
    csrf: CsrfCookie,
    Form(form): Form<SignupForm>,
) -> Result<impl IntoResponse, SignupError> {
    info!("signup attempt");

    if form.csrf_token != csrf.token() {
        return Err(SignupError::CsrfValidationFailed);
    }

    info!("hash password: {}", &form.password);
    let password_hash: String = tokio::task::spawn_blocking(async move || {
        let salt = SaltString::generate(&mut OsRng);
        let argon2 = Argon2::default();
        Ok::<String, anyhow::Error>(
            argon2
                .hash_password(form.password.as_bytes(), &salt)?
                .to_string(),
        )
    })
    .await
    .context("run password hashing operation")?
    .await?;
    info!("hashed password: {password_hash}");

    let mut txn = db.begin().await.context("start database transaction")?;

    let user_id = Uuid::new_v4();
    sqlx::query!(
        r#"
            INSERT INTO users (id, email, password, status)
            VALUES ($1, $2, $3, 'pending_confirmation')
        "#,
        user_id,
        form.email,
        password_hash,
    )
    .execute(&mut *txn)
    .await
    .context("insert new user into database")?;

    let token: String = {
        let mut rng = rng();
        std::iter::repeat_with(|| rng.sample(Alphanumeric))
            .map(char::from)
            .take(25)
            .collect()
    };

    sqlx::query!(
        r#"
            INSERT INTO signup_tokens (user_id, token)
            VALUES ($1, $2)
        "#,
        user_id,
        token
    )
    .execute(&mut *txn)
    .await
    .context("insert subscription token into database")?;

    email_client
        .send_email(
            form.email.parse().map_err(|_| SignupError::InvalidEmail)?,
            "Test".to_owned(),
            format!(
                "Please confirm your subscription: {}auth/confirm?token={token}",
                conf.app.public_url
            ),
        )
        .await?;

    debug!("email sent");

    txn.commit().await.context("commit transaction")?;

    Ok(SignupConfirmation { email: form.email })
}

#[derive(thiserror::Error, Debug)]
pub enum SignupError {
    #[error("Invalid Email")]
    InvalidEmail,
    #[error("CSRF Validation Failed")]
    CsrfValidationFailed,
    #[error("Unknown Error: {0}")]
    Unknown(#[from] anyhow::Error),
}

impl IntoResponse for SignupError {
    fn into_response(self) -> axum::response::Response {
        let (status, message) = match self {
            SignupError::InvalidEmail => (StatusCode::BAD_REQUEST, "Invalid Email"),
            SignupError::CsrfValidationFailed => {
                (StatusCode::BAD_REQUEST, "CSRF Validation Failed, Try Again")
            }
            SignupError::Unknown(e) => {
                error!(?e, "returning INTERNAL SERVER ERROR");
                (StatusCode::INTERNAL_SERVER_ERROR, "Unknown Error")
            }
        };

        (
            status,
            ErrorPage {
                error: message.to_string(),
            },
        )
            .into_response()
    }
}

#[derive(Template, WebTemplate, Default)]
#[template(path = "login.html")]
struct LoginPage {
    error: Option<String>,
    csrf_token: String,
}

#[tracing::instrument]
pub async fn login_page(csrf: CsrfCookie) -> impl IntoResponse {
    info!("get login page");

    let csrf_token = csrf.token().to_string();

    (
        csrf,
        LoginPage {
            error: None,
            csrf_token,
        },
    )
}

#[derive(Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct LoginForm {
    email: String,
    password: String,
    csrf_token: String,
}

impl fmt::Debug for LoginForm {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("LoginForm")
            .field("email", &self.email)
            .field("password", &"REDACTED")
            .finish()
    }
}

#[tracing::instrument(skip(db, auth))]
pub async fn login(
    State(AppState { db, .. }): State<AppState>,
    mut auth: Auth,
    csrf: CsrfCookie,
    Form(form): Form<LoginForm>,
) -> Result<impl IntoResponse, LoginError> {
    info!("login attempt");

    if form.csrf_token != csrf.token() {
        return Err(LoginError::CsrfValidationFailed);
    }

    let user = sqlx::query!(
        r#"
            SELECT id, password FROM users WHERE status = 'confirmed' AND email = $1 LIMIT 1;
        "#,
        form.email
    )
    .fetch_one(&db)
    .await;

    if matches!(user, Err(sqlx::Error::RowNotFound)) {
        return Err(LoginError::UnknownUser);
    }

    let user = user.context("get user info from db")?;

    tokio::task::spawn_blocking(async move || {
        let parsed_hash =
            PasswordHash::new(&user.password).context("parse stored password hash")?;

        match Argon2::default().verify_password(form.password.as_bytes(), &parsed_hash) {
            Ok(()) => Ok(()),
            Err(password_hash::Error::Password) => Err(LoginError::InvalidPassword),
            Err(e) => Err(e).context("verify password hash")?,
        }
    })
    .await
    .context("spawn password verifier task")?
    .await?;

    auth.user = Some(User::new(user.id));
    auth.session
        .cycle_id()
        .await
        .context("refresh session id")?;
    auth.session
        .insert("user", &auth.user)
        .await
        .context("set user data in session")?;

    Ok(Redirect::to("/"))
}

#[derive(thiserror::Error, Debug)]
pub enum LoginError {
    #[error("Invalid Password")]
    InvalidPassword,
    #[error("Unknown User")]
    UnknownUser,
    #[error("CSRF Validation Failed")]
    CsrfValidationFailed,
    #[error("Unknown Error: {0}")]
    Unknown(#[from] anyhow::Error),
}

impl IntoResponse for LoginError {
    fn into_response(self) -> axum::response::Response {
        let (status, message) = match self {
            LoginError::InvalidPassword => (StatusCode::UNAUTHORIZED, "Invalid Password"),
            LoginError::UnknownUser => (StatusCode::UNAUTHORIZED, "Unknown User"),
            LoginError::CsrfValidationFailed => {
                (StatusCode::BAD_REQUEST, "CSRF Validation Failed, Try Again")
            }
            LoginError::Unknown(e) => {
                error!(?e, "returning INTERNAL SERVER ERROR");
                (StatusCode::INTERNAL_SERVER_ERROR, "Unknown Error")
            }
        };

        (
            status,
            ErrorPage {
                error: message.to_string(),
            },
        )
            .into_response()
    }
}

#[derive(Template, WebTemplate)]
#[template(path = "logout.html")]
struct LogoutPage {
    csrf_token: String,
}

#[tracing::instrument]
pub async fn logout_page(csrf: CsrfCookie) -> impl IntoResponse {
    info!("get logout page");

    let csrf_token = csrf.token().to_string();

    (csrf, LogoutPage { csrf_token })
}

#[derive(Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct LogoutForm {
    csrf_token: String,
}

pub async fn logout(
    mut user: Auth,
    csrf: CsrfCookie,
    Form(form): Form<LogoutForm>,
) -> Result<impl IntoResponse, LogoutError> {
    info!("logout attempt");

    if form.csrf_token != csrf.token() {
        return Err(LogoutError::CsrfValidationFailed);
    }

    if user.user.is_none() {
        return Err(LogoutError::NotLoggedIn);
    }

    user.user = None;
    user.session.flush().await.context("flush user session")?;

    Ok(Redirect::to("/"))
}

#[derive(thiserror::Error, Debug)]
pub enum LogoutError {
    #[error("Not Logged In")]
    NotLoggedIn,
    #[error("CSRF Validation Failed")]
    CsrfValidationFailed,
    #[error("Unknown Error: {0}")]
    Unknown(#[from] anyhow::Error),
}

impl IntoResponse for LogoutError {
    fn into_response(self) -> axum::response::Response {
        let (status, message) = match self {
            LogoutError::NotLoggedIn => (StatusCode::UNAUTHORIZED, "Unknown User"),
            LogoutError::CsrfValidationFailed => {
                (StatusCode::BAD_REQUEST, "CSRF Validation Failed, Try Again")
            }
            LogoutError::Unknown(e) => {
                error!(?e, "returning INTERNAL SERVER ERROR");
                (StatusCode::INTERNAL_SERVER_ERROR, "Unknown Error")
            }
        };

        (
            status,
            ErrorPage {
                error: message.to_string(),
            },
        )
            .into_response()
    }
}

#[derive(Template, WebTemplate)]
#[template(path = "signup-confirm.html")]
struct SignupConfirmPage {
    csrf_token: String,
}

#[tracing::instrument]
pub async fn signup_confirm_page(csrf: CsrfCookie) -> impl IntoResponse {
    info!("get signup confirm page");

    let csrf_token = csrf.token().to_string();

    (csrf, SignupConfirmPage { csrf_token })
}

#[derive(Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct SignupConfirmForm {
    csrf_token: String,
}

#[derive(Deserialize)]
pub struct SignupConfirmQuery {
    token: String,
}

pub async fn signup_confirm(
    State(AppState { db, .. }): State<AppState>,
    csrf: CsrfCookie,
    Query(query): Query<SignupConfirmQuery>,
    Form(form): Form<SignupConfirmForm>,
) -> Result<impl IntoResponse, SignupConfirmError> {
    info!("signup confirm attempt");

    if form.csrf_token != csrf.token() {
        return Err(SignupConfirmError::CsrfValidationFailed);
    }

    let rows = sqlx::query!(
        r#"
            UPDATE users SET status = 'confirmed' FROM signup_tokens
            WHERE status = 'pending_confirmation' AND signup_tokens.user_id = id AND signup_tokens.token = $1;
        "#,
        query.token
    )
    .execute(&db)
    .await
    .context("set user to confirmed")?.rows_affected();

    if rows < 1 {
        return Err(SignupConfirmError::InvalidToken);
    }

    Ok(Redirect::to("/auth/login"))
}

#[derive(thiserror::Error, Debug)]
pub enum SignupConfirmError {
    #[error("Invalid Token")]
    InvalidToken,
    #[error("CSRF Validation Failed")]
    CsrfValidationFailed,
    #[error("Unknown Error: {0}")]
    Unknown(#[from] anyhow::Error),
}

impl IntoResponse for SignupConfirmError {
    fn into_response(self) -> axum::response::Response {
        let (status, message) = match self {
            SignupConfirmError::InvalidToken => (StatusCode::UNAUTHORIZED, "Invalid Token"),
            SignupConfirmError::CsrfValidationFailed => {
                (StatusCode::BAD_REQUEST, "CSRF Validation Failed, Try Again")
            }
            SignupConfirmError::Unknown(e) => {
                error!(?e, "returning INTERNAL SERVER ERROR");
                (StatusCode::INTERNAL_SERVER_ERROR, "Unknown Error")
            }
        };

        (
            status,
            ErrorPage {
                error: message.to_string(),
            },
        )
            .into_response()
    }
}