From 5afb0eaf40642bfda87f2681a650a919d43163a4 Mon Sep 17 00:00:00 2001
From: violet360 <mazedk2k@gmail.com>
Date: Mon, 28 Feb 2022 19:31:02 +0530
Subject: [PATCH] API to get an iterator for funding utxos in psbt

---
 src/util/psbt/error.rs |  6 ++++++
 src/util/psbt/mod.rs   | 30 ++++++++++++++++++++++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/util/psbt/error.rs b/src/util/psbt/error.rs
index 3d46a118..fbcb2977 100644
--- a/src/util/psbt/error.rs
+++ b/src/util/psbt/error.rs
@@ -37,8 +37,12 @@ pub enum Error {
     /// Magic bytes for a PSBT must be the ASCII for "psbt" serialized in most
     /// significant byte order.
     InvalidMagic,
+    /// Missing both the witness and non-witness utxo.
+    MissingUtxo,
     /// The separator for a PSBT must be `0xff`.
     InvalidSeparator,
+    /// Returned when output index is out of bounds in relation to the output in non-witness UTXO.
+    PsbtUtxoOutOfbounds,
     /// Known keys must be according to spec.
     InvalidKey(raw::Key),
     /// Non-proprietary key type found when proprietary key was expected
@@ -98,6 +102,8 @@ impl fmt::Display for Error {
             }
             Error::NoMorePairs => f.write_str("no more key-value pairs for this psbt map"),
             Error::HashParseError(e) => write!(f, "Hash Parse Error: {}", e),
+            Error::MissingUtxo => f.write_str("UTXO information is not present in PSBT"),
+            Error::PsbtUtxoOutOfbounds => f.write_str("output index is out of bounds of non witness script output array"),
             Error::InvalidPreimageHashPair{ref preimage, ref hash, ref hash_type} => {
                 // directly using debug forms of psbthash enums
                 write!(f, "Preimage {:?} does not match {:?} hash {:?}", preimage, hash_type, hash )
diff --git a/src/util/psbt/mod.rs b/src/util/psbt/mod.rs
index e1cf4cca..89495cc4 100644
--- a/src/util/psbt/mod.rs
+++ b/src/util/psbt/mod.rs
@@ -20,14 +20,14 @@
 //!
 
 use blockdata::script::Script;
-use blockdata::transaction::Transaction;
+use blockdata::transaction::{ TxOut, Transaction};
 use consensus::{encode, Encodable, Decodable};
 use consensus::encode::MAX_VEC_SIZE;
+pub use util::sighash::Prevouts;
 
 use prelude::*;
 
 use io;
-
 mod error;
 pub use self::error::Error;
 
@@ -72,6 +72,32 @@ pub struct PartiallySignedTransaction {
 }
 
 impl PartiallySignedTransaction {
+    /// Returns an iterator for the funding UTXOs of the psbt
+    ///
+    /// For each PSBT input that contains UTXO information `Ok` is returned containing that information.
+    /// The order of returned items is same as the order of inputs.
+    ///
+    /// ## Errors
+    ///
+    /// The function returns error when UTXO information is not present or is invalid.
+    ///
+    /// ## Panics
+    ///
+    /// The function panics if the length of transaction inputs is not equal to the length of PSBT inputs.
+    pub fn iter_funding_utxos(&self) -> impl Iterator<Item = Result<&TxOut, Error>> {
+        assert_eq!(self.inputs.len(), self.unsigned_tx.input.len());
+        self.unsigned_tx.input.iter().zip(&self.inputs).map(|(tx_input, psbt_input)| {
+            match (&psbt_input.witness_utxo, &psbt_input.non_witness_utxo) {
+                (Some(witness_utxo), _) => Ok(witness_utxo),
+                (None, Some(non_witness_utxo)) => {
+                    let vout = tx_input.previous_output.vout as usize;
+                    non_witness_utxo.output.get(vout).ok_or(Error::PsbtUtxoOutOfbounds)
+                },
+                (None, None) => Err(Error::MissingUtxo),
+            }
+        })
+    }
+
     /// Checks that unsigned transaction does not have scriptSig's or witness
     /// data
     fn unsigned_tx_checks(&self) -> Result<(), Error> {