spark_rust/wallet/handlers/

create_tree.rs

use std::sync::Arc;

use bitcoin::{
    absolute::LockTime, key::Secp256k1, transaction::Version, Address, OutPoint, ScriptBuf,
    Transaction, TxIn, TxOut,
};
use parking_lot::RwLock;
use spark_protos::{
    common::SignatureIntent,
    spark::{
        FinalizeNodeSignaturesRequest, NodeSignatures, SigningJob, StartTreeCreationRequest,
        TreeNode, Utxo,
    },
};

use crate::{
    error::{CryptoError, SparkSdkError, TransactionError},
    signer::{default_signer::marshal_frost_commitments, traits::SparkSigner},
    wallet::{
        internal_handlers::traits::create_tree::CreateTreeInternalHandlers,
        leaf_manager::SparkLeaf,
        utils::{
            bitcoin::{
                bitcoin_tx_from_bytes, parse_public_key, serialize_bitcoin_transaction,
                sighash_from_tx,
            },
            sequence::initial_sequence,
        },
    },
    SparkSdk,
};

pub struct CreateTreeSdkResponse {
    pub nodes: Vec<TreeNode>,
}

impl<S: SparkSigner + Send + Sync + Clone + 'static> SparkSdk<S> {
    /// Creates a new tree of deposit addresses from either a deposit transaction or parent node.
    ///
    /// This function handles the creation of a new tree of deposit addresses by:
    /// 1. Generating the deposit address tree structure based on the split level
    /// 2. Finalizing the tree creation by signing all nodes and registering with the Spark network
    ///
    /// # Arguments
    ///
    /// * `deposit_transaction` - Optional Bitcoin transaction containing the deposit UTXO
    /// * `parent_node` - Optional parent tree node to create child nodes from
    /// * `vout` - The output index in the deposit transaction or parent node to use
    /// * `split_level` - The number of levels to split the tree into (determines number of leaves)
    /// * `parent_signing_public_key` - The public key of the parent node or deposit UTXO
    ///
    /// # Returns
    ///
    /// * `Ok(CreateTreeSdkResponse)` - Contains the finalized tree response including data for all created nodes, both branch and leaf nodes.
    /// * `Err(SparkSdkError)` - If there was an error during tree creation
    ///
    /// # Errors
    ///
    /// Returns [`SparkSdkError`] if:
    /// - Neither deposit transaction nor parent node is provided
    /// - Failed to generate deposit addresses
    /// - Failed to finalize tree creation with signatures
    /// - Network errors when communicating with Spark operators
    ///
    /// # Example
    ///
    /// ```
    /// # use spark_rust::SparkSdk;
    /// # use bitcoin::{Transaction, transaction::Version, absolute::LockTime};
    /// # async fn example(mut sdk: SparkSdk) -> Result<(), Box<dyn std::error::Error>> {
    /// // Assume we have a pre-existing deposit transaction
    /// let deposit_tx = Transaction { version: Version::TWO, lock_time: LockTime::ZERO, input: vec![], output: vec![] };
    /// let split_level = 2; // Creates 4 leaf nodes
    /// let parent_pubkey = vec![/* public key bytes */];
    ///
    /// let tree = sdk.create_tree(
    ///     Some(deposit_tx),
    ///     None,
    ///     0,
    ///     split_level,
    ///     parent_pubkey
    /// ).await?;
    /// # Ok(())
    /// # }
    /// ```
    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    async fn create_tree(
        &self,
        deposit_transaction: Option<Transaction>,
        parent_node: Option<Arc<RwLock<TreeNode>>>,
        vout: u32,
        split_level: u32,
        parent_signing_public_key: Vec<u8>,
    ) -> Result<CreateTreeSdkResponse, SparkSdkError> {
        // generate the tree
        let time_start = std::time::Instant::now();
        let deposit_address_response = self
            .generate_deposit_address_for_tree(
                deposit_transaction.clone(),
                parent_node.clone(),
                vout,
                parent_signing_public_key.clone(),
                split_level,
            )
            .await?;
        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "generate_deposit_address_for_tree");

        // finalize the tree
        let time_start = std::time::Instant::now();
        let finalize_tree_response = self
            .finalize_tree_creation(
                deposit_transaction,
                parent_node,
                vout,
                deposit_address_response.tree,
            )
            .await?;
        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "finalize_tree_creation");

        let finalize_tree_response = finalize_tree_response.finalize_tree_response;
        let nodes = finalize_tree_response.nodes;

        Ok(CreateTreeSdkResponse { nodes })
    }

    pub(crate) async fn create_tree_root(
        &self,
        signing_pubkey_bytes: Vec<u8>,
        verifying_pubkey_bytes: Vec<u8>,
        deposit_tx_bytes: Vec<u8>,
        vout: u32,
    ) -> Result<CreateTreeSdkResponse, SparkSdkError> {
        let deposit_tx = bitcoin_tx_from_bytes(&deposit_tx_bytes)?;

        let mut root_tx = Transaction {
            version: Version::TWO,
            lock_time: LockTime::ZERO,
            input: vec![],
            output: vec![],
        };

        let output = deposit_tx
            .output
            .get(vout as usize)
            .ok_or(SparkSdkError::from(TransactionError::InvalidDepositTx(
                format!("Invalid deposit transaction: {}", vout),
            )))?;

        let script = output.script_pubkey.clone();
        let amount = output.value;

        let txid = deposit_tx.compute_txid();
        root_tx.input.push(TxIn {
            previous_output: OutPoint { txid, vout },
            script_sig: ScriptBuf::new(),
            sequence: Default::default(),
            witness: Default::default(),
        });

        root_tx.output.push(TxOut {
            script_pubkey: script,
            value: amount,
        });

        // Get random signing commitment for root nonce
        let root_nonce_commitment = self.signer.new_frost_signing_noncepair()?;

        // Calculate sighash for root transaction
        let root_tx_sighash = sighash_from_tx(&root_tx, 0, output)?;

        // Create refund transaction
        let mut refund_tx = Transaction {
            version: Version::TWO,
            lock_time: LockTime::ZERO,
            input: vec![],
            output: vec![],
        };

        let root_txid = root_tx.compute_txid();
        refund_tx.input.push(TxIn {
            previous_output: OutPoint {
                txid: root_txid,
                vout: 0,
            },
            script_sig: ScriptBuf::new(),
            sequence: initial_sequence(),
            witness: Default::default(),
        });

        // Create refund output
        let secp = Secp256k1::new();
        let signing_pubkey = parse_public_key(&signing_pubkey_bytes)?;
        let refund_address = Address::p2tr(
            &secp,
            signing_pubkey.x_only_public_key().0,
            None,
            self.config.spark_config.network.to_bitcoin_network(),
        );
        let refund_script = refund_address.script_pubkey();

        refund_tx.output.push(TxOut {
            script_pubkey: refund_script,
            value: amount,
        });

        // Get random signing commitment for refund nonce
        let refund_nonce_commitment = self.signer.new_frost_signing_noncepair()?;

        // Calculate sighash for refund transaction
        let refund_tx_sighash = sighash_from_tx(&refund_tx, 0, &root_tx.output[0])?;

        // Get spark client
        let root_tx_bytes = serialize_bitcoin_transaction(&root_tx)?;
        let refund_tx_bytes = serialize_bitcoin_transaction(&refund_tx)?;

        // start tree creation
        let request = StartTreeCreationRequest {
            identity_public_key: self.get_spark_address()?.serialize().to_vec(),
            on_chain_utxo: Some(Utxo {
                vout,
                raw_tx: deposit_tx_bytes,
                network: self.config.spark_config.network.marshal_proto(),
            }),
            root_tx_signing_job: Some(SigningJob {
                raw_tx: root_tx_bytes.clone(),
                signing_public_key: signing_pubkey_bytes.clone(),
                signing_nonce_commitment: Some(marshal_frost_commitments(&root_nonce_commitment)?),
            }),
            refund_tx_signing_job: Some(SigningJob {
                raw_tx: refund_tx_bytes.clone(),
                signing_public_key: signing_pubkey_bytes.clone(),
                signing_nonce_commitment: Some(marshal_frost_commitments(
                    &refund_nonce_commitment,
                )?),
            }),
        };

        let tree_creation_response = self
            .config
            .spark_config
            .call_with_retry(
                request,
                |mut client, req| Box::pin(async move { client.start_tree_creation(req).await }),
                None,
            )
            .await?;

        let node_id = tree_creation_response
            .root_node_signature_shares
            .clone()
            .ok_or(SparkSdkError::from(
                CryptoError::MissingRootNodeSignatureShares,
            ))?
            .node_id
            .clone();

        let sigantures = self.signer.sign_root_creation(
            signing_pubkey_bytes.clone(),
            verifying_pubkey_bytes.clone(),
            root_tx_bytes.clone(),
            refund_tx_bytes.clone(),
            root_tx_sighash.to_vec(),
            refund_tx_sighash.to_vec(),
            root_nonce_commitment,
            refund_nonce_commitment,
            tree_creation_response,
        )?;

        let request = FinalizeNodeSignaturesRequest {
            intent: SignatureIntent::Creation as i32,
            node_signatures: vec![NodeSignatures {
                node_id,
                node_tx_signature: sigantures[0].clone(),
                refund_tx_signature: sigantures[1].clone(),
            }],
        };

        let finalize_node_signatures_response = self
            .config
            .spark_config
            .call_with_retry(
                request,
                |mut client, req| {
                    Box::pin(async move { client.finalize_node_signatures(req).await })
                },
                None,
            )
            .await?;

        let finalized_root_node = finalize_node_signatures_response.nodes[0].clone();

        self.leaf_manager
            .insert_leaves(vec![SparkLeaf::Bitcoin(finalized_root_node.clone())], false)?;

        let response = CreateTreeSdkResponse {
            nodes: vec![finalized_root_node],
        };

        Ok(response)
    }
}