spark_sdk/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::SparkSdkError;
use crate::signer::traits::secp256k1::KeygenMethod;
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::internal_handlers::utils::frost_commitment_to_proto_commitment;
use crate::wallet::leaf_manager::LeafNode;
use parking_lot::RwLock;
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;

#[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.
    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 = self.signer.new_secp256k1_keypair(KeygenMethod::Random)?;

        // 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(target_pubkey, &left_pubkey, true)?;

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

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

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

    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,
        })
    }

    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(())
    }

    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 mut spark_client = self.config.spark_config.get_spark_connection(None).await?;
        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 mut deposit_address_tree =
            self.create_deposit_address_binary_tree(split_level, &parent_public_key)?;
        let duration = time_start.elapsed();
        println!(
            "[create_deposit_address_binary_tree] duration: {:?}",
            duration
        );

        // 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(&mut deposit_address_tree.tree)?;
        let duration = time_start.elapsed();
        println!(
            "[create_address_request_node_from_tree_nodes] duration: {:?}",
            duration
        );

        // 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::InvalidArgument(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 mut request = tonic::Request::new(PrepareTreeAddressRequest {
            user_identity_public_key: self.get_identity_public_key().to_vec(),
            node,
            source,
        });
        self.add_authorization_header_to_request(&mut request, None);

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

        let local_time = chrono::Local::now();
        println!(
            "Sending prepare_tree_address request at: {}",
            local_time.format("%Y-%m-%d %H:%M:%S")
        );
        let spark_tree_response_ = spark_client.prepare_tree_address(request).await?;
        let spark_tree_response = spark_tree_response_.into_inner();
        let duration = time_start.elapsed();
        println!("[prepare_tree_address] duration: {:?}", duration);

        // 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();
        println!("Applying address nodes to tree at time: {:?}", time_start);
        self.apply_address_nodes_to_tree(&mut root_in_vec, vec![response_address_node])?;
        let duration = time_start.elapsed();
        println!("[apply_address_nodes_to_tree] duration: {:?}", duration);

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

        Ok(GenerateDepositAddressForTreeSdkResponse { tree: root })
    }

    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::InvalidArgument(e.to_string()))?;

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

                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::InvalidArgument(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::InvalidArgument(e.to_string()))?;

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

                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,
        })
    }

    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_identity_public_key().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::InvalidArgument(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 mut tx_buf = parent_node.read().node_tx.clone();
            let ptx: Transaction = bitcoin::consensus::deserialize(&mut tx_buf).map_err(|_| {
                SparkSdkError::InvalidArgument("Failed to parse parent node_tx".into())
            })?;
            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::InvalidArgument(
                "No parent_tx or parent_node provided to create_tree".into(),
            ));
        };

        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();
        println!("[build_creation_nodes_from_tree] duration: {:?}", duration);

        request.node = Some(creation_node_response.creation_nodes.clone());
        let mut tonic_request = tonic::Request::new(request);
        self.add_authorization_header_to_request(&mut tonic_request, None);

        let mut spark_client = self.config.spark_config.get_spark_connection(None).await?;
        let time_start = std::time::Instant::now();
        let resp = spark_client.create_tree(tonic_request).await?;
        let duration = time_start.elapsed();
        println!("[create_tree] duration: {:?}", duration);

        let create_tree_result = resp.into_inner().node.ok_or_else(|| {
            SparkSdkError::InvalidArgument("Coordinator returned no creation node".into())
        })?;

        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,
            create_tree_result,
        )?;
        let duration = time_start.elapsed();
        println!("[sign_created_tree_in_bfs_order] duration: {:?}", duration);

        let mut spark_client = self.config.spark_config.get_spark_connection(None).await?;
        let mut spark_signatures_request = tonic::Request::new(FinalizeNodeSignaturesRequest {
            node_signatures: node_signatures.clone(),
            ..Default::default()
        });
        self.add_authorization_header_to_request(&mut spark_signatures_request, None);

        let time_start = std::time::Instant::now();
        let spark_signatures_response_ = spark_client
            .finalize_node_signatures(spark_signatures_request)
            .await?;
        let spark_signatures_response = spark_signatures_response_.into_inner();
        let duration = time_start.elapsed();
        println!("[finalize_node_signatures] duration: {:?}", duration);

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

        // Get all node ids
        let node_ids = spark_signatures_response
            .nodes
            .iter()
            .map(|node| node.id.clone())
            .collect::<Vec<_>>();

        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() {
            println!("node.tree_id: {:?}", node.tree_id);

            let node_id = node.id.clone();
            let i = i + starting_index;
            if node_id != node_signatures[i].clone().node_id {
                return Err(SparkSdkError::InvalidArgument("Node ID mismatch".into()));
            }

            let verifying_public_key = node.verifying_public_key.clone();
            let node_tx_transaction = node.node_tx.clone();
            let refund_tx_transaction = node.refund_tx.clone();
            let value = node.value;
            let node_tx_signature = node_signatures[i].node_tx_signature.clone();
            let refund_tx_signature = node_signatures[i].refund_tx_signature.clone();

            let _verifying_key = node.verifying_public_key.clone();

            // 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(LeafNode::new(
                node_id,
                "".to_string(),
                value,
                node.parent_node_id.clone(),
                node.vout,
                verifying_public_key,
                signing_public_keys[i].clone(),
                // node_tx_signature,
                // refund_tx_signature,
                node_tx_transaction,
                refund_tx_transaction,
                None,   // default will be the Bitcoin token public key
                vec![], // No revocation public key for tree creation
                vec![], // No token transaction hash for tree creation
            ));
        }

        self.leaf_manager
            .insert_leaves_in_batch(leaf_nodes_to_insert)?;

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