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