From f0aca23590c50019ea1c06caf69ace29b8293c81 Mon Sep 17 00:00:00 2001
From: Sonkeng Maldini <sdmg15@pm.me>
Date: Tue, 12 May 2026 06:24:56 +0100
Subject: [PATCH] fix(psbt): sanity check psbt before signing

---
 src/psbt/mod.rs | 25 ++++++++++++++++++++
 tests/wallet.rs | 62 +++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/src/psbt/mod.rs b/src/psbt/mod.rs
index 3655d9e4..fc35fcac 100644
--- a/src/psbt/mod.rs
+++ b/src/psbt/mod.rs
@@ -41,6 +41,12 @@ impl PsbtUtils for Psbt {
         let tx = &self.unsigned_tx;
         let input = self.inputs.get(input_index)?;
 
+        // Validate that the input index is within bounds of the transaction's inputs.
+        // This prevents out-of-bounds panic when PSBT is malformed (inputs.len() > tx.input.len()).
+        if input_index >= tx.input.len() {
+            return None;
+        }
+
         match (&input.witness_utxo, &input.non_witness_utxo) {
             (Some(_), _) => input.witness_utxo.clone(),
             (_, Some(_)) => input.non_witness_utxo.as_ref().and_then(|prev_tx| {
@@ -154,4 +160,23 @@ mod tests {
         // Must return None — vout out of bounds, no panic
         assert_eq!(psbt.get_utxo_for(0), None);
     }
+
+    #[test]
+    fn get_utxo_doesnt_panic_on_input_count_mismatch() {
+        let prev_tx = build_tx(Amount::from_sat(100_000));
+
+        // Create a valid PSBT with 1 input
+        let mut psbt = build_psbt(&prev_tx, 0);
+
+        // Add more inputs to psbt.inputs without adding to unsigned_tx.input
+        let extra_input = Input {
+            non_witness_utxo: Some(prev_tx.clone()),
+            ..Default::default()
+        };
+        psbt.inputs.push(extra_input);
+
+        assert_eq!(psbt.inputs.len(), 2);
+        assert_eq!(psbt.unsigned_tx.input.len(), 1);
+        assert_eq!(None, psbt.get_utxo_for(1));
+    }
 }
diff --git a/tests/wallet.rs b/tests/wallet.rs
index 268c66f8..18735cfb 100644
--- a/tests/wallet.rs
+++ b/tests/wallet.rs
@@ -17,8 +17,8 @@ use bitcoin::script::PushBytesBuf;
 use bitcoin::sighash::{EcdsaSighashType, TapSighashType};
 use bitcoin::taproot::TapNodeHash;
 use bitcoin::{
-    absolute, transaction, Address, Amount, BlockHash, FeeRate, Network, OutPoint, ScriptBuf,
-    Sequence, SignedAmount, Transaction, TxIn, TxOut, Txid,
+    absolute, transaction, Address, Amount, BlockHash, FeeRate, Network, OutPoint, Psbt, ScriptBuf,
+    Sequence, SignedAmount, Transaction, TxIn, TxOut, Txid, Witness,
 };
 use rand::rngs::StdRng;
 use rand::SeedableRng;
@@ -3007,3 +3007,61 @@ fn test_tx_ordering_untouched_preserves_insertion_ordering_bnb_success() {
         "UTXOs should be ordered with required first, then selected"
     );
 }
+
+#[test]
+fn test_wallet_sign_rejects_malformed_psbt_input_count_mismatch() {
+    let (wallet, _) = get_funded_wallet_wpkh();
+
+    // Create a transaction with 1 input
+    let prev_tx = Transaction {
+        version: transaction::Version::TWO,
+        lock_time: absolute::LockTime::ZERO,
+        input: vec![TxIn {
+            previous_output: OutPoint::null(),
+            script_sig: ScriptBuf::default(),
+            sequence: Sequence::MAX,
+            witness: Witness::default(),
+        }],
+        output: vec![TxOut {
+            value: Amount::from_sat(50_000),
+            script_pubkey: ScriptBuf::default(),
+        }],
+    };
+
+    let unsigned_tx = Transaction {
+        version: transaction::Version::TWO,
+        lock_time: absolute::LockTime::ZERO,
+        input: vec![TxIn {
+            previous_output: OutPoint {
+                txid: prev_tx.compute_txid(),
+                vout: 0,
+            },
+            script_sig: ScriptBuf::default(),
+            sequence: Sequence::MAX,
+            witness: Witness::default(),
+        }],
+        output: vec![TxOut {
+            value: Amount::from_sat(40_000),
+            script_pubkey: ScriptBuf::default(),
+        }],
+    };
+
+    let mut psbt = Psbt::from_unsigned_tx(unsigned_tx).unwrap();
+
+    // Add extra PSBT inputs without adding to transaction
+    use bitcoin::psbt::Input;
+    psbt.inputs.push(Input {
+        non_witness_utxo: Some(prev_tx.clone()),
+        ..Default::default()
+    });
+    psbt.inputs.push(Input {
+        non_witness_utxo: Some(prev_tx),
+        ..Default::default()
+    });
+
+    assert_eq!(psbt.inputs.len(), 3);
+    assert_eq!(psbt.unsigned_tx.input.len(), 1);
+
+    let result = wallet.sign(&mut psbt.clone(), SignOptions::default());
+    assert!(matches!(result, Err(SignerError::MissingNonWitnessUtxo)));
+}