tld-registration/src/main.rs

mod email;
mod password;
mod session;

use std::str::FromStr;

use crate::email::send_registration_email;
use crate::password::Password;
use crate::session::Session;

use actix_csrf::extractor::{CsrfCookie, CsrfToken};
use actix_csrf::Csrf;
use actix_session::CookieSession;
use actix_web::cookie::SameSite;
use actix_web::error::InternalError;
use actix_web::http::header::LOCATION;
use actix_web::http::{Method, StatusCode, Uri};
use actix_web::middleware::Logger;
use actix_web::web::{Form, Query};
use actix_web::{get, post, App, HttpRequest, HttpResponse, HttpServer, Responder};
use handlebars::Handlebars;
use lettre::Address as EmailAddress;
use once_cell::sync::{Lazy, OnceCell};
use rand::prelude::StdRng;
use serde::{Deserialize, Serialize};
use sodiumoxide::crypto::pwhash::argon2id13::{self, HashedPassword, HASHEDPASSWORDBYTES};
use sqlx::sqlite::SqliteConnectOptions;
use sqlx::SqlitePool;
use tracing::error;

static TEMPLATE_ENGINE: Lazy<Handlebars> = Lazy::new(|| {
    let mut handlebars = Handlebars::new();

    handlebars.set_strict_mode(true);
    handlebars.set_dev_mode(true);

    handlebars
        .register_templates_directory(".hbs", "src/templates")
        .expect("failed to load template directory");
    handlebars
});

static DB_POOL: OnceCell<SqlitePool> = OnceCell::new();

#[actix_rt::main]
async fn main() -> std::io::Result<()> {
    tracing_subscriber::fmt::init();
    sodiumoxide::init().unwrap();

    let db = {
        let db_options = SqliteConnectOptions::from_str("db.sqlite")
            .unwrap()
            .create_if_missing(true);
        let pool = SqlitePool::connect_with(db_options).await.unwrap();
        sqlx::query_file!("db_queries/init.sql")
            .execute(&pool)
            .await
            .unwrap();
        pool
    };
    DB_POOL.set(db).unwrap();

    HttpServer::new(|| {
        App::new()
            .wrap(Logger::default())
            .wrap(
                CookieSession::private(&[0; 32])
                    .name("session")
                    .path("/")
                    .secure(true)
                    .http_only(true)
                    .same_site(SameSite::Strict),
            )
            .wrap(
                Csrf::<StdRng>::new()
                    .set_cookie(Method::GET, "/login")
                    .validate_cookie(Method::POST, "/login")
                    .set_cookie(Method::GET, "/register")
                    .validate_cookie(Method::POST, "/register"),
            )
            .service(index)
            .service(login_ui)
            .service(login)
            .service(register_ui)
            .service(register)
            .service(registration_confirmation)
            .service(account_ui)
            .service(logout)
            .service(actix_files::Files::new("/static", "src/static"))
    })
    .bind("127.0.0.1:8080")?
    .run()
    .await
}

/// Macro meme to render a template without any context or with a provided one.
macro_rules! render {
    ($template_name:literal) => {
        render!($template_name, &())
    };
    ($template_name:literal, $template_args:expr $(,)?) => {
        match TEMPLATE_ENGINE.render($template_name, $template_args) {
            Ok(resp) => Ok(HttpResponse::Ok().body(resp)),
            Err(e) => {
                error!("{}", e);
                Err(InternalError::new(e, StatusCode::INTERNAL_SERVER_ERROR))
            }
        }
    };
}

#[get("/")]
async fn index() -> impl Responder {
    render!("index")
}

#[derive(Deserialize)]
struct LoginQuery {
    error: Option<String>,
}

#[get("/login")]
async fn login_ui(csrf: CsrfToken, mut query: Query<LoginQuery>) -> impl Responder {
    #[derive(Serialize)]
    struct TemplateArgs {
        error: Option<String>,
        csrf: CsrfToken,
    }

    render!(
        "login",
        &TemplateArgs {
            error: query.error.take(),
            csrf,
        },
    )
}

#[derive(Deserialize)]
struct Login {
    csrf: CsrfToken,
    email: EmailAddress,
    password: Password,
}

#[post("/login")]
async fn login(csrf_cookie: CsrfCookie, form: Form<Login>, session: Session) -> impl Responder {
    if !csrf_cookie.validate(form.csrf.as_ref()) {
        return HttpResponse::BadRequest().finish();
    }

    let email: &str = form.email.as_ref();
    let verified = sqlx::query!("SELECT password FROM users WHERE email = ?", email)
        .fetch_one(DB_POOL.get().expect("db connection to be set"))
        .await;

    if let Ok(record) = verified {
        if form
            .password
            .verify(&HashedPassword::from_slice(&record.password).unwrap())
        {
            let redirect_to = session.get_redirect_url();
            session.init(&form.email);
            return redirect(
                &redirect_to
                    .as_ref()
                    .map(Uri::to_string)
                    .unwrap_or("/account".into()),
            );
        }
    } else {
        // To guard against timing attacks, we'll construct a fake password to
        // hash. We won't even check if it's successful, we just need to compute
        // the hash. Since we don't check the result, this must be in a separate
        // branch from the success branch, else it's possible to actually
        // succeed from a bogey hash.
        let mut data = [0_u8; HASHEDPASSWORDBYTES];
        assert!(argon2id13::STRPREFIX.len() < HASHEDPASSWORDBYTES);
        for (i, c) in argon2id13::STRPREFIX.iter().enumerate() {
            data[i] = *c;
        }

        // Rust shouldn't optimize this out since it ultimately calls out to
        // a C function, so it shouldn't find out that the function is pure.
        form.password
            .verify(&HashedPassword::from_slice(&data).expect("Should be valid password data"));
    }

    redirect("/login?error=true")
}

#[get("/register")]
async fn register_ui(csrf: CsrfToken) -> impl Responder {
    #[derive(Serialize)]
    struct TemplateArgs {
        csrf: CsrfToken,
    }

    render!("register", &TemplateArgs { csrf })
}

#[derive(Deserialize)]
struct RegistrationInfo {
    csrf: CsrfToken,
    email: EmailAddress,
    password: Password,
}

#[post("/register")]
async fn register(csrf_cookie: CsrfCookie, form: Form<RegistrationInfo>) -> impl Responder {
    if !csrf_cookie.validate(form.csrf.as_ref()) {
        return HttpResponse::BadRequest().finish();
    }

    let hashed = if let Ok(res) = form.password.hashed() {
        res
    } else {
        return HttpResponse::InternalServerError().finish();
    };
    let hashed = hashed.as_ref();

    let email: &str = form.email.as_ref();
    let insert_res = sqlx::query!(
        "INSERT INTO users (email, password) VALUES (?, ?)",
        email,
        hashed,
    )
    .execute(DB_POOL.get().expect("db connection to be set"))
    .await;

    if insert_res.is_ok() {
        send_registration_email(form.email.clone()).await.unwrap();
    }

    redirect("/registration_confirmation")
}

#[get("/registration_confirmation")]
async fn registration_confirmation() -> impl Responder {
    render!("registration_confirmation")
}

#[get("/account")]
async fn account_ui(req: HttpRequest, session: Session) -> impl Responder {
    if let Err(error) = session.validate_or_redirect(req.uri()) {
        return Ok(error);
    }

    render!("account")
}

#[get("/logout")]
async fn logout(session: Session) -> impl Responder {
    // While this is a state-mutating endpoint, it is fine to not have a CSRF
    // token; the worst case is that the user is logged out, which is fail-safe.
    session.purge();
    redirect("/")
}

fn redirect(location: &str) -> HttpResponse {
    HttpResponse::SeeOther()
        .append_header((LOCATION, location))
        .finish()
}