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