spark_rust/wallet/internal_handlers/implementations/

timelock.rs

use std::collections::HashMap;

use bitcoin::secp256k1::PublicKey;
use bitcoin::{Sequence, Transaction};
use frost_secp256k1_tr_unofficial::round1::SigningCommitments;
use spark_protos::common::SignatureIntent;
use spark_protos::frost::{AggregateFrostRequest, FrostSigningJob};
use spark_protos::spark::{
    FinalizeNodeSignaturesRequest, NodeSignatures, RefreshTimelockRequest, SigningJob, TreeNode,
};
use tonic::{async_trait, Request};

use crate::constants::spark::TIME_LOCK_INTERVAL;
use crate::error::{network::NetworkError, validation::ValidationError, SparkSdkError};
use crate::signer::default_signer::{
    create_user_key_package, marshal_frost_commitments, marshal_frost_nonces,
};
use crate::signer::traits::SparkSigner;
use crate::wallet::internal_handlers::traits::timelock::TimelockInternalHandlers;
use crate::wallet::utils::bitcoin::{
    bitcoin_tx_from_bytes, serialize_bitcoin_transaction, sighash_from_tx,
};
use crate::wallet::utils::frost::frost_commitment_to_proto_commitment;
use crate::wallet::utils::sequence::{initial_sequence, next_sequence};
use crate::SparkSdk;

#[async_trait]
impl<S: SparkSigner + Send + Sync> TimelockInternalHandlers for SparkSdk<S> {
    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    async fn refresh_timelock_refund_tx(
        &self,
        leaf: &TreeNode,
        signing_public_key: &PublicKey,
    ) -> Result<(), SparkSdkError> {
        // create and encode the new refund tx
        let mut new_refund_tx = bitcoin_tx_from_bytes(&leaf.refund_tx)?;
        let curr_sequence = new_refund_tx.input[0].sequence;
        let next_sequence = next_sequence(curr_sequence.0);
        new_refund_tx.input[0].sequence = Sequence(next_sequence);
        let new_refund_tx_buf = serialize_bitcoin_transaction(&new_refund_tx)?;

        // prepare the signing job
        let commitment = self.signer.new_frost_signing_noncepair()?;
        let commitment_bytes = commitment.serialize().unwrap();
        let commitment_proto = frost_commitment_to_proto_commitment(&commitment)?;
        let signing_jobs = vec![SigningJob {
            signing_public_key: signing_public_key.serialize().to_vec(),
            raw_tx: new_refund_tx_buf,
            signing_nonce_commitment: Some(commitment_proto),
        }];

        // get the nonce, save the signing data
        let nonce = self
            .signer
            .sensitive_expose_nonces_from_commitments(&commitment_bytes)?;
        let signing_datas = vec![nonce];

        // call Spark to refresh timelock on the leaf
        let signing_jobs_len = signing_jobs.len();
        let mut spark_client = self.config.spark_config.get_spark_connection(None).await?;
        let mut request = Request::new(RefreshTimelockRequest {
            leaf_id: leaf.id.clone(),
            owner_identity_public_key: self
                .config
                .spark_config
                .ssp_identity_public_key
                .serialize()
                .to_vec(),
            signing_jobs: signing_jobs.clone(),
        });
        self.add_authorization_header_to_request(&mut request, None);
        let refresh_timelock_response = spark_client
            .refresh_timelock(request)
            .await
            .map_err(|status| SparkSdkError::from(NetworkError::Status(status)))?
            .into_inner();

        // validate response
        if signing_jobs_len != refresh_timelock_response.signing_results.len() {
            return Err(SparkSdkError::from(NetworkError::InvalidResponse));
        }

        // prepare signing and aggregation jobs
        let mut user_signing_jobs = Vec::new();
        let mut job_to_aggregate_request_map = HashMap::new();
        let mut job_to_node_id_map = HashMap::new();

        let signing_results = refresh_timelock_response.signing_results;
        for (i, signing_result) in signing_results.iter().enumerate() {
            let signing_data = &signing_datas[i];
            let signing_job = &signing_jobs[i];
            let refund_tx = bitcoin_tx_from_bytes(&signing_job.raw_tx)?;
            let node_tx = bitcoin_tx_from_bytes(&leaf.node_tx)?;

            let refund_tx_sighash = sighash_from_tx(&refund_tx, 0, &node_tx.output[0])?;

            let proto_signing_nonce = marshal_frost_nonces(signing_data)?;
            let proto_signing_commitment = marshal_frost_commitments(signing_data.commitments())?;
            let signing_secret_key = self.signer.sensitive_expose_secret_key_from_pubkey(
                &PublicKey::from_slice(&signing_job.signing_public_key)?,
                false,
            )?;
            let user_key_package = create_user_key_package(&signing_secret_key.secret_bytes());

            let signing_result_inner = signing_result.signing_result.as_ref().unwrap();
            let operator_commitments = &signing_result_inner.signing_nonce_commitments;

            let user_signing_job_id = uuid::Uuid::now_v7().to_string();

            user_signing_jobs.push(FrostSigningJob {
                job_id: user_signing_job_id.clone(),
                message: refund_tx_sighash.to_vec(),
                key_package: Some(user_key_package),
                verifying_key: signing_result.verifying_key.clone(),
                nonce: Some(proto_signing_nonce),
                commitments: operator_commitments.clone(),
                user_commitments: Some(proto_signing_commitment.clone()),
                adaptor_public_key: Default::default(),
            });

            job_to_aggregate_request_map.insert(
                user_signing_job_id.clone(),
                AggregateFrostRequest {
                    message: refund_tx_sighash.to_vec(),
                    signature_shares: signing_result_inner.signature_shares.clone(),
                    public_shares: signing_result_inner.public_keys.clone(),
                    verifying_key: signing_result.verifying_key.clone(),
                    commitments: operator_commitments.clone(),
                    user_commitments: Some(proto_signing_commitment),
                    user_public_key: signing_public_key.serialize().to_vec(),
                    adaptor_public_key: Default::default(),
                    user_signature_share: Default::default(),
                },
            );

            job_to_node_id_map.insert(user_signing_job_id, leaf.id.clone());
        }

        // sign with user
        let user_signatures = self.signer.sign_frost(user_signing_jobs)?;

        // aggregate signatures
        let mut node_signatures = Vec::new();
        for (job_id, user_signature) in user_signatures.results.iter() {
            let request = job_to_aggregate_request_map.get_mut(job_id).unwrap();
            request.user_signature_share = user_signature.signature_share.clone();

            let response = self.signer.aggregate_frost(request.clone())?;
            node_signatures.push(NodeSignatures {
                node_id: job_to_node_id_map[job_id].clone(),
                refund_tx_signature: response.signature,
                node_tx_signature: Default::default(),
            });
        }

        // finalize node signatures for the flow
        let mut request = Request::new(FinalizeNodeSignaturesRequest {
            intent: SignatureIntent::Refresh.into(),
            node_signatures,
        });
        self.add_authorization_header_to_request(&mut request, None);
        spark_client
            .finalize_node_signatures(request)
            .await
            .map_err(|status| SparkSdkError::from(NetworkError::Status(status)))?;

        Ok(())
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    async fn refresh_timelock_nodes(
        &self,
        nodes: &Vec<TreeNode>,
        parent_nodes: &Vec<TreeNode>,
        signing_public_key: &PublicKey,
    ) -> Result<(), SparkSdkError> {
        if nodes.is_empty() {
            return Err(SparkSdkError::from(ValidationError::InvalidInput {
                field: "no nodes to refresh timelock".to_string(),
            }));
        }

        let mut signing_jobs = Vec::with_capacity(nodes.len() + 1);
        let mut nonces = Vec::with_capacity(nodes.len() + 1);

        for (i, node) in nodes.iter().enumerate() {
            let mut new_node_tx = bitcoin_tx_from_bytes(&node.node_tx)?;
            if i == 0 {
                let curr_sequence = new_node_tx.input[0].sequence;
                new_node_tx.input[0].sequence = Sequence(next_sequence(curr_sequence.0));
            } else {
                new_node_tx.input[0].sequence = Sequence(TIME_LOCK_INTERVAL);
            }

            let (signing_job, signing_commitment) =
                self.signing_job_from_tx(&new_node_tx, signing_public_key)?;

            signing_jobs.push(signing_job);

            let commitment_bytes = signing_commitment.serialize().unwrap();
            let nonce = self
                .signer
                .sensitive_expose_nonces_from_commitments(&commitment_bytes)?;
            nonces.push(nonce);
        }

        // add one more job for the refund tx
        let leaf = &nodes[nodes.len() - 1];
        let mut new_refund_tx = bitcoin_tx_from_bytes(&leaf.refund_tx)?;
        new_refund_tx.input[0].sequence = initial_sequence(); // set the sequence to the initial sequence
        let (signing_job, signing_commitment) =
            self.signing_job_from_tx(&new_refund_tx, signing_public_key)?;
        signing_jobs.push(signing_job);
        let commitment_bytes = signing_commitment.serialize().unwrap();
        let refund_nonce = self
            .signer
            .sensitive_expose_nonces_from_commitments(&commitment_bytes)?;
        nonces.push(refund_nonce);

        // call Spark to refresh timelock on the nodes
        let mut spark_client = self.config.spark_config.get_spark_connection(None).await?;
        let mut request = Request::new(RefreshTimelockRequest {
            leaf_id: leaf.id.clone(),
            owner_identity_public_key: self
                .config
                .spark_config
                .ssp_identity_public_key
                .serialize()
                .to_vec(),
            signing_jobs: signing_jobs.clone(),
        });
        self.add_authorization_header_to_request(&mut request, None);

        let refresh_timelock_response = spark_client
            .refresh_timelock(request)
            .await
            .map_err(|status| SparkSdkError::from(NetworkError::Status(status)))?
            .into_inner();

        if signing_jobs.len() != refresh_timelock_response.signing_results.len() {
            return Err(SparkSdkError::from(ValidationError::InvalidInput {
                field: format!(
                    "number of signing jobs and signing results do not match: {} != {}",
                    signing_jobs.len(),
                    refresh_timelock_response.signing_results.len()
                ),
            }));
        }

        // sign and aggregate
        let mut user_signing_jobs = Vec::new();
        let mut job_to_aggregate_request_map = HashMap::new();
        let mut job_to_node_id_map = HashMap::new();
        let mut job_to_refund_map = HashMap::new();

        for (i, signing_result) in refresh_timelock_response.signing_results.iter().enumerate() {
            let nonce = nonces[i].clone();
            let signing_job = &signing_jobs[i];
            let raw_tx = bitcoin_tx_from_bytes(&signing_job.raw_tx)?;

            // get parent node for txout for sighash
            let (parent_node, node, refund, vout) = if i == nodes.len() {
                // Refund tx
                let node = &nodes[i - 1];
                (node, node, true, 0)
            } else {
                let node = &nodes[i];
                let parent_node = &parent_nodes[i];
                (parent_node, node, false, node.vout)
            };
            let parent_tx = bitcoin_tx_from_bytes(&parent_node.node_tx)?;
            let txout = parent_tx.output[vout as usize].clone();

            let raw_tx_sighash = sighash_from_tx(&raw_tx, 0, &txout)?;
            let proto_signing_nonce = marshal_frost_nonces(&nonce)?;
            let proto_signing_commitment = marshal_frost_commitments(nonce.commitments())?;
            let signing_secret_key = self.signer.sensitive_expose_secret_key_from_pubkey(
                &PublicKey::from_slice(&signing_job.signing_public_key)?,
                false,
            )?;
            let user_key_package = create_user_key_package(&signing_secret_key.secret_bytes());

            let signing_result_inner = signing_result.signing_result.as_ref().unwrap();
            let operator_commitments = &signing_result_inner.signing_nonce_commitments;

            let user_signing_job_id = uuid::Uuid::now_v7().to_string();
            user_signing_jobs.push(FrostSigningJob {
                job_id: user_signing_job_id.clone(),
                message: raw_tx_sighash.to_vec(),
                key_package: Some(user_key_package),
                verifying_key: signing_result.verifying_key.clone(),
                nonce: Some(proto_signing_nonce),
                commitments: operator_commitments.clone(),
                user_commitments: Some(proto_signing_commitment.clone()),
                adaptor_public_key: Default::default(),
            });

            job_to_aggregate_request_map.insert(
                user_signing_job_id.clone(),
                AggregateFrostRequest {
                    message: raw_tx_sighash.to_vec(),
                    signature_shares: signing_result_inner.signature_shares.clone(),
                    public_shares: signing_result_inner.public_keys.clone(),
                    verifying_key: signing_result.verifying_key.clone(),
                    commitments: operator_commitments.clone(),
                    user_commitments: Some(proto_signing_commitment),
                    user_public_key: signing_public_key.serialize().to_vec(),
                    adaptor_public_key: Default::default(),
                    user_signature_share: Default::default(),
                },
            );

            job_to_node_id_map.insert(user_signing_job_id.clone(), node.id.clone());
            job_to_refund_map.insert(user_signing_job_id, refund);
        }

        let user_signatures = self.signer.sign_frost(user_signing_jobs)?;
        let mut node_signatures = Vec::new();

        for (job_id, user_signature) in user_signatures.results.iter() {
            let request = job_to_aggregate_request_map.get_mut(job_id).unwrap();
            request.user_signature_share = user_signature.signature_share.clone();

            let response = self.signer.aggregate_frost(request.clone())?;

            if job_to_refund_map[job_id] {
                node_signatures.push(NodeSignatures {
                    node_id: job_to_node_id_map[job_id].clone(),
                    refund_tx_signature: response.signature,
                    node_tx_signature: Default::default(),
                });
            } else {
                node_signatures.push(NodeSignatures {
                    node_id: job_to_node_id_map[job_id].clone(),
                    node_tx_signature: response.signature,
                    refund_tx_signature: Default::default(),
                });
            }
        }

        // finalize node signatures for the flow
        let mut request = Request::new(FinalizeNodeSignaturesRequest {
            intent: SignatureIntent::Refresh.into(),
            node_signatures,
        });
        self.add_authorization_header_to_request(&mut request, None);
        spark_client
            .finalize_node_signatures(request)
            .await
            .map_err(|status| SparkSdkError::from(NetworkError::Status(status)))?;

        Ok(())
    }

    #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
    fn signing_job_from_tx(
        &self,
        new_tx: &Transaction,
        signing_public_key: &PublicKey,
    ) -> Result<(SigningJob, SigningCommitments), SparkSdkError> {
        let tx_bytes = serialize_bitcoin_transaction(new_tx)?;

        let signing_commitment = self.signer.new_frost_signing_noncepair()?;
        let proto_signing_commitment = marshal_frost_commitments(&signing_commitment)?;

        let signing_job = SigningJob {
            signing_public_key: signing_public_key.serialize().to_vec(),
            raw_tx: tx_bytes,
            signing_nonce_commitment: Some(proto_signing_commitment),
        };

        Ok((signing_job, signing_commitment))
    }
}