dynamic-waas-sdk 0.0.3

//! Ed25519 (`ExportableEd25519`) keygen orchestration. SVM equivalent of
//! `keygen.rs::run_keygen` for ECDSA.
//!
//! Flow is identical at the SSE / API layer:
//!   1. `init_keygen` (sync, `Ed25519Signer`) — generates 32-byte secret + base58 `keygen_id`
//!   2. POST /waas/create (SSE)
//!   3. SSE `keygen_complete` → run `Ed25519Signer.receive_key` against the relay
//!   4. SSE `ceremony_complete` → carries `walletId`
//!   5. `derive_pubkey` (sync) — 32-byte raw Ed25519 pubkey
//!   6. `mark_key_shares_as_backed_up` (location: external)

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::{Ed25519Signer, KeygenId, RoomUuid, SecretShare};
use tracing::{debug, instrument};

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

#[derive(Debug)]
pub struct KeygenOutputEd25519 {
    pub wallet_properties: WalletProperties,
    pub external_server_key_shares: Vec<ServerKeyShare>,
    pub raw_public_key: [u8; 32],
}

#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct KeygenOptsEd25519 {
    pub chain_name: String,
    pub threshold_signature_scheme: ThresholdSignatureScheme,
    pub derivation_path: Vec<u32>,
    /// See [`KeygenOpts::back_up_to_dynamic`](crate::KeygenOpts).
    pub back_up_to_dynamic: bool,
    /// See [`KeygenOpts::password`](crate::KeygenOpts).
    pub password: Option<String>,
}

impl KeygenOptsEd25519 {
    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,
            back_up_to_dynamic: true,
            password: None,
        }
    }

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

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

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

    crate::keygen::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.
    let signer = Ed25519Signer::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: None,
    };
    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 receive_key");

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

        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
                .receive_key(&room, scheme_n, scheme_t, &init.keygen_secret, &all_others)
                .await?;
            results.push(result);
        }
        Ok::<_, Error>(results)
    })
    .await?;

    // 6. derive pubkey.
    let first = mpc_results
        .first()
        .ok_or(Error::Mpc(dynamic_waas_sdk_mpc::MpcError::Unknown))?;
    let signer = Ed25519Signer::new(client.base_mpc_relay_url().to_string());
    let raw_pubkey = signer.derive_pubkey(&first.secret_share)?;

    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 account_address
    )
    .with_threshold(opts.threshold_signature_scheme)
    .with_derivation_path(opts.derivation_path.clone());

    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. See `keygen.rs` for the dual-path rationale —
    // same logic, but with Ed25519Signer::export_id for the keygen_id.
    if !wallet_id.is_empty() {
        let signer = Ed25519Signer::new(client.base_mpc_relay_url().to_string());
        let keygen_id = signer
            .export_id(&SecretShare::from_string(
                key_shares[0].secret_share.clone(),
            ))?
            .into_string();
        if opts.back_up_to_dynamic {
            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(KeygenOutputEd25519 {
        wallet_properties: wp,
        external_server_key_shares: key_shares,
        raw_public_key: raw_pubkey,
    })
}