ordinary-app 0.7.0

// Copyright (C) 2026 Ordinary Labs, LLC.
//
// SPDX-License-Identifier: AGPL-3.0-only

use std::sync::Arc;

use aes_gcm::{
    Aes256Gcm, Key,
    aead::{Aead, AeadCore, OsRng},
};
use axum::body::Bytes;
use axum::extract::{Form, State};
use axum::http::header::CONTENT_TYPE;
use axum::http::{StatusCode, Uri};
use axum::response::IntoResponse;
use base64::{Engine as B64Engine, engine::general_purpose::URL_SAFE_NO_PAD as b64};
use bytes::{BufMut, BytesMut};
use hyper::HeaderMap;
use ordinary_auth::AuthClient;
use ordinary_config::ClientPasswordHash;
use ordinary_template::TemplateResult;
use ordinary_utils::{get_bearer_token_as_bytes, get_host};
use serde::Deserialize;
use sha2::{Digest, Sha256};
use tracing::Instrument;
use x25519_dalek::{EphemeralSecret, PublicKey};

pub async fn start(
    State(state): State<Arc<crate::server::OrdinaryAppServerState>>,
    headers: HeaderMap,
    body: Bytes,
) -> impl IntoResponse {
    let span = tracing::info_span!("auth", flv = %"wasm");

    span.in_scope(|| {
        let invite_token: Vec<u8>;

        let checked_claims = if state.auth.config.invite.is_some() {
            let Ok(token) = get_bearer_token_as_bytes(&headers) else {
                return StatusCode::UNAUTHORIZED.into_response();
            };

            invite_token = token;

            let Ok(checked_claims) = state.auth.invite_check(invite_token.as_ref()) else {
                return StatusCode::UNAUTHORIZED.into_response();
            };

            Some(checked_claims)
        } else {
            None
        };

        match state
            .auth
            .registration_start(body, None, None, checked_claims)
        {
            Ok(v) => (StatusCode::OK, v).into_response(),
            Err(err) => {
                tracing::error!(%err);
                StatusCode::INTERNAL_SERVER_ERROR.into_response()
            }
        }
    })
}

pub async fn finish(
    State(state): State<Arc<crate::server::OrdinaryAppServerState>>,
    body: Bytes,
) -> impl IntoResponse {
    let span = tracing::info_span!("auth", flv = %"wasm");

    span.in_scope(|| match state.auth.registration_finish(body, None) {
        Ok((v, account_bytes, _invite_claims)) => {
            if let Ok(account) = std::str::from_utf8(&account_bytes) {
                let mut builder =
                    flexbuffers::Builder::new(&flexbuffers::BuilderOptions::SHARE_NONE);
                let mut vec_builder = builder.start_vector();
                vec_builder.push(account);

                // empty claims
                vec_builder.push(());

                vec_builder.end_vector();

                for action in &state.registration_actions {
                    if let Some(action) = state.actions.get(*action as usize) {
                        let span = tracing::info_span!("action");

                        span.in_scope(|| {
                            if let Err(err) = action.call(builder.view(), &state.actions) {
                                tracing::error!(%err);
                            }
                        });
                    }
                }
            }

            (StatusCode::OK, v)
        }
        Err(e) => {
            tracing::error!("{e}");
            (StatusCode::INTERNAL_SERVER_ERROR, Bytes::new())
        }
    })
}

#[derive(Deserialize, Debug)]
pub struct RegisterForm {
    account: String,
    password: String,
    invite_token: Option<String>,
}

#[allow(clippy::too_many_lines)]
/// for <noscript> users
pub async fn form(
    uri: Uri,
    headers: HeaderMap,
    State(state): State<Arc<crate::server::OrdinaryAppServerState>>,
    Form(register_form): Form<RegisterForm>,
) -> impl IntoResponse {
    let span = tracing::info_span!("auth", flv = %"noscript");

    async {
        let Some(host) = get_host(&headers, &uri) else {
            tracing::error!("no host");
            return StatusCode::BAD_REQUEST.into_response();
        };

        let account = register_form.account.as_bytes();

        if account.len() > 255 {
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                [(CONTENT_TYPE, "text/plain".to_string())],
                Bytes::copy_from_slice(b"encoded account cannot be more than 255 bytes!"),
            )
                .into_response();
        }

        let password = register_form.password.as_bytes();
        let app_name = state.config.domain.as_bytes();

        let mut input = app_name.to_vec();
        input.extend_from_slice(account);
        input.extend_from_slice(password);

        let mut password = vec![];

        if let Some(auth) = &state.config.auth {
            match &auth.client_hash {
                ClientPasswordHash::Sha256 => {
                    let mut hasher = Sha256::new();
                    hasher.update(&input);
                    password = hasher.finalize().to_vec();
                }
            }
        }

        let invite_token: Vec<u8>;

        let checked_claims = if state.auth.config.invite.is_some()
            && let Some(token) = register_form.invite_token
        {
            match b64.decode(token) {
                Ok(token) => {
                    invite_token = token;

                    let Ok(checked_claims) = state.auth.invite_check(invite_token.as_ref()) else {
                        return StatusCode::UNAUTHORIZED.into_response();
                    };

                    Some(checked_claims)
                }
                Err(err) => {
                    tracing::error!(%err);
                    return (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        [(CONTENT_TYPE, "text/plain".to_string())],
                        Bytes::copy_from_slice(b"error!"),
                    )
                        .into_response();
                }
            }
        } else {
            None
        };

        match AuthClient::registration_start_req(account, &password[..]) {
            Ok((client_state, req)) => {
                match state
                    .auth
                    .registration_start(req, None, None, checked_claims)
                {
                    Ok(server_message) => match AuthClient::registration_finish_req(
                        account,
                        &password[..],
                        &client_state,
                        &server_message[..],
                    ) {
                        Ok((private_key, req)) => match state.auth.registration_finish(req, None) {
                            Ok((res, account_bytes, _invite_claims)) => {
                                if let Ok(account) = std::str::from_utf8(&account_bytes) {
                                    let mut builder = flexbuffers::Builder::new(
                                        &flexbuffers::BuilderOptions::SHARE_NONE,
                                    );
                                    let mut vec_builder = builder.start_vector();
                                    vec_builder.push(account);

                                    // empty claims
                                    vec_builder.push(());

                                    vec_builder.end_vector();

                                    for action in &state.registration_actions {
                                        if let Some(action) = state.actions.get(*action as usize) {
                                            let span = tracing::info_span!("action");

                                            span.in_scope(|| {
                                                if let Err(err) =
                                                    action.call(builder.view(), &state.actions)
                                                {
                                                    tracing::error!(%err);
                                                }
                                            });
                                        }
                                    }
                                }

                                match AuthClient::decrypt_totp_mfa_to_qr_svg(
                                    &res,
                                    private_key,
                                    state.config.domain.clone(),
                                    register_form.account,
                                ) {
                                    Ok((qr_code, recovery_codes)) => {
                                        if let Some(idx) = state.mfa_totp_template_idx
                                            && let Some(template) =
                                                state.templates.get(idx as usize)
                                        {
                                            match std::str::from_utf8(&account_bytes) {
                                                Ok(account) => {
                                                    match template.render(
                                                        host.as_str(),
                                                        "/totp".into(),
                                                        None,
                                                        None,
                                                        Some((
                                                            qr_code.clone(),
                                                            account.to_string(),
                                                            recovery_codes,
                                                        )),
                                                        &None,
                                                    ) {
                                                        Ok(res) => {
                                                            return match res {
                                                                TemplateResult::Result(bytes) => (
                                                                    StatusCode::OK,
                                                                    [(
                                                                        CONTENT_TYPE,
                                                                        template
                                                                            .config
                                                                            .mime
                                                                            .clone(),
                                                                    )],
                                                                    bytes,
                                                                )
                                                                    .into_response(),
                                                                TemplateResult::StatusCode(
                                                                    code,
                                                                ) => code.into_response(),
                                                            };
                                                        }
                                                        Err(err) => tracing::error!(%err),
                                                    }
                                                }
                                                Err(err) => tracing::error!(%err),
                                            }
                                        }

                                        return (
                                            StatusCode::OK,
                                            [(CONTENT_TYPE, "text/html".to_string())],
                                            Bytes::copy_from_slice(qr_code.as_bytes()),
                                        )
                                            .into_response();
                                    }
                                    Err(err) => tracing::error!(%err),
                                }
                            }
                            Err(err) => tracing::error!(%err),
                        },
                        Err(err) => tracing::error!(%err),
                    },
                    Err(err) => tracing::error!(%err),
                }
            }
            Err(err) => tracing::error!(%err),
        }

        (
            StatusCode::INTERNAL_SERVER_ERROR,
            [(CONTENT_TYPE, "text/plain".to_string())],
            Bytes::copy_from_slice(b"error!"),
        )
            .into_response()
    }
    .instrument(span)
    .await
}

const ZEROED_KEY: [u8; 32] = [0u8; 32];

#[allow(clippy::too_many_lines)]
pub async fn hash_only(
    State(state): State<Arc<crate::server::OrdinaryAppServerState>>,
    headers: HeaderMap,
    body: Bytes,
) -> impl IntoResponse {
    let span = tracing::info_span!("auth", flv = %"js");

    span.in_scope(|| {
        let invite_token: Vec<u8>;

        let checked_claims = if state.auth.config.invite.is_some() {
            let Ok(token) = get_bearer_token_as_bytes(&headers) else {
                return StatusCode::UNAUTHORIZED.into_response();
            };

            invite_token = token;

            let Ok(checked_claims) = state.auth.invite_check(invite_token.as_ref()) else {
                return StatusCode::UNAUTHORIZED.into_response();
            };

            Some(checked_claims)
        } else {
            None
        };

        if let Some(account_len) = body.first() {
            let account_len = *account_len as usize;

            if body.len() < account_len + 1 + 32 + 32 {
                return StatusCode::BAD_REQUEST.into_response();
            }

            let account_bytes = &body[1..=account_len];
            let password = &body[account_len + 1..account_len + 1 + 32];
            let public_key: [u8; 32] =
                match body[account_len + 1 + 32..account_len + 1 + 32 + 32].try_into() {
                    Ok(pk) => {
                        if pk == ZEROED_KEY {
                            tracing::error!("public key cannot be all 0s");
                            return StatusCode::BAD_REQUEST.into_response();
                        }

                        pk
                    }
                    Err(err) => {
                        tracing::error!(%err);
                        return StatusCode::EXPECTATION_FAILED.into_response();
                    }
                };

            match AuthClient::registration_start_req(account_bytes, password) {
                Ok((client_state, req)) => {
                    match state
                        .auth
                        .registration_start(req, None, None, checked_claims)
                    {
                        Ok(server_message) => match AuthClient::registration_finish_req(
                            account_bytes,
                            password,
                            &client_state,
                            &server_message[..],
                        ) {
                            Ok((private_key, req)) => {
                                match state.auth.registration_finish(req, None) {
                                    Ok((res, account_bytes, _invite_claims)) => {
                                        if let Ok(account) = std::str::from_utf8(&account_bytes) {
                                            let mut builder = flexbuffers::Builder::new(
                                                &flexbuffers::BuilderOptions::SHARE_NONE,
                                            );
                                            let mut vec_builder = builder.start_vector();
                                            vec_builder.push(account);

                                            // empty claims
                                            vec_builder.push(());

                                            vec_builder.end_vector();

                                            for action in &state.registration_actions {
                                                if let Some(action) =
                                                    state.actions.get(*action as usize)
                                                {
                                                    let span = tracing::info_span!("action");

                                                    span.in_scope(|| {
                                                        if let Err(err) = action
                                                            .call(builder.view(), &state.actions)
                                                        {
                                                            tracing::error!(%err);
                                                        }
                                                    });
                                                }
                                            }

                                            match AuthClient::decrypt_totp_mfa_to_qr_svg(
                                                &res,
                                                private_key,
                                                state.config.domain.clone(),
                                                account.into(),
                                            ) {
                                                Ok((qr_code, recovery_codes)) => {
                                                    use aes_gcm::aead::KeyInit;

                                                    let public_key = PublicKey::from(public_key);
                                                    let ephemeral_secret =
                                                        EphemeralSecret::random_from_rng(OsRng);
                                                    let ephemeral_public_key =
                                                        PublicKey::from(&ephemeral_secret);

                                                    let shared_secret = ephemeral_secret
                                                        .diffie_hellman(&public_key);
                                                    if !shared_secret.was_contributory() {
                                                        tracing::error!(
                                                            "non-contributory shared secret"
                                                        );
                                                        return StatusCode::BAD_REQUEST
                                                            .into_response();
                                                    }

                                                    let key = Key::<Aes256Gcm>::from_slice(
                                                        shared_secret.as_bytes(),
                                                    );

                                                    let cipher = Aes256Gcm::new(key);
                                                    let nonce =
                                                        Aes256Gcm::generate_nonce(&mut OsRng); // 96-bits; unique per message

                                                    let ciphertext = match cipher.encrypt(
                                                        &nonce,
                                                        format!("{recovery_codes}__{qr_code}")
                                                            .as_ref(),
                                                    ) {
                                                        Ok(v) => v,
                                                        Err(err) => {
                                                            tracing::error!(%err);
                                                            return StatusCode::INTERNAL_SERVER_ERROR
                                                            .into_response();
                                                        }
                                                    };

                                                    let mut res = BytesMut::new();

                                                    res.put(&ephemeral_public_key.as_bytes()[..]);
                                                    res.put(&nonce[..]);
                                                    res.put(&ciphertext[..]);

                                                    return (StatusCode::OK, res).into_response();
                                                }
                                                Err(err) => tracing::error!(%err),
                                            }
                                        }
                                    }
                                    Err(err) => tracing::error!(%err),
                                }
                            }
                            Err(err) => tracing::error!(%err),
                        },
                        Err(err) => tracing::error!(%err),
                    }
                }
                Err(err) => tracing::error!(%err),
            }
        }

        StatusCode::INTERNAL_SERVER_ERROR.into_response()
    })
}