nom-nom-nix-gc/src/handlers/webauthn.rs

use actix_web::{web, error, HttpResponse, http::header::ContentType, cookie::{Cookie, SameSite}};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use webauthn_rs::prelude::{RegisterPublicKeyCredential, PasskeyRegistration, Passkey, RequestChallengeResponse, PublicKeyCredential};

use crate::{models::{AppState, RegistrationUuid, Key, LoginUuid, SessionUuid}, templates};

/// Page you get after clicking a register-user link.
///
/// The path contains the registration Uuid segment.
pub async fn webauthn_registration(app_state: web::Data<AppState<'_>>, path: web::Path<String>) -> HttpResponse {
    let uuid: RegistrationUuid = match Uuid::parse_str(&path.into_inner()) {
        Ok(p) => RegistrationUuid(p),
        Err(_) => return HttpResponse::from_error(error::ErrorBadRequest("This registration token is invalid: invalid UUID")),
    };
    let user = app_state.retrieve_registration_user(&uuid).await;
    let user = match user {
        Ok(Some(user)) => user,
        Ok(None) => return HttpResponse::from_error(error::ErrorBadRequest("This registration token is invalid: cannot find it in the database")),
        Err(e) => return HttpResponse::from_error(error::ErrorInternalServerError(e.to_string())),
    };
    let keys = match app_state.get_user_keys(&user.uuid).await {
        Ok(keys) => keys.into_iter().map(|k| k.name).collect(),
        Err(e) => {
            eprintln!("webauthn_registration: error while retrieving keys: {}", e);
            return HttpResponse::from_error(error::ErrorInternalServerError("Cannot retrieve keys associated to your user."));
        }
    };
    let response = templates::register_user_start(app_state.hbs.clone(), uuid, user.name, keys).unwrap();
    HttpResponse::Ok()
        .content_type(ContentType::html())
        .body(response)
}

#[derive(Debug, Deserialize)]
pub struct InitRegistrationBody {
    uuid: RegistrationUuid
}

/// First phase of the webauthn key enrolling.
///
/// For now, we don't save anything to the DB. We just generate the
/// challenge, the UUID and store everything in the session hashmap
/// server-side, in the response and cookie on the client-side.
pub async fn start_webauthn_registration(app_state: web::Data<AppState<'_>>, body: web::Json<InitRegistrationBody>) -> HttpResponse {
    let res = app_state.retrieve_registration_user(&body.uuid).await;
    let user = match res {
        Err(e) => {
            eprintln!("start_webauthn_registration: internal server error {}", e);
            return HttpResponse::from_error(error::ErrorInternalServerError("Internal server error"))
        },
        // Note: there's a constraint set in the DB. We shouldn't end up in this branch.
        Ok(None) => return HttpResponse::from_error(error::ErrorInternalServerError("Can't find the user attached to this registration.")),
        Ok(Some(user)) => user
    };
    let (creation_challenge_response, passkey_registration) = app_state.webauthn.start_passkey_registration(user.uuid.0, &user.name, &user.name, None).unwrap();
    {
        let mut user_registrations = app_state.session.user_registrations.write().await;
        user_registrations.insert(body.uuid.clone(), passkey_registration);
    }
    let res = serde_json::to_string(&creation_challenge_response).unwrap();
    HttpResponse::Ok()
        .content_type(ContentType::json())
        .body(res)
}

#[derive(Clone, Debug, Deserialize)]
pub struct RegisterWebauthnPayload {
    uuid: RegistrationUuid,
    keyname: String,
    credentials: RegisterPublicKeyCredential
}

/// Final phase of the webauthn key enrolling.
///
/// We verify the key enrolling challenge, then store the new key in the DB.
pub async fn finish_webauthn_registration(app_state: web::Data<AppState<'_>>, register: web::Json<RegisterWebauthnPayload>) -> impl actix_web::Responder {
    let passkey_registration: Option<PasskeyRegistration> = {
        let mut user_registrations = app_state.session.user_registrations.write().await;
        let pr = user_registrations.get(&register.uuid).cloned();
        user_registrations.remove(&register.uuid);
        pr
    };
    let db_usr = app_state.retrieve_registration_user(&register.uuid).await;
    let user = match db_usr {
        Ok(Some(usr)) => usr,
        // Note: there's a constraint set in the DB. We shouldn't end up in this branch.
        Ok(None) => return HttpResponse::from_error(error::ErrorInternalServerError("Can't find the user attached to this registration.")),
        Err(e) => {
            eprintln!("finish_webauthn_registration: internal server error {}", e);
            return HttpResponse::from_error(error::ErrorInternalServerError("Internal server error"))
        }
    };

    match passkey_registration {
        Some(passkey) => {

            let registration_result = app_state.webauthn.finish_passkey_registration(&register.credentials, &passkey);
            match registration_result {
                Ok(passkey) => {
                    let passkey = Key {
                        name: register.keyname.clone(),
                        key: passkey,
                    };
                    // TODO: add cookie to the answer to auth the
                    // user? Maybe rather later, when we're sure we
                    // have all the keys required.
                    app_state.save_user_key(&user.uuid, &passkey).await.unwrap();
                    HttpResponse::Ok()
                        .body("ok")
                },
                Err(_) =>
                    HttpResponse::from_error(error::ErrorUnauthorized("Webauthn challenge failed"))
            }
        },
        None => {
            HttpResponse::from_error(error::ErrorInternalServerError("Session expired"))
        }

    }
}


pub async fn webauthn_login(app_state: web::Data<AppState<'_>>) -> impl actix_web::Responder {
    let response = templates::login(app_state.hbs.clone()).unwrap();
    HttpResponse::Ok()
        .content_type(ContentType::html())
        .body(response)
}


#[derive(Deserialize, Debug)]
pub struct InitLoginBody {
    username: String
}

#[derive(Serialize, Debug)]
pub struct InitLoginResp {
    uuid: LoginUuid,
    challenge: RequestChallengeResponse
}

pub async fn webauthn_login_init(app_state: web::Data<AppState<'_>>, body: web::Json<InitLoginBody>) -> impl actix_web::Responder {
    let Ok(user) = app_state.get_user(&body.username).await else {
        return HttpResponse::from_error(error::ErrorBadRequest("Invalid username"))
    };
    let Ok(keys) = app_state.get_user_keys(&user.uuid).await else {
        return HttpResponse::from_error(error::ErrorInternalServerError("Cannot find keys attached to user"));
    };
    let keys: Vec<Passkey> = keys.into_iter().map(|k| k.key).collect();
    let (challenge, passkey_auth) = match app_state.webauthn.start_passkey_authentication(keys.as_slice()) {
        Ok((c, p)) => (c, p),
        Err(e) => {
            eprintln!("webauthn_login_init: webauthn error: {}", e);
            return HttpResponse::from_error(error::ErrorInternalServerError("Webauthn error"));
        }
    };
    let uuid = LoginUuid(Uuid::new_v4());
    {
        let mut user_logins =  app_state.session.user_pending_logins.write().await;
        user_logins.insert(uuid.clone(), (passkey_auth, user));
    }
    let res = InitLoginResp {
        uuid,
        challenge,
    };
    let res = serde_json::to_string(&res).unwrap();
    HttpResponse::Ok()
        .content_type(ContentType::json())
        .body(res)
}

#[derive(Deserialize, Debug)]
pub struct FinishLoginBody {
    challenge: PublicKeyCredential,
    uuid: LoginUuid
}

pub async fn webauthn_login_finish(app_state: web::Data<AppState<'_>>, body: web::Json<FinishLoginBody>) -> impl actix_web::Responder {
    let passkey_auth = {
        let mut user_logins = app_state.session.user_pending_logins.write().await;
        let pk = user_logins.get(&body.uuid).cloned();
        user_logins.remove(&body.uuid);
        pk
    };
    let Some((passkey_auth,user)) = passkey_auth else {
        eprintln!("webauthn_login_finish: cannot find login state");
        return HttpResponse::from_error(error::ErrorBadRequest("Cannot find login state in session"));
    };
    let res = match app_state.webauthn.finish_passkey_authentication(&body.challenge, &passkey_auth) {
        Ok(res) => res,
        Err(e) => {
            eprintln!("webauthn_login_finish: webauthn error: {}", e);
            return HttpResponse::from_error(error::ErrorForbidden("Failed webauthn auth"));
        },
    };
    let session_uuid = {
        let mut user_sessions = app_state.session.user_sessions.write().await;
        let session_uuid = SessionUuid(Uuid::new_v4());
        user_sessions.insert(session_uuid, user.clone());
        session_uuid
    };
    if res.needs_update() {
        match app_state.update_user_keys(&user.uuid, res).await {
            Ok(()) => (),
            Err(e) => {
                eprintln!("Error while updating the key states: {}", e);
                return HttpResponse::from_error(error::ErrorInternalServerError("Error while updating the key states."));
            }
        }
    }
    let response = templates::login(app_state.hbs.clone()).unwrap();
    #[cfg(debug_assertions)]
    let secure = false;
    #[cfg(not(debug_assertions))]
    let secure = true;
    let auth_cookie = Cookie::build("auth-uuid", session_uuid.0.to_string())
        .http_only(true)
        .same_site(SameSite::Strict)
        .secure(secure)
        .finish();
    HttpResponse::Ok()
        .content_type(ContentType::html())
        .cookie(auth_cookie)
        .body(response)
}