spark_rust/wallet/handlers/

timelock.rs

use hashbrown::HashMap;
use spark_protos::spark::{query_nodes_request::Source, QueryNodesRequest, TreeNodeIds};

use crate::{
    error::{SparkSdkError, WalletError},
    signer::traits::{derivation_path::SparkKeyType, SparkSigner},
    wallet::{
        internal_handlers::traits::timelock::TimelockInternalHandlers,
        leaf_manager::{SparkLeaf, SparkNodeStatus},
        utils::{bitcoin::bitcoin_tx_from_bytes, sequence::next_sequence},
    },
    SparkSdk,
};

impl<S: SparkSigner + Send + Sync + Clone + 'static> SparkSdk<S> {
    /// Refreshes timelock nodes in the wallet to handle timelock-related operations.
    ///
    /// This internal method updates the timelock state of Bitcoin leaves in the wallet.
    /// It handles the following timelock-related operations:
    ///
    /// 1. If a specific node ID is provided, it refreshes only that node
    /// 2. If no node ID is provided, it refreshes all Bitcoin leaves that are eligible for
    ///    timelock updates (based on sequence number progression)
    /// 3. Queries the Spark network for node information
    /// 4. Generates new signing keys for refreshed nodes
    /// 5. Updates the leaf manager with the refreshed nodes
    ///
    /// Timelock nodes are part of Spark's security model, where certain operations become
    /// valid only after a specific time has passed. This method ensures that the wallet
    /// properly tracks and updates the timelock state of its leaves.
    ///
    /// # Arguments
    ///
    /// * `node_id` - Optional ID of a specific node to refresh. If None, all eligible nodes are refreshed.
    ///
    /// # Returns
    ///
    /// * `Ok(())` - If the refresh operation was successful
    /// * `Err(SparkSdkError)` - If there was an error during the refresh process
    ///
    /// # Internal Usage
    ///
    /// This method is called by several public methods in the SDK:
    /// - `sync_wallet()` - To ensure all timelock nodes are up-to-date
    /// - `transfer()` - Before selecting leaves for transfer
    /// - `pay_lightning_invoice()` - Before selecting leaves for payment
    /// - Other operations that require fresh leaf state
    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    pub(crate) async fn refresh_timelock_nodes(
        &self,
        node_id: Option<String>,
    ) -> Result<(), SparkSdkError> {
        let mut nodes_to_refresh = vec![];
        let mut node_ids = vec![];

        match node_id {
            Some(leaf_id) => {
                let leaf_in_array = self
                    .leaf_manager
                    .filter_nodes_by_ids(&vec![leaf_id.clone()]);
                if leaf_in_array.is_empty() {
                    return Err(SparkSdkError::from(
                        WalletError::LeafNotFoundAfterOperation {
                            leaf_id: leaf_id.clone(),
                        },
                    ));
                }
                let leaf = leaf_in_array.first().unwrap();
                if !leaf.is_bitcoin() {
                    return Err(SparkSdkError::from(WalletError::LeafIsNotBitcoin {
                        leaf_id: leaf.get_id().clone(),
                    }));
                }

                nodes_to_refresh.push(leaf.get_tree_node()?);
                node_ids.push(leaf_id);
            }
            None => {
                let lock_callback = Some(Box::new(|leaf: &SparkLeaf| -> bool {
                    if !leaf.is_bitcoin() {
                        return false;
                    }

                    let leaf_node = leaf.get_tree_node().unwrap();
                    let refund_tx = bitcoin_tx_from_bytes(&leaf_node.refund_tx);
                    if refund_tx.is_err() {
                        return false;
                    }

                    let refund_tx = refund_tx.unwrap();
                    let current_sequence = refund_tx.input[0].sequence.0;
                    let next_sequence = next_sequence(current_sequence);

                    next_sequence == 0 || next_sequence > current_sequence // If 0 or if overflowed
                }) as Box<dyn Fn(&SparkLeaf) -> bool>);

                let (leaves, _unlocking_id) = self.leaf_manager.get_all_available_leaves(
                    lock_callback,
                    Some(SparkNodeStatus::RefreshTimelock),
                );
                node_ids = leaves.iter().map(|n| n.get_id().clone()).collect();
                nodes_to_refresh = leaves.iter().map(|n| n.get_tree_node().unwrap()).collect();
            }
        };

        if nodes_to_refresh.is_empty() {
            #[cfg(feature = "telemetry")]
            tracing::debug!("No leaf needs timelock refresh");

            return Ok(());
        }

        let request = QueryNodesRequest {
            source: Some(Source::NodeIds(TreeNodeIds { node_ids })),
            include_parents: true,
            network: self.config.spark_config.network.marshal_proto(),
        };

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

        let nodes_response = node_response.nodes;
        let mut nodes_map = HashMap::new();
        for (id, node) in nodes_response {
            nodes_map.insert(id, node);
        }

        for node in nodes_to_refresh {
            let parent_node_id = node.clone().parent_node_id.ok_or(SparkSdkError::from(
                WalletError::LeafHasNoParent {
                    leaf_id: node.id.clone(),
                },
            ))?;

            let parent_node = nodes_map.get(&parent_node_id).ok_or(SparkSdkError::from(
                WalletError::LeafParentNotFound {
                    leaf_id: node.id.clone(),
                },
            ))?;

            let signing_public_key = self.signer.new_secp256k1_keypair(
                node.id.clone(),
                SparkKeyType::BaseSigning,
                0,
                self.config.spark_config.network.to_bitcoin_network(),
            )?;

            let nodes = self
                .refresh_timelock_transfer_nodes(
                    &vec![node.clone()],
                    parent_node.clone(),
                    &signing_public_key,
                )
                .await?;

            if nodes.len() != 1 {
                return Err(SparkSdkError::from(
                    WalletError::PostRefreshNodeSignatureLengthMismatch,
                ));
            }

            let new_node = nodes.first().unwrap();

            // Update the leaves collection by removing the old node
            let leaf_ids_to_delete = vec![node.id.clone()];
            let leaves_to_insert = vec![SparkLeaf::Bitcoin(new_node.clone())];
            self.leaf_manager
                .delete_and_insert_leaves_atomically(&leaf_ids_to_delete, &leaves_to_insert)?;
        }

        Ok(())
    }
}