dynamic-waas-sdk 0.0.3

//! Generic MPC keygen orchestration. Used by chain clients to implement
//! `create_wallet_account`. Mirrors `python/dynamic_wallet_sdk/wallet_client.py:keygen`.
//!
//! Flow:
//!   1. `init_keygen` (sync) — produce client-side `(keygen_id,
//!      keygen_secret)` pairs, one per client share.
//!   2. POST `/waas/create` with the `clientKeygenIds` — opens an SSE stream.
//!   3. SSE delivers `keygen_complete` with `roomId` + `serverKeygenIds`.
//!   4. While the stream stays open, run the MPC ceremony: each client
//!      share calls `EcdsaSigner::keygen` against the relay room.
//!   5. SSE delivers `ceremony_complete` with `walletId`.
//!   6. `derive_pubkey` (sync) from the first share's `secret_share`.
//!   7. `mark_key_shares_as_backed_up` (empty `locations` — flips the
//!      shares to "active" without Dynamic-side encryption).

use std::sync::Arc;

use dynamic_waas_sdk_core::{
    api::{CreateWalletKeygenReq, KeygenCompleteEvent},
    sse::{stream_sse_keygen, SseEventData},
    Error, Result, ServerKeyShare, ThresholdSignatureScheme, WalletProperties,
};
use dynamic_waas_sdk_mpc::{EcdsaSigner, KeygenId, RoomUuid};
use tracing::{debug, instrument};

use crate::client::DynamicWalletClient;
use crate::mpc_config::{threshold_wire, MpcSchemeConfig};

/// Output of [`run_keygen`] — the building block chain clients use to
/// implement `create_wallet_account`.
#[derive(Debug)]
pub struct KeygenOutput {
    /// Identity-only `WalletProperties`. Chain clients fill in
    /// `account_address` from the derived pubkey before returning.
    pub wallet_properties: WalletProperties,
    /// Server key shares the customer must persist in their vault.
    pub external_server_key_shares: Vec<ServerKeyShare>,
    /// Uncompressed 65-byte secp256k1 pubkey (0x04 prefix). Chain clients
    /// derive their address representation from this.
    pub raw_public_key_uncompressed: Vec<u8>,
    /// Compressed 33-byte secp256k1 pubkey.
    pub raw_public_key_compressed: Vec<u8>,
}

/// Options for [`run_keygen`].
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct KeygenOpts {
    pub chain_name: String,
    pub threshold_signature_scheme: ThresholdSignatureScheme,
    pub derivation_path: Vec<u32>,
    pub address_type: Option<String>,
    /// When true (the default), the customer's shares are encrypted with
    /// `password` and uploaded to Dynamic's backup store; the wallet is
    /// activated with `BackupLocation::Dynamic`. When false, no backup is
    /// uploaded — the wallet is activated with `BackupLocation::External`
    /// and the caller is responsible for vaulting the returned shares.
    /// Mirrors `back_up_to_dynamic` on the Node and Python SDKs (PR #856).
    pub back_up_to_dynamic: bool,
    /// Required when `back_up_to_dynamic=true` — used to AES-GCM-encrypt
    /// the customer's shares before upload. Ignored when
    /// `back_up_to_dynamic=false`.
    pub password: Option<String>,
}

impl KeygenOpts {
    pub fn new(
        chain_name: impl Into<String>,
        threshold_signature_scheme: ThresholdSignatureScheme,
        derivation_path: Vec<u32>,
    ) -> Self {
        Self {
            chain_name: chain_name.into(),
            threshold_signature_scheme,
            derivation_path,
            address_type: None,
            back_up_to_dynamic: true,
            password: None,
        }
    }

    #[must_use]
    pub fn with_address_type(mut self, t: impl Into<String>) -> Self {
        self.address_type = Some(t.into());
        self
    }

    /// Override `back_up_to_dynamic` (default `true`).
    #[must_use]
    pub fn with_back_up_to_dynamic(mut self, value: bool) -> Self {
        self.back_up_to_dynamic = value;
        self
    }

    /// Provide the password used to encrypt the customer's share for the
    /// Dynamic-side backup. Required iff `back_up_to_dynamic=true`.
    #[must_use]
    pub fn with_password(mut self, password: impl Into<String>) -> Self {
        self.password = Some(password.into());
        self
    }
}

/// Returns `Err(Error::InvalidArgument)` when `back_up_to_dynamic=true`
/// and no password was supplied. Mirrors
/// `_validate_password_for_backup` on the Python SDK so the failure
/// mode matches.
pub(crate) fn validate_password_for_backup(
    password: Option<&str>,
    back_up_to_dynamic: bool,
) -> Result<()> {
    if back_up_to_dynamic && password.unwrap_or("").is_empty() {
        return Err(Error::InvalidArgument(
            "password is required when back_up_to_dynamic=true".into(),
        ));
    }
    Ok(())
}

#[instrument(skip(client), fields(chain = %opts.chain_name))]
#[allow(clippy::too_many_lines)]
pub async fn run_keygen(client: &DynamicWalletClient, opts: KeygenOpts) -> Result<KeygenOutput> {
    if !client.is_authenticated() {
        return Err(Error::Authentication(crate::AUTH_REQUIRED_MSG.into()));
    }

    // Fail fast before the MPC ceremony spins up: if the caller asked
    // for Dynamic-side backup, the password must be present. Matches
    // PR #856 on Node + the Python validate_password_for_backup helper.
    validate_password_for_backup(opts.password.as_deref(), opts.back_up_to_dynamic)?;

    let scheme_cfg = MpcSchemeConfig::from(opts.threshold_signature_scheme);

    // 1. Client init keygen — synchronous.
    let signer = EcdsaSigner::new(client.base_mpc_relay_url().to_string());
    let mut init_results = Vec::with_capacity(scheme_cfg.client_threshold as usize);
    for _ in 0..scheme_cfg.client_threshold {
        init_results.push(signer.init_keygen()?);
    }
    let client_keygen_ids: Vec<String> = init_results
        .iter()
        .map(|r| r.keygen_id.as_str().to_owned())
        .collect();

    // 2. POST /waas/create — opens SSE stream.
    let body = CreateWalletKeygenReq {
        chain: opts.chain_name.clone(),
        client_keygen_ids: client_keygen_ids.clone(),
        threshold_signature_scheme: threshold_wire(opts.threshold_signature_scheme).to_string(),
        derivation_path: Some(opts.derivation_path.clone()),
        address_type: opts.address_type.clone(),
    };
    let response = client.api().create_wallet_keygen(&body).await?;

    // 3+4+5. Stream stays open during MPC ceremony.
    let host_url = client.base_mpc_relay_url().to_string();
    let init_results = Arc::new(init_results);
    let init_results_for_cb = Arc::clone(&init_results);
    let client_keygen_ids_for_cb = client_keygen_ids.clone();
    let scheme_n = scheme_cfg.n;
    let scheme_t = scheme_cfg.t;

    let (mpc_results, ceremony_data) = stream_sse_keygen(response, move |trigger| async move {
        let event: KeygenCompleteEvent = match trigger {
            SseEventData::Json(v) => serde_json::from_value(v).map_err(Error::from)?,
            SseEventData::Raw(s) => {
                return Err(Error::Sse(format!(
                    "keygen_complete payload was not JSON: {s}"
                )))
            }
        };
        debug!(room_id = %event.room_id, "running MPC keygen");

        let signer = EcdsaSigner::new(host_url);
        let room = RoomUuid::new(event.room_id);

        // For TWO_OF_TWO there's exactly one client share; loop kept for
        // future TWO_OF_THREE support (where client_threshold = 2).
        let mut results = Vec::with_capacity(init_results_for_cb.len());
        for (i, init) in init_results_for_cb.iter().enumerate() {
            let other_external_ids: Vec<KeygenId> = client_keygen_ids_for_cb
                .iter()
                .enumerate()
                .filter(|(j, _)| *j != i)
                .map(|(_, id)| KeygenId::new(id.clone()))
                .collect();
            let server_ids: Vec<KeygenId> = event
                .server_keygen_ids
                .iter()
                .map(|s| KeygenId::new(s.clone()))
                .collect();
            let all_others: Vec<KeygenId> =
                server_ids.into_iter().chain(other_external_ids).collect();

            let result = signer
                .keygen(&room, scheme_n, scheme_t, &init.keygen_secret, &all_others)
                .await?;
            results.push(result);
        }
        Ok::<_, Error>(results)
    })
    .await?;

    // 6. derive pubkey from the first share.
    let first = mpc_results
        .first()
        .ok_or(Error::Mpc(dynamic_waas_sdk_mpc::MpcError::Unknown))?;
    let signer = EcdsaSigner::new(client.base_mpc_relay_url().to_string());
    let (raw_uncompressed, raw_compressed) =
        signer.derive_pubkey(&first.secret_share, &opts.derivation_path)?;

    // 7. Build the (still-incomplete) WalletProperties — chain client fills
    // in account_address.
    let mut wallet_id = String::new();
    if let Some(SseEventData::Json(v)) = &ceremony_data {
        if let Some(id) = v.get("walletId").and_then(|x| x.as_str()) {
            id.clone_into(&mut wallet_id);
        }
    }
    let mut wp = WalletProperties::new(
        opts.chain_name.clone(),
        wallet_id.clone(),
        String::new(), // chain client fills this in from raw_uncompressed
    )
    .with_threshold(opts.threshold_signature_scheme)
    .with_derivation_path(opts.derivation_path.clone());

    // Build the ServerKeyShare list, one per client share.
    let key_shares: Vec<ServerKeyShare> = mpc_results
        .into_iter()
        .enumerate()
        .map(|(i, kr)| {
            ServerKeyShare::new(client_keygen_ids[i].clone(), kr.secret_share.into_string())
        })
        .collect();

    // Activate the wallet. The server requires
    // `mark_key_shares_as_backed_up` before signing — without it the
    // relay rejects sign requests with "Wallet is not active".
    //
    // Two paths:
    //  - back_up_to_dynamic=true (v1 default): encrypt shares with the
    //    customer password and upload to Dynamic's backup store
    //    (`BackupLocation::Dynamic`). Returns the new `KeyShareBackupInfo`
    //    so the caller can persist the pointers.
    //  - back_up_to_dynamic=false: skip the backup upload. Mark with
    //    `BackupLocation::External` only — the caller is responsible for
    //    vaulting the shares themselves.
    if !wallet_id.is_empty() {
        let signer = EcdsaSigner::new(client.base_mpc_relay_url().to_string());
        let keygen_id = signer
            .export_id(&dynamic_waas_sdk_mpc::SecretShare::from_string(
                key_shares[0].secret_share.clone(),
            ))?
            .into_string();
        if opts.back_up_to_dynamic {
            // Safe: validate_password_for_backup ensured this is Some.
            let password = opts.password.as_deref().unwrap_or("");
            let backup_info = crate::backup::run_backup_dynamic(
                client,
                &wallet_id,
                &key_shares,
                &keygen_id,
                password,
            )
            .await?;
            wp.external_server_key_shares_backup_info = Some(backup_info);
        } else {
            crate::backup::run_mark_external_no_backup(client, &wallet_id, &key_shares, &keygen_id)
                .await?;
        }
    }

    Ok(KeygenOutput {
        wallet_properties: wp,
        external_server_key_shares: key_shares,
        raw_public_key_uncompressed: raw_uncompressed,
        raw_public_key_compressed: raw_compressed,
    })
}