spark_rust/wallet/internal_handlers/implementations/

create_tree.rs

use crate::common_types::types::Address;
use crate::common_types::types::Encodable;
use crate::common_types::types::Secp256k1;
use crate::common_types::types::Transaction;
use crate::common_types::types::TxIn;
use crate::common_types::types::TxOut;
use crate::common_types::types::Txid;
use crate::error::{validation::ValidationError, SparkSdkError};
use crate::signer::traits::SparkSigner;
use crate::wallet::client::SparkSdk;
use crate::wallet::internal_handlers::traits::create_tree::BuildCreationNodesFromTreeSdkResponse;
use crate::wallet::internal_handlers::traits::create_tree::CreateAddressRequestNodeFromTreeNodesSdkResponse;
use crate::wallet::internal_handlers::traits::create_tree::CreateDepositAddressBinaryTreeSdkResponse;
use crate::wallet::internal_handlers::traits::create_tree::CreateTreeInternalHandlers;
use crate::wallet::internal_handlers::traits::create_tree::DepositAddressTree;
use crate::wallet::internal_handlers::traits::create_tree::FinalizeTreeCreationSdkResponse;
use crate::wallet::internal_handlers::traits::create_tree::GenerateDepositAddressForTreeSdkResponse;
use crate::wallet::leaf_manager::SparkLeaf;
use crate::wallet::utils::frost::frost_commitment_to_proto_commitment;
use bitcoin::secp256k1::PublicKey;
use parking_lot::RwLock;
use spark_cryptography::derivation_path::SparkKeyType;
use spark_protos::spark::prepare_tree_address_request::Source as SourceProto;
use spark_protos::spark::AddressNode as AddressNodeProto;
use spark_protos::spark::AddressRequestNode as AddressRequestNodeProto;
use spark_protos::spark::CreateTreeRequest;
use spark_protos::spark::FinalizeNodeSignaturesRequest;
use spark_protos::spark::NodeOutput as NodeOutputProto;
use spark_protos::spark::PrepareTreeAddressRequest;
use spark_protos::spark::TreeNode as TreeNodeProto;
use spark_protos::spark::Utxo as UtxoProto;
use std::collections::VecDeque;
use std::str::FromStr;
use std::sync::Arc;
use tonic::async_trait;
use uuid::Uuid;

#[async_trait]
impl<S: SparkSigner + Send + Sync + Clone + 'static> CreateTreeInternalHandlers<S> for SparkSdk<S> {
    /// Creates a binary tree of deposit addresses.
    ///
    /// The tree is created by recursively splitting the target signing private key into two halves.
    /// Each node in the tree represents a deposit address, and the children of each node are the
    /// next level of the tree.
    ///
    /// # Arguments
    ///
    /// * `split_level` - The level of the tree to create.
    /// * `target_pubkey` - The public key to split into the tree.
    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    fn create_deposit_address_binary_tree(
        &self,
        split_level: u32,
        target_pubkey: &Vec<u8>,
    ) -> Result<CreateDepositAddressBinaryTreeSdkResponse, SparkSdkError> {
        if split_level == 0 {
            return Ok(CreateDepositAddressBinaryTreeSdkResponse { tree: vec![] });
        }

        // generate left pubkey
        let left_pubkey_uuid = Uuid::now_v7().to_string();
        let left_keypair = self.signer.derive_spark_key(
            left_pubkey_uuid,
            0,
            SparkKeyType::Deposit,
            self.config.spark_config.network.to_bitcoin_network(),
        )?;
        let left_pk = left_keypair.public_key();
        let left_pubkey = left_pk.serialize().to_vec();

        // left node
        let mut left_node = DepositAddressTree {
            address: None,
            verification_key: None,
            signing_public_key: left_pubkey.clone(),
            children: vec![],
        };

        // create left children recursively
        let left_children =
            self.create_deposit_address_binary_tree(split_level - 1, &left_pubkey)?;
        left_node.children = left_children.tree;

        // calculate right pubkey
        let right_pubkey = self.signer.subtract_secret_keys_given_pubkeys(
            &PublicKey::from_slice(target_pubkey)?,
            &PublicKey::from_slice(&left_pubkey)?,
            true,
        )?;

        // right node
        let mut right_node = DepositAddressTree {
            address: None,
            verification_key: None,
            signing_public_key: right_pubkey.serialize().to_vec(),
            children: vec![],
        };

        // create right children recursively
        let right_children = self.create_deposit_address_binary_tree(
            split_level - 1,
            &right_pubkey.serialize().to_vec(),
        )?;
        right_node.children = right_children.tree;

        Ok(CreateDepositAddressBinaryTreeSdkResponse {
            tree: vec![
                Arc::new(RwLock::new(left_node)),
                Arc::new(RwLock::new(right_node)),
            ],
        })
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    fn create_address_request_node_from_tree_nodes(
        &self,
        tree_nodes: &Vec<Arc<RwLock<DepositAddressTree>>>,
    ) -> Result<CreateAddressRequestNodeFromTreeNodesSdkResponse, SparkSdkError> {
        let mut results = Vec::<AddressRequestNodeProto>::new();

        for node in tree_nodes {
            let node = node.read();

            let address_request_node =
                self.create_address_request_node_from_tree_nodes(&node.children)?;
            let address_request_node = AddressRequestNodeProto {
                user_public_key: node.signing_public_key.clone(),
                children: address_request_node.address_request_nodes,
            };
            results.push(address_request_node);
        }

        Ok(CreateAddressRequestNodeFromTreeNodesSdkResponse {
            address_request_nodes: results,
        })
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    fn apply_address_nodes_to_tree(
        &self,
        tree: &mut Vec<Arc<RwLock<DepositAddressTree>>>,
        address_nodes: Vec<AddressNodeProto>,
    ) -> Result<(), SparkSdkError> {
        for (i, node) in tree.iter_mut().enumerate() {
            let mut node = node.write();
            let node_address_data = address_nodes[i].address.clone().unwrap();
            node.address = Some(node_address_data.address);
            node.verification_key = Some(node_address_data.verifying_key);

            if !node.children.is_empty() {
                self.apply_address_nodes_to_tree(
                    &mut node.children,
                    address_nodes[i].children.clone(),
                )?;
            }
        }

        Ok(())
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    async fn generate_deposit_address_for_tree(
        &self,
        parent_tx: Option<Transaction>,
        parent_node: Option<Arc<RwLock<TreeNodeProto>>>,
        vout: u32,
        parent_public_key: Vec<u8>,
        split_level: u32,
    ) -> Result<GenerateDepositAddressForTreeSdkResponse, SparkSdkError> {
        let network_proto = self.config.spark_config.network.marshal_proto();

        // 1. Create the binary tree given the user request
        let time_start = std::time::Instant::now();
        let deposit_address_tree =
            self.create_deposit_address_binary_tree(split_level, &parent_public_key)?;
        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "create_deposit_address_binary_tree");

        // If split_level = 0, len = 1. if split_level = 1, len = 1 + 2 = 3. If split_level = 2, len = 1 + 2 + 4 = 7.
        // This is because the tree is a binary tree, and the number of nodes is 2^split_level.
        // assert!(tree.len() == 2u32.pow(split_level) as usize);

        // 2. Create the address request nodes (in proto format) from the tree nodes
        let time_start = std::time::Instant::now();
        let address_nodes =
            self.create_address_request_node_from_tree_nodes(&deposit_address_tree.tree)?;
        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "create_address_request_node_from_tree_nodes");

        // 3. Send PrepareTreeAddressRequest to Spark. This is the first step of tree creation.
        let request_source = match parent_node {
            Some(parent_node) => {
                let node_output = NodeOutputProto {
                    node_id: parent_node.read().id.clone(),
                    vout,
                };
                SourceProto::ParentNodeOutput(node_output)
            }
            None => {
                let mut raw_tx = Vec::new();
                parent_tx
                    .as_ref()
                    .unwrap()
                    .consensus_encode(&mut raw_tx)
                    .map_err(|e| {
                        SparkSdkError::from(ValidationError::InvalidArgument {
                            argument: e.to_string(),
                        })
                    })?;
                let utxo = UtxoProto {
                    vout,
                    raw_tx,
                    network: network_proto,
                };
                SourceProto::OnChainUtxo(utxo)
            }
        };
        let source = Some(request_source);
        let request_node = AddressRequestNodeProto {
            user_public_key: parent_public_key.clone(),
            children: address_nodes.address_request_nodes,
        };
        let node = Some(request_node);

        let request_data = PrepareTreeAddressRequest {
            user_identity_public_key: self.get_spark_address()?.serialize().to_vec(),
            node,
            source,
        };

        let time_start = std::time::Instant::now();
        let spark_tree_response = self
            .config
            .spark_config
            .call_with_retry(
                request_data,
                |mut client, req| Box::pin(async move { client.prepare_tree_address(req).await }),
                None,
            )
            .await?;

        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "prepare_tree_address");

        // 4. Create the root node
        let response_address_node = spark_tree_response.node.unwrap();
        let root = DepositAddressTree {
            address: None,
            verification_key: None,
            signing_public_key: parent_public_key.clone(),
            children: deposit_address_tree.tree.clone(),
        };
        let root = Arc::new(RwLock::new(root));

        // 5. Apply the address nodes to the tree
        let mut root_in_vec = vec![root];
        let time_start = std::time::Instant::now();
        self.apply_address_nodes_to_tree(&mut root_in_vec, vec![response_address_node])?;
        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "apply_address_nodes_to_tree");

        // 6. Extract the root back from the vector and return it
        let root = root_in_vec.remove(0);

        Ok(GenerateDepositAddressForTreeSdkResponse { tree: root })
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    fn build_creation_nodes_from_tree(
        &self,
        parent_txid: Txid,
        txout_: &TxOut,
        vout: u32,
        root: Arc<RwLock<DepositAddressTree>>,
    ) -> Result<BuildCreationNodesFromTreeSdkResponse, SparkSdkError> {
        struct TreeNode {
            parent_txid: Txid,
            txout: TxOut,
            vout: u32,
            node: Arc<RwLock<DepositAddressTree>>,
        }

        let mut creation_node = spark_protos::spark::CreationNode::default();
        let mut queue = VecDeque::<(TreeNode, *mut spark_protos::spark::CreationNode)>::new();

        queue.push_back((
            TreeNode {
                parent_txid,
                txout: txout_.clone(),
                vout,
                node: root,
            },
            &mut creation_node,
        ));

        let network = self.config.spark_config.network.to_bitcoin_network();
        let secp = Secp256k1::new();

        while let Some((current, creation_ptr)) = queue.pop_front() {
            let creation_ref = unsafe { &mut *creation_ptr };
            let local_node = current.node.read();

            // let node_signing_key = local_node.verification_key.clone();
            // let user_signing_key = SecretKey::from_slice(&node_signing_key).unwrap();
            let user_verifying_key_ = local_node.signing_public_key.clone();
            let user_verifying_key =
                bitcoin::secp256k1::PublicKey::from_slice(&user_verifying_key_).unwrap();

            if !local_node.children.is_empty() {
                let tx = {
                    let tx_input = bitcoin::TxIn {
                        previous_output: bitcoin::OutPoint {
                            txid: current.parent_txid,
                            vout: current.vout,
                        },
                        script_sig: bitcoin::ScriptBuf::default(),
                        sequence: bitcoin::Sequence::ZERO,
                        witness: bitcoin::Witness::default(),
                    };

                    let mut output = Vec::new();
                    for child in local_node.children.iter() {
                        let child_address = child.read().address.clone().unwrap();
                        let child_address = Address::from_str(&child_address).unwrap();
                        let child_address = child_address.require_network(network).unwrap();
                        output.push(TxOut {
                            value: bitcoin::Amount::from_sat(
                                current.txout.value.to_sat() / local_node.children.len() as u64,
                            ),
                            script_pubkey: child_address.script_pubkey(),
                        });
                    }

                    bitcoin::Transaction {
                        version: bitcoin::transaction::Version::TWO,
                        lock_time: bitcoin::absolute::LockTime::ZERO,
                        input: vec![tx_input],
                        output,
                    }
                };

                let mut tx_buf = vec![];
                tx.consensus_encode(&mut tx_buf).map_err(|e| {
                    SparkSdkError::from(ValidationError::InvalidArgument {
                        argument: e.to_string(),
                    })
                })?;

                let commitment = self.signer.generate_frost_signing_commitments()?;

                let signing_job = spark_protos::spark::SigningJob {
                    signing_public_key: user_verifying_key.serialize().to_vec(),
                    raw_tx: tx_buf,
                    signing_nonce_commitment: Some(frost_commitment_to_proto_commitment(
                        &commitment,
                    )?),
                };

                creation_ref.node_tx_signing_job = Some(signing_job);
                creation_ref.children = vec![Default::default(); local_node.children.len()];

                let txid = tx.compute_txid();
                for (i, child) in local_node.children.iter().enumerate() {
                    queue.push_back((
                        TreeNode {
                            parent_txid: txid,
                            txout: tx.output[i].clone(),
                            vout: i as u32,
                            node: child.clone(),
                        },
                        &mut creation_ref.children[i] as *mut _,
                    ));
                }
            } else {
                let aggregated_address = local_node.address.clone().unwrap();
                let aggregated_address = Address::from_str(&aggregated_address).unwrap();
                let aggregated_address = aggregated_address.require_network(network).unwrap();

                let node_tx = {
                    let input = bitcoin::TxIn {
                        previous_output: bitcoin::OutPoint {
                            txid: current.parent_txid,
                            vout: current.vout,
                        },
                        script_sig: bitcoin::ScriptBuf::default(),
                        // sequence: bitcoin::Sequence::ZERO,
                        sequence: bitcoin::Sequence::ZERO,
                        witness: bitcoin::Witness::default(),
                    };

                    bitcoin::Transaction {
                        version: bitcoin::transaction::Version::TWO,
                        lock_time: bitcoin::absolute::LockTime::ZERO,
                        input: vec![input],
                        output: vec![TxOut {
                            value: current.txout.value,
                            script_pubkey: aggregated_address.script_pubkey(),
                        }],
                    }
                };

                let mut node_tx_buf = vec![];
                node_tx.consensus_encode(&mut node_tx_buf).map_err(|e| {
                    SparkSdkError::from(ValidationError::InvalidArgument {
                        argument: e.to_string(),
                    })
                })?;

                let refund_tx = {
                    let user_self_xonly = user_verifying_key.x_only_public_key().0;
                    let user_self_address = Address::p2tr(&secp, user_self_xonly, None, network);

                    bitcoin::Transaction {
                        version: bitcoin::transaction::Version::TWO,
                        lock_time: bitcoin::absolute::LockTime::ZERO,
                        input: vec![TxIn {
                            previous_output: bitcoin::OutPoint {
                                txid: node_tx.compute_txid(),
                                vout: 0,
                            },
                            script_sig: bitcoin::ScriptBuf::default(),
                            // TODO: this must be the default sequence. For tree creation here, we can set it to MAX, since this is an SSP feature and is unlikely to be used here. Yet, this will affect unilateral exits.
                            sequence: bitcoin::Sequence::MAX,
                            witness: bitcoin::Witness::default(),
                        }],
                        output: vec![TxOut {
                            value: current.txout.value,
                            script_pubkey: user_self_address.script_pubkey(),
                        }],
                    }
                };

                let mut refund_tx_buf = vec![];
                refund_tx
                    .consensus_encode(&mut refund_tx_buf)
                    .map_err(|e| {
                        SparkSdkError::from(ValidationError::InvalidArgument {
                            argument: e.to_string(),
                        })
                    })?;

                let node_commitment = self.signer.generate_frost_signing_commitments()?;
                let refund_commitment = self.signer.generate_frost_signing_commitments()?;

                creation_ref.node_tx_signing_job = Some(spark_protos::spark::SigningJob {
                    signing_public_key: user_verifying_key.serialize().to_vec(),
                    raw_tx: node_tx_buf,
                    signing_nonce_commitment: Some(frost_commitment_to_proto_commitment(
                        &node_commitment,
                    )?),
                });

                creation_ref.refund_tx_signing_job = Some(spark_protos::spark::SigningJob {
                    signing_public_key: user_verifying_key.serialize().to_vec(),
                    raw_tx: refund_tx_buf,
                    signing_nonce_commitment: Some(frost_commitment_to_proto_commitment(
                        &refund_commitment,
                    )?),
                });
            }
        }

        Ok(BuildCreationNodesFromTreeSdkResponse {
            creation_nodes: creation_node,
        })
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    async fn finalize_tree_creation(
        &self,
        parent_tx: Option<Transaction>,
        parent_node: Option<Arc<RwLock<TreeNodeProto>>>,
        vout: u32,
        root: Arc<RwLock<DepositAddressTree>>,
    ) -> Result<FinalizeTreeCreationSdkResponse, SparkSdkError> {
        let mut request = CreateTreeRequest {
            user_identity_public_key: self.get_spark_address()?.serialize().to_vec(),
            ..Default::default()
        };

        let final_parent_tx = if let Some(ptx) = parent_tx {
            let mut raw_tx = Vec::new();
            ptx.consensus_encode(&mut raw_tx).map_err(|e| {
                SparkSdkError::from(ValidationError::InvalidArgument {
                    argument: e.to_string(),
                })
            })?;

            request.source = Some(
                spark_protos::spark::create_tree_request::Source::OnChainUtxo(
                    spark_protos::spark::Utxo {
                        // txid,
                        vout,
                        raw_tx,
                        network: self.config.spark_config.network.marshal_proto(),
                    },
                ),
            );
            ptx
        } else if let Some(parent_node) = parent_node {
            let tx_buf = parent_node.read().node_tx.clone();
            let ptx: Transaction = bitcoin::consensus::deserialize(&tx_buf).map_err(|_| {
                SparkSdkError::from(ValidationError::InvalidArgument {
                    argument: "Failed to parse parent node_tx".to_string(),
                })
            })?;
            let node_id = parent_node.read().id.clone();
            request.source = Some(
                spark_protos::spark::create_tree_request::Source::ParentNodeOutput(
                    spark_protos::spark::NodeOutput { node_id, vout },
                ),
            );
            ptx
        } else {
            return Err(SparkSdkError::from(ValidationError::InvalidArgument {
                argument: "No parent_tx or parent_node provided to create_tree".to_string(),
            }));
        };

        let parent_txid = final_parent_tx.compute_txid();
        let time_start = std::time::Instant::now();
        let creation_node_response = self.build_creation_nodes_from_tree(
            parent_txid,
            &final_parent_tx.output[vout as usize],
            vout,
            root.clone(),
        )?;
        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "build_creation_nodes_from_tree");

        request.node = Some(creation_node_response.creation_nodes.clone());

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

        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "create_tree");

        let tree_node = spark_tree_response.node.ok_or_else(|| {
            SparkSdkError::from(ValidationError::InvalidArgument {
                argument: "Coordinator returned no creation node".to_string(),
            })
        })?;

        let time_start = std::time::Instant::now();
        let (node_signatures, signing_public_keys) = self.signer.sign_created_tree_in_bfs_order(
            final_parent_tx,
            vout,
            root,
            creation_node_response.creation_nodes,
            tree_node,
        )?;
        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "sign_created_tree_in_bfs_order");
        let spark_signatures_request = FinalizeNodeSignaturesRequest {
            node_signatures: node_signatures.clone(),
            ..Default::default()
        };

        let time_start = std::time::Instant::now();

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

        let duration = time_start.elapsed();
        #[cfg(feature = "telemetry")]
        tracing::debug!(duration = ?duration, "finalize_node_signatures");

        // Slice the array to get the second half (including the center)
        let starting_index = spark_signatures_response.nodes.len() / 2;

        let leaf_nodes = spark_signatures_response
            .nodes
            .iter()
            .skip(spark_signatures_response.nodes.len() / 2)
            .collect::<Vec<_>>();

        let mut leaf_nodes_to_insert = vec![];
        for (i, node) in leaf_nodes.iter().enumerate() {
            let node_id = node.id.clone();
            let i = i + starting_index;
            if node_id != node_signatures[i].clone().node_id {
                return Err(SparkSdkError::from(ValidationError::InvalidArgument {
                    argument: "Node ID mismatch".to_string(),
                }));
            }

            // TODO: this is an error-prone approach. For roots, this approach is correct.
            // For splits, this approach is incorrect because now it should be strictly the right half, excluding the center.

            // TODO: For the parents, add aggregated.
            leaf_nodes_to_insert.push(SparkLeaf::Bitcoin((*node).clone()));
        }

        self.leaf_manager
            .insert_leaves(leaf_nodes_to_insert, false)?;

        Ok(FinalizeTreeCreationSdkResponse {
            finalize_tree_response: spark_signatures_response,
            signing_public_keys,
        })
    }
}