From 33a5ba6cd2da0a37178c4226942b8f2c437a248b Mon Sep 17 00:00:00 2001
From: Alekos Filini <alekos.filini@gmail.com>
Date: Wed, 16 Sep 2020 17:31:43 +0200
Subject: [PATCH] [signer] Fix signing for `ShWpkh` inputs

---
 src/wallet/mod.rs    | 15 +++++++++++++++
 src/wallet/signer.rs | 27 +++++++++++++++++++--------
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs
index e45b97c7..d84a6206 100644
--- a/src/wallet/mod.rs
+++ b/src/wallet/mod.rs
@@ -2403,6 +2403,21 @@ mod test {
         assert_eq!(extracted.input[0].witness.len(), 2);
     }
 
+    #[test]
+    fn test_sign_single_xprv_sh_wpkh() {
+        let (wallet, _, _) = get_funded_wallet("sh(wpkh(tprv8ZgxMBicQKsPd3EupYiPRhaMooHKUHJxNsTfYuScep13go8QFfHdtkG9nRkFGb7busX4isf6X9dURGCoKgitaApQ6MupRhZMcELAxTBRJgS/*))");
+        let addr = wallet.get_new_address().unwrap();
+        let (psbt, _) = wallet
+            .create_tx(TxBuilder::with_recipients(vec![(addr.script_pubkey(), 0)]).send_all())
+            .unwrap();
+
+        let (signed_psbt, finalized) = wallet.sign(psbt, None).unwrap();
+        assert_eq!(finalized, true);
+
+        let extracted = signed_psbt.extract_tx();
+        assert_eq!(extracted.input[0].witness.len(), 2);
+    }
+
     #[test]
     fn test_sign_single_wif() {
         let (wallet, _, _) =
diff --git a/src/wallet/signer.rs b/src/wallet/signer.rs
index c290126c..12905151 100644
--- a/src/wallet/signer.rs
+++ b/src/wallet/signer.rs
@@ -100,7 +100,7 @@ use bitcoin::hashes::{hash160, Hash};
 use bitcoin::secp256k1::{Message, Secp256k1};
 use bitcoin::util::bip32::{ExtendedPrivKey, Fingerprint};
 use bitcoin::util::{bip143, psbt};
-use bitcoin::{PrivateKey, SigHash, SigHashType};
+use bitcoin::{PrivateKey, Script, SigHash, SigHashType};
 
 use miniscript::descriptor::{DescriptorPublicKey, DescriptorSecretKey, DescriptorXKey, KeyMap};
 use miniscript::{Legacy, MiniscriptKey, Segwitv0};
@@ -444,6 +444,16 @@ impl ComputeSighash for Legacy {
     }
 }
 
+fn p2wpkh_script_code(script: &Script) -> Script {
+    ScriptBuilder::new()
+        .push_opcode(opcodes::all::OP_DUP)
+        .push_opcode(opcodes::all::OP_HASH160)
+        .push_slice(&script[2..])
+        .push_opcode(opcodes::all::OP_EQUALVERIFY)
+        .push_opcode(opcodes::all::OP_CHECKSIG)
+        .into_script()
+}
+
 impl ComputeSighash for Segwitv0 {
     fn sighash(
         psbt: &psbt::PartiallySignedTransaction,
@@ -467,13 +477,14 @@ impl ComputeSighash for Segwitv0 {
             &Some(ref witness_script) => witness_script.clone(),
             &None => {
                 if witness_utxo.script_pubkey.is_v0_p2wpkh() {
-                    ScriptBuilder::new()
-                        .push_opcode(opcodes::all::OP_DUP)
-                        .push_opcode(opcodes::all::OP_HASH160)
-                        .push_slice(&witness_utxo.script_pubkey[2..])
-                        .push_opcode(opcodes::all::OP_EQUALVERIFY)
-                        .push_opcode(opcodes::all::OP_CHECKSIG)
-                        .into_script()
+                    p2wpkh_script_code(&witness_utxo.script_pubkey)
+                } else if psbt_input
+                    .redeem_script
+                    .as_ref()
+                    .map(Script::is_v0_p2wpkh)
+                    .unwrap_or(false)
+                {
+                    p2wpkh_script_code(&psbt_input.redeem_script.as_ref().unwrap())
                 } else {
                     return Err(SignerError::MissingWitnessScript);
                 }