use std::convert::Infallible; use argon2::Argon2; use either::Either; use password_hash::{PasswordHasher, SaltString, rand_core::OsRng}; use rocket::{ FromForm, Request, Route, TypedError, async_trait, form::Form, get, outcome::{Outcome, try_outcome}, post, request::FromRequest, response::Redirect, routes, serde::{Deserialize, Serialize}, trace::info, uri, }; use rocket_db_pools::Connection; use rocket_dyn_templates::{Template, context}; use crate::{ Db, SqlError, auth::{self, AdminUser, BasicAuth}, }; #[get("/", rank = 1)] fn dash(auth: Option) -> Template { info!("user: {:?}", auth.as_ref().map(|a| a.username())); Template::render( "dash", context! { username: auth.as_ref().map(|a| a.username()), admin: auth.as_ref().map(|a| a.is_admin()).unwrap_or(false), mode: "dark", }, ) } #[get("/login")] fn login(auth: Option) -> Either { if auth.is_some() { Either::Right(Redirect::temporary(uri!(dash))) } else { Either::Left(Template::render( "login", context! { mode: "dark", }, )) } } #[get("/account")] fn account(auth: BasicAuth) -> Template { Template::render( "account", context! { username: auth.username(), admin: auth.is_admin(), mode: "dark", }, ) } #[get("/admin")] async fn admin(auth: AdminUser, mut db: Connection) -> Result { let users = sqlx::query!("SELECT * from users") .fetch_all(&mut **db) .await?; Ok(Template::render( "admin", context! { username: auth.username(), admin: auth.is_admin(), users: users.into_iter().filter(|r| r.name != auth.username()).map(|r| r.name).collect::>(), mode: "dark", }, )) } struct InitialSetup; #[derive(Debug, TypedError)] pub enum InitialSetupError { #[error(status = 500)] InternalError, #[error(status = 500)] DbError(#[error(source)] SqlError), } impl From for InitialSetupError { fn from(value: SqlError) -> Self { Self::DbError(value) } } impl From for InitialSetupError { fn from(value: sqlx::Error) -> Self { Self::DbError(value.into()) } } #[derive(Debug, TypedError)] struct Empty; impl From for Empty { fn from(value: Infallible) -> Self { match value {} } } #[async_trait] impl<'r> FromRequest<'r> for InitialSetup { type Forward = Empty; type Error = InitialSetupError; async fn from_request(req: &'r Request<'_>) -> Outcome { // TODO: cache result in managed state? let mut db = try_outcome!( req.guard::>() .await .map_error(|_| InitialSetupError::InternalError) ); let user_count = match sqlx::query!("SELECT count(name) as count from users") .fetch_one(&mut **db) .await { Ok(v) => v, Err(e) => return Outcome::Error(e.into()), }; if user_count.count == 0 { Outcome::Success(Self) } else { Outcome::Forward(Empty) } } } #[get("/")] fn initial_setup(_initial: InitialSetup) -> Template { Template::render("initial_setup", context! {}) } #[derive(Debug, Serialize, Deserialize, FromForm)] #[serde(crate = "rocket::serde")] pub struct InitialSetupOptions<'a> { admin_username: &'a str, admin_password: &'a str, } #[post("/initial_setup", data = "")] async fn initial_complete( _initial: InitialSetup, opts: Form>, mut db: Connection, ) -> Result { let salt = SaltString::generate(&mut OsRng); let pw_hash = Argon2::default() .hash_password(opts.admin_password.as_bytes(), salt.as_salt()) .expect("Failed to hash password"); sqlx::query("INSERT INTO users (name, role, password) VALUES (?1, \"admin\", ?2)") .bind(opts.admin_username) .bind(pw_hash.to_string()) .execute(&mut **db) .await?; Ok(Redirect::to("/")) } pub fn routes() -> Vec { routes![ dash, login, auth::login_form, auth::update_password, account, admin, initial_setup, initial_complete, ] }