From 3d80a0c962fbda0ad368b48ef81ab26b1e8f1dbd Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Tue, 25 Aug 2020 20:06:28 -0400
Subject: [PATCH] Expose the witnesses of the tx being hashed in SigHashCache

See docuemntation of the new method for more, but this allows
certain use patterns which were broken with the introduction of
SigHashCache.
---
 src/util/bip143.rs | 54 +++++++++++++++++++++++++++++++++++-----------
 1 file changed, 42 insertions(+), 12 deletions(-)

diff --git a/src/util/bip143.rs b/src/util/bip143.rs
index 8b805e61..d558fb3e 100644
--- a/src/util/bip143.rs
+++ b/src/util/bip143.rs
@@ -25,6 +25,8 @@ use blockdata::script::Script;
 use blockdata::transaction::{Transaction, TxIn, SigHashType};
 use consensus::encode::Encodable;
 
+use std::ops::{Deref, DerefMut};
+
 /// Parts of a sighash which are common across inputs or signatures, and which are
 /// sufficient (in conjunction with a private key) to sign the transaction
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -102,9 +104,9 @@ impl SighashComponents {
 }
 
 /// A replacement for SigHashComponents which supports all sighash modes
-pub struct SigHashCache<'a> {
+pub struct SigHashCache<R: Deref<Target=Transaction>> {
     /// Access to transaction required for various introspection
-    tx: &'a Transaction,
+    tx: R,
     /// Hash of all the previous outputs, computed as required
     hash_prevouts: Option<sha256d::Hash>,
     /// Hash of all the input sequence nos, computed as required
@@ -113,12 +115,12 @@ pub struct SigHashCache<'a> {
     hash_outputs: Option<sha256d::Hash>,
 }
 
-impl<'a> SigHashCache<'a> {
+impl<R: Deref<Target=Transaction>> SigHashCache<R> {
     /// Compute the sighash components from an unsigned transaction and auxiliary
     /// in a lazy manner when required.
     /// For the generated sighashes to be valid, no fields in the transaction may change except for
     /// script_sig and witnesses.
-    pub fn new(tx: &Transaction) -> SigHashCache {
+    pub fn new(tx: R) -> Self {
         SigHashCache {
             tx: tx,
             hash_prevouts: None,
@@ -189,15 +191,17 @@ impl<'a> SigHashCache<'a> {
             zero_hash.consensus_encode(&mut enc).unwrap();
         }
 
-        let txin = &self.tx.input[input_index];
+        {
+            let txin = &self.tx.input[input_index];
 
-        txin
-            .previous_output
-            .consensus_encode(&mut enc)
-            .unwrap();
-        script_code.consensus_encode(&mut enc).unwrap();
-        value.consensus_encode(&mut enc).unwrap();
-        txin.sequence.consensus_encode(&mut enc).unwrap();
+            txin
+                .previous_output
+                .consensus_encode(&mut enc)
+                .unwrap();
+            script_code.consensus_encode(&mut enc).unwrap();
+            value.consensus_encode(&mut enc).unwrap();
+            txin.sequence.consensus_encode(&mut enc).unwrap();
+        }
 
         if sighash != SigHashType::Single && sighash != SigHashType::None {
             self.hash_outputs().consensus_encode(&mut enc).unwrap();
@@ -215,6 +219,32 @@ impl<'a> SigHashCache<'a> {
     }
 }
 
+impl<R: DerefMut<Target=Transaction>> SigHashCache<R> {
+    /// When the SigHashCache is initialized with a mutable reference to a transaction instead of a
+    /// regular reference, this method is available to allow modification to the witnesses.
+    ///
+    /// This allows in-line signing such as
+    /// ```
+    /// use bitcoin::blockdata::transaction::{Transaction, SigHashType};
+    /// use bitcoin::util::bip143::SigHashCache;
+    /// use bitcoin::Script;
+    ///
+    /// let mut tx_to_sign = Transaction { version: 2, lock_time: 0, input: Vec::new(), output: Vec::new() };
+    /// let input_count = tx_to_sign.input.len();
+    ///
+    /// let mut sig_hasher = SigHashCache::new(&mut tx_to_sign);
+    /// for inp in 0..input_count {
+    ///     let prevout_script = Script::new();
+    ///     let _sighash = sig_hasher.signature_hash(inp, &prevout_script, 42, SigHashType::All);
+    ///     // ... sign the sighash
+    ///     sig_hasher.access_witness(inp).push(Vec::new());
+    /// }
+    /// ```
+    pub fn access_witness(&mut self, input_index: usize) -> &mut Vec<Vec<u8>> {
+        &mut self.tx.input[input_index].witness
+    }
+}
+
 #[cfg(test)]
 #[allow(deprecated)]
 mod tests {