clementine_core/builder/transaction/
txhandler.rs

1//! # Transaction Handler Module
2//!
3//! This module defines the [`TxHandler`] abstraction, which wraps a protocol transaction and its metadata.
4//! Metadata includes taproot scripts and protocol specific data to enable signing of the transactions.
5//! [`TxHandlerBuilder`] is used to create [`TxHandler`]s.
6//!
7
8use super::input::{SpendableTxIn, SpentTxIn, UtxoVout};
9use super::output::UnspentTxOut;
10use crate::builder::script::SpendPath;
11use crate::builder::sighash::{PartialSignatureInfo, SignatureInfo};
12use crate::builder::transaction::deposit_signature_owner::{DepositSigKeyOwner, EntityType};
13use crate::rpc::clementine::tagged_signature::SignatureId;
14#[cfg(test)]
15use bitcoin::sighash::Annex;
16use bitcoin::sighash::SighashCache;
17use bitcoin::taproot::{self, LeafVersion};
18use bitcoin::transaction::Version;
19use bitcoin::{absolute, OutPoint, Script, Sequence, TapNodeHash, Transaction, Witness};
20use bitcoin::{TapLeafHash, TapSighash, TapSighashType, TxOut, Txid};
21use clementine_errors::BridgeError;
22use clementine_errors::{TransactionType, TxError};
23use eyre::{Context, OptionExt};
24use std::collections::BTreeMap;
25use std::marker::PhantomData;
26
27pub const DEFAULT_SEQUENCE: Sequence = Sequence::ENABLE_RBF_NO_LOCKTIME;
28
29#[derive(Debug, Clone)]
30/// Handler for protocol transactions, wrapping inputs, outputs, and cached transaction data.
31pub struct TxHandler<T: State = Unsigned> {
32    transaction_type: TransactionType,
33    txins: Vec<SpentTxIn>,
34    txouts: Vec<UnspentTxOut>,
35
36    /// Cached and immutable, same as other fields
37    cached_tx: bitcoin::Transaction,
38    cached_txid: bitcoin::Txid,
39
40    /// Annex bytes for test-only large tx scenarios.
41    #[cfg(test)]
42    test_annex: Option<Vec<u8>>,
43
44    phantom: PhantomData<T>,
45}
46
47pub trait State: Clone + std::fmt::Debug {}
48
49// #[derive(Debug, Clone)]
50// pub struct PartialInputs;
51#[derive(Debug, Clone)]
52/// Marker type for signed transactions.
53pub struct Signed;
54#[derive(Debug, Clone)]
55/// Marker type for unsigned transactions.
56pub struct Unsigned;
57
58// impl State for PartialInputs {}
59impl State for Unsigned {}
60impl State for Signed {}
61pub type SighashCalculator<'a> =
62    Box<dyn Fn(TapSighashType) -> Result<TapSighash, BridgeError> + 'a>;
63
64impl<T: State> TxHandler<T> {
65    /// Returns a spendable input for the specified output index in this transaction.
66    ///
67    /// # Arguments
68    /// * `vout` - The protocol-specific output index.
69    ///
70    /// # Returns
71    /// A [`SpendableTxIn`] for the specified output, or a [`BridgeError`] if not found.
72    pub fn get_spendable_output(&self, vout: UtxoVout) -> Result<SpendableTxIn, BridgeError> {
73        let idx = vout.get_vout();
74        let txout = self
75            .txouts
76            .get(idx as usize)
77            .ok_or_else(|| eyre::eyre!("Could not find output {idx} in transaction"))?;
78        Ok(SpendableTxIn::new(
79            OutPoint {
80                txid: self.cached_txid,
81                vout: idx,
82            },
83            txout.txout().clone(),
84            txout.scripts().clone(),
85            txout.spendinfo().clone(),
86        ))
87    }
88
89    /// Returns the Taproot merkle root of the specified input, if available.
90    ///
91    /// # Arguments
92    /// * `idx` - The input index.
93    ///
94    /// # Returns
95    /// The Taproot merkle root, or a [`BridgeError`] if not found.
96    pub fn get_merkle_root_of_txin(&self, idx: usize) -> Result<Option<TapNodeHash>, BridgeError> {
97        let txin = self
98            .txins
99            .get(idx)
100            .ok_or(TxError::TxInputNotFound)?
101            .get_spendable();
102        let merkle_root = txin
103            .get_spend_info()
104            .as_ref()
105            .ok_or(eyre::eyre!(
106                "Spend info not found for requested txin in get_merkle_root_of_txin"
107            ))?
108            .merkle_root();
109        Ok(merkle_root)
110    }
111
112    /// Returns the signature ID for the specified input.
113    ///
114    /// # Arguments
115    /// * `idx` - The input index.
116    ///
117    /// # Returns
118    /// The signature ID, or a [`BridgeError`] if not found.
119    pub fn get_signature_id(&self, idx: usize) -> Result<SignatureId, BridgeError> {
120        let txin = self.txins.get(idx).ok_or(TxError::TxInputNotFound)?;
121        Ok(txin.get_signature_id())
122    }
123
124    /// Returns the protocol transaction type for this handler.
125    pub fn get_transaction_type(&self) -> TransactionType {
126        self.transaction_type
127    }
128
129    /// Returns a reference to the cached Bitcoin transaction.
130    pub fn get_cached_tx(&self) -> &Transaction {
131        &self.cached_tx
132    }
133
134    /// Returns a reference to the cached transaction ID.
135    pub fn get_txid(&self) -> &Txid {
136        // Not sure if this should be public
137        &self.cached_txid
138    }
139
140    #[cfg(test)]
141    pub fn set_test_annex(&mut self, annex: Option<Vec<u8>>) {
142        self.test_annex = annex;
143    }
144
145    /// Returns a lambda function that calculates the sighash for the specified input, given the sighash type.
146    ///
147    /// # Arguments
148    /// * `idx` - The input index.
149    ///
150    /// # Returns
151    /// A lambda function that calculates the sighash for the specified input, given the sighash type.
152    fn get_sighash_calculator(
153        &self,
154        idx: usize,
155    ) -> impl Fn(TapSighashType) -> Result<TapSighash, BridgeError> + '_ {
156        move |sighash_type: TapSighashType| -> Result<TapSighash, BridgeError> {
157            match self.txins[idx].get_spend_path() {
158                SpendPath::KeySpend => self.calculate_pubkey_spend_sighash(idx, sighash_type),
159                SpendPath::ScriptSpend(script_idx) => {
160                    self.calculate_script_spend_sighash_indexed(idx, script_idx, sighash_type)
161                }
162                SpendPath::Unknown => Err(TxError::SpendPathNotSpecified.into()),
163            }
164        }
165    }
166
167    /// Signs all **unsigned** transaction inputs using the provided signer function.
168    ///
169    /// This function will skip all transaction inputs that already have a witness.
170    ///
171    /// # Arguments
172    /// * `signer` - A function that returns an optional witness for transaction inputs or returns an error
173    ///   if the signing fails. The function takes the input idx, input object, and a sighash calculator closure.
174    ///
175    /// # Returns
176    /// * `Ok(())` if signing is successful
177    /// * `Err(BridgeError)` if signing fails
178    pub fn sign_txins(
179        &mut self,
180        mut signer: impl for<'a> FnMut(
181            usize,
182            &'a SpentTxIn,
183            SighashCalculator<'a>,
184        ) -> Result<Option<Witness>, BridgeError>,
185    ) -> Result<(), BridgeError> {
186        for idx in 0..self.txins.len() {
187            let calc_sighash = Box::new(self.get_sighash_calculator(idx));
188            if self.txins[idx].get_witness().is_some() {
189                continue;
190            }
191
192            if let Some(witness) = signer(idx, &self.txins[idx], calc_sighash)
193                .wrap_err_with(|| format!("Failed to sign input {idx}"))?
194            {
195                self.cached_tx.input[idx].witness = witness.clone();
196                self.txins[idx].set_witness(witness);
197            }
198        }
199        Ok(())
200    }
201
202    /// Calculates the Taproot sighash for a key spend input for the given input and sighash type.
203    ///
204    /// # Arguments
205    /// * `txin_index` - The input index.
206    /// * `sighash_type` - The Taproot sighash type.
207    ///
208    /// # Returns
209    /// The calculated Taproot sighash, or a [`BridgeError`] if calculation fails.
210    pub fn calculate_pubkey_spend_sighash(
211        &self,
212        txin_index: usize,
213        sighash_type: TapSighashType,
214    ) -> Result<TapSighash, BridgeError> {
215        let prevouts_vec: Vec<&TxOut> = self
216            .txins
217            .iter()
218            .map(|s| s.get_spendable().get_prevout())
219            .collect();
220        let mut sighash_cache: SighashCache<&bitcoin::Transaction> =
221            SighashCache::new(&self.cached_tx);
222        let prevouts = match sighash_type {
223            TapSighashType::SinglePlusAnyoneCanPay
224            | TapSighashType::AllPlusAnyoneCanPay
225            | TapSighashType::NonePlusAnyoneCanPay => {
226                bitcoin::sighash::Prevouts::One(txin_index, prevouts_vec[txin_index])
227            }
228            _ => bitcoin::sighash::Prevouts::All(&prevouts_vec),
229        };
230
231        let sig_hash = sighash_cache
232            .taproot_key_spend_signature_hash(txin_index, &prevouts, sighash_type)
233            .wrap_err("Failed to calculate taproot sighash for key spend")?;
234
235        #[cfg(test)]
236        {
237            if let Some(ref annex_bytes) = self.test_annex {
238                if matches!(
239                    self.txins[txin_index]
240                        .get_signature_id()
241                        .get_deposit_sig_owner()?,
242                    DepositSigKeyOwner::Own(_)
243                ) {
244                    let annex =
245                        Annex::new(annex_bytes).map_err(|e| eyre::eyre!("Invalid annex: {e:?}"))?;
246                    return Ok(sighash_cache
247                        .taproot_signature_hash(
248                            txin_index,
249                            &prevouts,
250                            Some(annex),
251                            None,
252                            sighash_type,
253                        )
254                        .wrap_err(
255                            "Failed to calculate taproot sighash for key spend with annex",
256                        )?);
257                }
258            }
259        }
260
261        Ok(sig_hash)
262    }
263
264    /// Calculates the Taproot sighash for a script spend input by script index.
265    ///
266    /// # Arguments
267    /// * `txin_index` - The input index.
268    /// * `spend_script_idx` - The script index in the input's script list.
269    /// * `sighash_type` - The Taproot sighash type.
270    ///
271    /// # Returns
272    /// The calculated Taproot sighash, or a [`BridgeError`] if calculation fails.
273    pub fn calculate_script_spend_sighash_indexed(
274        &self,
275        txin_index: usize,
276        spend_script_idx: usize,
277        sighash_type: TapSighashType,
278    ) -> Result<TapSighash, BridgeError> {
279        let script = self
280            .txins
281            .get(txin_index)
282            .ok_or(TxError::TxInputNotFound)?
283            .get_spendable()
284            .get_scripts()
285            .get(spend_script_idx)
286            .ok_or(TxError::ScriptNotFound(spend_script_idx))?
287            .to_script_buf();
288
289        self.calculate_script_spend_sighash(txin_index, &script, sighash_type)
290    }
291
292    /// Calculates the Taproot sighash for a script spend input by script.
293    ///
294    /// # Arguments
295    /// * `txin_index` - The input index.
296    /// * `spend_script` - The script being spent.
297    /// * `sighash_type` - The Taproot sighash type.
298    ///
299    /// # Returns
300    /// The calculated Taproot sighash, or a [`BridgeError`] if calculation fails.
301    pub fn calculate_script_spend_sighash(
302        &self,
303        txin_index: usize,
304        spend_script: &Script,
305        sighash_type: TapSighashType,
306    ) -> Result<TapSighash, BridgeError> {
307        let prevouts_vec: Vec<&TxOut> = self
308            .txins
309            .iter()
310            .map(|s| s.get_spendable().get_prevout())
311            .collect();
312        let mut sighash_cache: SighashCache<&bitcoin::Transaction> =
313            SighashCache::new(&self.cached_tx);
314
315        let prevouts = match sighash_type {
316            TapSighashType::SinglePlusAnyoneCanPay
317            | TapSighashType::AllPlusAnyoneCanPay
318            | TapSighashType::NonePlusAnyoneCanPay => {
319                bitcoin::sighash::Prevouts::One(txin_index, prevouts_vec[txin_index])
320            }
321            _ => bitcoin::sighash::Prevouts::All(&prevouts_vec),
322        };
323        let leaf_hash = TapLeafHash::from_script(spend_script, LeafVersion::TapScript);
324        let sig_hash = sighash_cache
325            .taproot_script_spend_signature_hash(txin_index, &prevouts, leaf_hash, sighash_type)
326            .wrap_err("Failed to calculate taproot sighash for script spend")?;
327
328        #[cfg(test)]
329        {
330            if let Some(ref annex_bytes) = self.test_annex {
331                if matches!(
332                    self.txins[txin_index]
333                        .get_signature_id()
334                        .get_deposit_sig_owner()?,
335                    DepositSigKeyOwner::Own(_)
336                ) {
337                    let annex =
338                        Annex::new(annex_bytes).map_err(|e| eyre::eyre!("Invalid annex: {e:?}"))?;
339                    return Ok(sighash_cache
340                        .taproot_signature_hash(
341                            txin_index,
342                            &prevouts,
343                            Some(annex),
344                            Some((leaf_hash, 0xFFFFFFFF)),
345                            sighash_type,
346                        )
347                        .wrap_err(
348                            "Failed to calculate taproot sighash for script spend with annex",
349                        )?);
350                }
351            }
352        }
353
354        Ok(sig_hash)
355    }
356
357    /// Calculates the sighash for the specified input, based on its spend path stored inside [`SpentTxIn`].
358    ///
359    /// # Arguments
360    /// * `txin_index` - The input index.
361    /// * `sighash_type` - The Taproot sighash type.
362    ///
363    /// # Returns
364    /// The calculated Taproot sighash, or a [`BridgeError`] if calculation fails.
365    pub fn calculate_sighash_txin(
366        &self,
367        txin_index: usize,
368        sighash_type: TapSighashType,
369    ) -> Result<TapSighash, BridgeError> {
370        match self.txins[txin_index].get_spend_path() {
371            SpendPath::ScriptSpend(idx) => {
372                self.calculate_script_spend_sighash_indexed(txin_index, idx, sighash_type)
373            }
374            SpendPath::KeySpend => self.calculate_pubkey_spend_sighash(txin_index, sighash_type),
375            SpendPath::Unknown => Err(TxError::MissingSpendInfo.into()),
376        }
377    }
378
379    /// Calculates sighashes for all shared inputs for a given entity type.
380    ///
381    /// # Arguments
382    /// * `needed_entity` - The entity type (operator, verifier, etc.).
383    /// * `partial_signature_info` - Partial signature info for the entity.
384    ///
385    /// # Returns
386    /// A vector of (sighash, signature info) pairs, or a [`BridgeError`] if calculation fails.
387    pub fn calculate_shared_txins_sighash(
388        &self,
389        needed_entity: EntityType,
390        partial_signature_info: PartialSignatureInfo,
391    ) -> Result<Vec<(TapSighash, SignatureInfo)>, BridgeError> {
392        let mut sighashes = Vec::with_capacity(self.txins.len());
393        for idx in 0..self.txins.len() {
394            let sig_id = self.txins[idx].get_signature_id();
395            let spend_data = self.txins[idx].get_tweak_data();
396            let sig_owner = sig_id.get_deposit_sig_owner()?;
397            match (sig_owner, needed_entity) {
398                (
399                    DepositSigKeyOwner::OperatorSharedDeposit(sighash_type),
400                    EntityType::OperatorDeposit,
401                )
402                | (
403                    DepositSigKeyOwner::NofnSharedDeposit(sighash_type),
404                    EntityType::VerifierDeposit,
405                )
406                | (
407                    DepositSigKeyOwner::OperatorSharedSetup(sighash_type),
408                    EntityType::OperatorSetup,
409                ) => {
410                    sighashes.push((
411                        self.calculate_sighash_txin(idx, sighash_type)?,
412                        partial_signature_info.complete(sig_id, spend_data),
413                    ));
414                }
415                _ => {}
416            }
417        }
418        Ok(sighashes)
419    }
420
421    #[cfg(test)]
422    /// Returns the previous output (TxOut) for the specified input
423    pub fn get_input_txout(&self, input_idx: usize) -> &TxOut {
424        self.txins[input_idx].get_spendable().get_prevout()
425    }
426}
427
428impl TxHandler<Unsigned> {
429    /// Promotes an unsigned handler to a signed handler, checking that all witnesses are present.
430    ///
431    /// # Returns
432    /// A [`TxHandler<Signed>`] if all witnesses are present, or a [`BridgeError`] if not.
433    pub fn promote(self) -> Result<TxHandler<Signed>, BridgeError> {
434        if self.txins.iter().any(|s| s.get_witness().is_none()) {
435            return Err(eyre::eyre!("Missing witness data").into());
436        }
437
438        Ok(TxHandler {
439            transaction_type: self.transaction_type,
440            txins: self.txins,
441            txouts: self.txouts,
442            cached_tx: self.cached_tx,
443            cached_txid: self.cached_txid,
444            #[cfg(test)]
445            test_annex: self.test_annex,
446            phantom: PhantomData::<Signed>,
447        })
448    }
449
450    /// Sets the witness for a script path spend input.
451    ///
452    /// # Arguments
453    /// * `script_inputs` - The inputs to the tapscript.
454    /// * `txin_index` - The input index.
455    /// * `script_index` - The script index in the input's script list.
456    ///
457    /// # Returns
458    /// Ok(()) if successful, or a [`BridgeError`] if not.
459    pub fn set_p2tr_script_spend_witness<T: AsRef<[u8]>>(
460        &mut self,
461        script_inputs: &[T],
462        txin_index: usize,
463        script_index: usize,
464    ) -> Result<(), BridgeError> {
465        let txin = self
466            .txins
467            .get_mut(txin_index)
468            .ok_or(TxError::TxInputNotFound)?;
469
470        if txin.get_witness().is_some() {
471            return Err(TxError::WitnessAlreadySet.into());
472        }
473
474        let script = txin
475            .get_spendable()
476            .get_scripts()
477            .get(script_index)
478            .ok_or_else(|| {
479                eyre::eyre!("Could not find script {script_index} in input {txin_index}")
480            })?
481            .to_script_buf();
482
483        let spend_control_block = txin
484            .get_spendable()
485            .get_spend_info()
486            .as_ref()
487            .ok_or(TxError::MissingSpendInfo)?
488            .control_block(&(script.clone(), LeafVersion::TapScript))
489            .ok_or_eyre("Failed to find control block for script")?;
490
491        let mut witness = Witness::new();
492        script_inputs
493            .iter()
494            .for_each(|element| witness.push(element));
495        witness.push(script.clone());
496        witness.push(spend_control_block.serialize());
497
498        self.cached_tx.input[txin_index].witness = witness.clone();
499        txin.set_witness(witness);
500
501        Ok(())
502    }
503
504    /// Sets the witness for a key path spend input.
505    ///
506    /// # Arguments
507    /// * `signature` - The Taproot signature.
508    /// * `txin_index` - The input index.
509    ///
510    /// # Returns
511    /// Ok(()) if successful, or a [`BridgeError`] if not.
512    pub fn set_p2tr_key_spend_witness(
513        &mut self,
514        signature: &taproot::Signature,
515        txin_index: usize,
516    ) -> Result<(), BridgeError> {
517        let txin = self
518            .txins
519            .get_mut(txin_index)
520            .ok_or(TxError::TxInputNotFound)?;
521
522        if txin.get_witness().is_none() {
523            let witness = Witness::p2tr_key_spend(signature);
524            txin.set_witness(witness.clone());
525            self.cached_tx.input[txin_index].witness = witness;
526
527            Ok(())
528        } else {
529            Err(TxError::WitnessAlreadySet.into())
530        }
531    }
532}
533
534#[derive(Debug, Clone)]
535/// Builder for [`TxHandler`], allowing stepwise construction of inputs and outputs.
536pub struct TxHandlerBuilder {
537    transaction_type: TransactionType,
538    version: Version,
539    lock_time: absolute::LockTime,
540    txins: Vec<SpentTxIn>,
541    txouts: Vec<UnspentTxOut>,
542}
543
544impl TxHandlerBuilder {
545    /// Creates a new [`TxHandlerBuilder`] for the specified transaction type.
546    pub fn new(transaction_type: TransactionType) -> TxHandlerBuilder {
547        TxHandlerBuilder {
548            transaction_type,
549            version: Version::TWO,
550            lock_time: absolute::LockTime::ZERO,
551            txins: vec![],
552            txouts: vec![],
553        }
554    }
555
556    /// Sets the version for the transaction being built.
557    pub fn with_version(mut self, version: Version) -> Self {
558        self.version = version;
559        self
560    }
561
562    /// Adds an input to the transaction being built.
563    pub fn add_input(
564        mut self,
565        input_id: impl Into<SignatureId>,
566        spendable: SpendableTxIn,
567        spend_path: SpendPath,
568        sequence: Sequence,
569    ) -> Self {
570        self.txins.push(SpentTxIn::from_spendable(
571            input_id.into(),
572            spendable,
573            spend_path,
574            sequence,
575            None,
576        ));
577
578        self
579    }
580
581    /// Adds an output to the transaction being built.
582    pub fn add_output(mut self, output: UnspentTxOut) -> Self {
583        self.txouts.push(output);
584
585        self
586    }
587
588    /// Finalizes the transaction, returning an unsigned [`TxHandler`].
589    pub fn finalize(self) -> TxHandler<Unsigned> {
590        // construct cached Transaction
591        let tx = Transaction {
592            version: self.version,
593            lock_time: self.lock_time,
594            input: self.txins.iter().map(|s| s.to_txin()).collect(),
595            output: self.txouts.iter().map(|s| s.txout().clone()).collect(),
596        };
597        let txid = tx.compute_txid();
598
599        TxHandler::<Unsigned> {
600            transaction_type: self.transaction_type,
601            txins: self.txins,
602            txouts: self.txouts,
603            cached_tx: tx,
604            cached_txid: txid,
605            #[cfg(test)]
606            test_annex: None,
607            phantom: PhantomData,
608        }
609    }
610}
611
612/// Removes a [`TxHandler`] from a map by transaction type, returning an error if not found.
613///
614/// # Arguments
615/// * `txhandlers` - The map of transaction handlers.
616/// * `tx_type` - The transaction type to remove.
617///
618/// # Returns
619/// The removed [`TxHandler`], or a [`BridgeError`] if not found.
620pub fn remove_txhandler_from_map<T: State>(
621    txhandlers: &mut BTreeMap<TransactionType, TxHandler<T>>,
622    tx_type: TransactionType,
623) -> Result<TxHandler<T>, BridgeError> {
624    txhandlers
625        .remove(&tx_type)
626        .ok_or(TxError::TxHandlerNotFound(tx_type).into())
627}