iota/omikron-connector/src/user_ops.rs

use base64::{Engine as _, engine::general_purpose::STANDARD};
use iota_logger::{PrintType, log, log_cv, log_t};
use iota_storage::users::user_manager::try_add_user;
use iota_storage::users::user_profile::UserProfile;
use iota_storage::util::config_util::CONFIG;
use iota_util::crypto_helper::{self, hex_hash, public_key_bundle_to_base64};
use iota_util::file_util::write_user_credential;
use iota_util::tu::TuCredential;
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
use mtp::crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme};
use rand_core::{OsRng, RngCore};
use std::time::Duration;

use crate::OmikronClient;
use crate::omega_discovery;

#[derive(Debug)]
pub enum CreateUserError {
    InvalidUsername,
    Transport(crate::OmikronError),
    InvalidResponse,
    RemoteRejected,
    LocalPersistence(String),
}

#[derive(Debug)]
pub enum LifecycleUserError {
    InvalidCredential(String),
    OmegaHostMismatch,
    RemoteRejected,
    Transport(crate::OmikronError),
    LocalPersistence(String),
}

impl From<crate::OmikronError> for LifecycleUserError {
    fn from(value: crate::OmikronError) -> Self { Self::Transport(value) }
}

fn lifecycle_payload(domain: &[u8], user_id: i64, iota_id: i64, nonce: u64) -> Vec<u8> {
    let mut payload = Vec::with_capacity(domain.len() + 24);
    payload.extend_from_slice(domain);
    payload.extend_from_slice(&user_id.to_be_bytes());
    payload.extend_from_slice(&iota_id.to_be_bytes());
    payload.extend_from_slice(&nonce.to_be_bytes());
    payload
}

fn configured_iota_id() -> Result<i64, LifecycleUserError> {
    CONFIG.load().iota_id
        .and_then(|id| i64::try_from(id).ok())
        .filter(|id| *id > 0)
        .ok_or_else(|| LifecycleUserError::InvalidCredential("Iota identity is not registered".into()))
}

fn sign_lifecycle_payload(credential: &TuCredential, payload: &[u8]) -> Result<(Vec<u8>, Vec<u8>), LifecycleUserError> {
    let classical = Ed25519Signer::new(&credential.keyring.sig_cl_secret_key)
        .map_err(|error| LifecycleUserError::InvalidCredential(error.to_string()))?
        .sign(payload)
        .map_err(|error| LifecycleUserError::InvalidCredential(error.to_string()))?;
    let pq = MlDsaSigner::new(&credential.keyring.sig_pq_secret_key, &credential.keyring.sig_pq_public_key)
        .map_err(|error| LifecycleUserError::InvalidCredential(error.to_string()))?
        .sign(payload)
        .map_err(|error| LifecycleUserError::InvalidCredential(error.to_string()))?;
    Ok((classical, pq))
}

async fn inspect_credential_account(
    connection: &dyn OmikronClient,
    credential: &TuCredential,
) -> Result<(String, String), LifecycleUserError> {
    if credential.omega_host != omega_discovery::omega_host() {
        return Err(LifecycleUserError::OmegaHostMismatch);
    }
    let request = CommunicationValue::new(CommunicationType::GetUserData)
        .add_typed_default(DataType::UserId, DataValue::SignedNumber(credential.user_id.into()));
    let response = connection.await_response(&request, Duration::from_secs(20)).await?;
    if !response.is_type(CommunicationType::GetUserData) {
        return Err(LifecycleUserError::RemoteRejected);
    }
    let username = response.get_data(DataType::Username).as_str().map(str::to_owned)
        .ok_or(LifecycleUserError::RemoteRejected)?;
    let public_key = response.get_data(DataType::PublicKey).as_str().map(str::to_owned)
        .ok_or(LifecycleUserError::RemoteRejected)?;
    if public_key != public_key_bundle_to_base64(&credential.public_key_bundle()) {
        return Err(LifecycleUserError::RemoteRejected);
    }
    Ok((username, public_key))
}

async fn credential_proof(
    connection: &dyn OmikronClient,
    credential: &TuCredential,
    begin: CommunicationType,
    challenge: CommunicationType,
    complete: CommunicationType,
    domain: &[u8],
) -> Result<(), LifecycleUserError> {
    let iota_id = configured_iota_id()?;
    let begin_request = CommunicationValue::new(begin)
        .add_typed_default(DataType::UserId, DataValue::SignedNumber(credential.user_id.into()));
    let challenge_response = connection.await_response(&begin_request, Duration::from_secs(20)).await?;
    if !challenge_response.is_type(challenge) {
        return Err(LifecycleUserError::RemoteRejected);
    }
    let nonce = challenge_response.get_data(DataType::ServerNonce).as_signed_number()
        .and_then(|value| u64::try_from(value).ok())
        .ok_or(LifecycleUserError::RemoteRejected)?;
    let (signature, pq_signature) = sign_lifecycle_payload(credential, &lifecycle_payload(domain, credential.user_id, iota_id, nonce))?;
    let complete_request = CommunicationValue::new(complete)
        .add_typed_default(DataType::UserId, DataValue::SignedNumber(credential.user_id.into()))
        .add_typed_default(DataType::ServerNonce, DataValue::SignedNumber(nonce.into()))
        .add_typed_default(DataType::Signature, DataValue::Bytes(signature))
        .add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature));
    let response = connection.await_response(&complete_request, Duration::from_secs(20)).await?;
    if response.is_type(CommunicationType::Success) { Ok(()) } else { Err(LifecycleUserError::RemoteRejected) }
}

/// Attach or migrate an existing account. Local state is written only after
/// Omega has accepted the credential proof and changed its assignment.
pub async fn attach_user_from_tu(connection: &dyn OmikronClient, contents: &str) -> Result<UserProfile, LifecycleUserError> {
    let credential = TuCredential::parse(contents).map_err(|error| LifecycleUserError::InvalidCredential(error.to_string()))?;
    let (username, public_key) = inspect_credential_account(connection, &credential).await?;
    credential_proof(connection, &credential, CommunicationType::AttachUserBegin, CommunicationType::AttachUserChallenge, CommunicationType::AttachUserComplete, b"tensamin:user-attach:v1\0").await?;
    let profile = UserProfile::new(credential.user_id, username, None, public_key, hex_hash(contents), String::new());
    write_user_credential(profile.user_id, &credential.to_canonical_string())
        .map_err(|error| LifecycleUserError::LocalPersistence(error.to_string()))?;
    try_add_user(profile.clone()).map_err(|error| LifecycleUserError::LocalPersistence(error.to_string()))?;
    Ok(profile)
}

pub async fn complete_delete_user_with_tu(connection: &dyn OmikronClient, contents: &str, expected_user_id: i64) -> Result<(), LifecycleUserError> {
    let credential = TuCredential::parse(contents).map_err(|error| LifecycleUserError::InvalidCredential(error.to_string()))?;
    if credential.user_id != expected_user_id { return Err(LifecycleUserError::InvalidCredential("credential user ID does not match deletion target".into())); }
    inspect_credential_account(connection, &credential).await?;
    credential_proof(connection, &credential, CommunicationType::DeleteUserCredentialBegin, CommunicationType::DeleteUserCredentialChallenge, CommunicationType::DeleteUserCredentialComplete, b"tensamin:user-delete:v1\0").await
}

/// Repair local management state after a release or migration committed in
/// Omega but local cleanup was interrupted. Hosted data is retained.
pub async fn reconcile_managed_users(connection: &dyn OmikronClient) {
    let Ok(local_iota_id) = configured_iota_id() else { return; };
    for user in iota_storage::users::user_manager::get_users() {
        let request = CommunicationValue::new(CommunicationType::GetUserData)
            .add_typed_default(DataType::UserId, DataValue::SignedNumber(user.user_id.into()));
        let Ok(response) = connection.await_response(&request, Duration::from_secs(10)).await else { continue; };
        let remote_iota_id = response.get_data(DataType::IotaId).as_signed_number().and_then(|value| i64::try_from(value).ok());
        if remote_iota_id != Some(local_iota_id) {
            let _ = iota_storage::users::user_manager::release_user(user.user_id);
        }
    }
}

fn valid_username(username: &str) -> bool {
    !username.is_empty()
        && username.chars().count() <= 15
        && !username.chars().any(char::is_control)
        && !username.contains(['/', '\\'])
}

async fn request_user_id(connection: &dyn OmikronClient) -> Result<(i64, String), CreateUserError> {
    let request = CommunicationValue::new(CommunicationType::GetRegister);
    let response = connection
        .await_response(&request, Duration::from_secs(20))
        .await
        .map_err(CreateUserError::Transport)?;

    if !response.is_type(CommunicationType::GetRegister) {
        return Err(CreateUserError::InvalidResponse);
    }

    let user_id = response
        .get_data(DataType::UserId)
        .as_number()
        .and_then(|id| i64::try_from(id).ok())
        .filter(|id| (1..(1_i64 << 48)).contains(id))
        .ok_or(CreateUserError::InvalidResponse)?;
    let registration_token = response
        .get_data(DataType::RegisterId)
        .as_str()
        .filter(|token| uuid::Uuid::parse_str(token).is_ok())
        .map(str::to_owned)
        .ok_or(CreateUserError::InvalidResponse)?;
    Ok((user_id, registration_token))
}

/// A completion response can be lost after Omega commits the user.  Confirm
/// the exact remote record before treating that transport failure as success.
async fn registration_committed(connection: &dyn OmikronClient, profile: &UserProfile) -> bool {
    let request = CommunicationValue::new(CommunicationType::GetUserData).add_typed_default(
        DataType::UserId,
        DataValue::SignedNumber(profile.user_id.into()),
    );
    let Ok(response) = connection
        .await_response(&request, Duration::from_secs(5))
        .await
    else {
        return false;
    };
    response.get_data(DataType::UserId).as_number() == Some(profile.user_id.into())
        && response.get_data(DataType::Username).as_str() == Some(profile.username.as_str())
        && response.get_data(DataType::PublicKey).as_str() == Some(profile.public_key.as_str())
}

pub async fn create_user(
    connection: &dyn OmikronClient,
    username: &str,
) -> Result<UserProfile, CreateUserError> {
    if !valid_username(username) {
        return Err(CreateUserError::InvalidUsername);
    }
    let (user_id, registration_token) = request_user_id(connection).await?;
    log!("User creation: Omega allocated user ID {user_id}");
    let keyring = crypto_helper::generate_keyring();
    let pub_key_bundle = keyring.public_key_bundle();
    let keyring_b64 = crypto_helper::keyring_to_base64(&keyring);

    let private_key_hash = hex_hash(&keyring_b64);

    let mut bytes = [0u8; 192];
    OsRng.fill_bytes(&mut bytes);
    let reset_token = STANDARD.encode(&bytes);

    let user_profile = UserProfile::new(
        user_id,
        username.to_string(),
        None,
        public_key_bundle_to_base64(&pub_key_bundle),
        private_key_hash,
        reset_token.clone(),
    );

    let communication_value = CommunicationValue::new(CommunicationType::CompleteRegisterUser)
        .add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into()))
        .add_typed_default(DataType::Username, DataValue::Str(username.to_string()))
        .add_typed_default(
            DataType::PublicKey,
            DataValue::Str(public_key_bundle_to_base64(&pub_key_bundle)),
        )
        .add_typed_default(DataType::ResetToken, DataValue::Str(reset_token))
        .add_typed_default(DataType::RegisterId, DataValue::Str(registration_token));

    let response_communication_value = connection
        .await_response(&communication_value, Duration::from_secs(20))
        .await;

    match response_communication_value {
        Ok(response) => {
            log_cv!(PrintType::Omega, response);
            if !response.is_type(CommunicationType::Success) {
                return Err(CreateUserError::RemoteRejected);
            }
        }
        Err(error) => {
            if registration_committed(connection, &user_profile).await {
                log!(
                    "User creation: completion response was lost; verified user {} remotely",
                    user_id
                );
            } else {
                log_t!("User creation: {}", error.to_string());
                return Err(match error {
                    crate::OmikronError::Internal(_) => CreateUserError::RemoteRejected,
                    error => CreateUserError::Transport(error),
                });
            }
        }
    }
    log!("Created User");
    write_user_credential(
        user_id,
        &format!("{}@{}::{}", user_id, omega_discovery::omega_host(), keyring_b64),
    )
    .map_err(|error| CreateUserError::LocalPersistence(error.to_string()))?;

    try_add_user(user_profile.clone())
        .map_err(|error| CreateUserError::LocalPersistence(error.to_string()))?;
    Ok(user_profile)
}

#[cfg(test)]
mod tests {
    use super::{CreateUserError, request_user_id, valid_username};
    use crate::{OmikronClient, OmikronError};
    use async_trait::async_trait;
    use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
    use std::time::Duration;

    struct RegistrationClient {
        response: CommunicationValue,
    }

    #[async_trait]
    impl OmikronClient for RegistrationClient {
        async fn send_message(&self, _: &CommunicationValue) -> Result<(), OmikronError> {
            unreachable!()
        }

        async fn await_response(
            &self,
            request: &CommunicationValue,
            _: Duration,
        ) -> Result<CommunicationValue, OmikronError> {
            assert!(request.is_type(CommunicationType::GetRegister));
            Ok(self.response.clone().with_id(request.get_id()))
        }

        async fn reconnect(&self) -> Result<(), OmikronError> {
            unreachable!()
        }

        async fn rotate_identity(&self) -> Result<(), OmikronError> {
            unreachable!()
        }

        async fn is_connected(&self) -> bool {
            true
        }
    }

    #[test]
    fn validates_usernames_before_remote_registration() {
        assert!(valid_username("alice"));
        assert!(valid_username("fifteen_char_ok"));
        assert!(!valid_username(""));
        assert!(!valid_username("sixteen_chars_bad"));
        assert!(!valid_username("path/name"));
        assert!(!valid_username("line\nbreak"));
    }

    #[tokio::test]
    async fn uses_user_id_allocated_by_omega() {
        let client = RegistrationClient {
            response: CommunicationValue::new(CommunicationType::GetRegister)
                .add_typed_default(DataType::UserId, DataValue::SignedNumber(4_294_967_311))
                .add_typed_default(
                    DataType::RegisterId,
                    DataValue::Str("00000000-0000-4000-8000-000000000001".into()),
                ),
        };

        assert_eq!(
            request_user_id(&client).await.unwrap(),
            (4_294_967_311, "00000000-0000-4000-8000-000000000001".into())
        );
    }

    #[tokio::test]
    async fn rejects_registration_response_without_a_positive_user_id() {
        let client = RegistrationClient {
            response: CommunicationValue::new(CommunicationType::GetRegister)
                .add_typed_default(DataType::UserId, DataValue::SignedNumber(0)),
        };

        assert!(matches!(
            request_user_id(&client).await,
            Err(CreateUserError::InvalidResponse)
        ));
    }
}