From 9615d9427048b368bb43480ce045b3a64a32eb59 Mon Sep 17 00:00:00 2001
From: Satoshi Otomakan <satoshiotomakan8@gmail.com>
Date: Wed, 18 Sep 2024 17:08:25 +0200
Subject: [PATCH] feat(btc): Add PSBT signing of `non_witness_utxo`

---
 .../src/modules/psbt_request/utxo_psbt.rs     | 34 ++++++++++++++---
 .../tests/chains/bitcoin/bitcoin_sign/psbt.rs | 38 ++++++++++++++++---
 2 files changed, 60 insertions(+), 12 deletions(-)

diff --git a/rust/chains/tw_bitcoin/src/modules/psbt_request/utxo_psbt.rs b/rust/chains/tw_bitcoin/src/modules/psbt_request/utxo_psbt.rs
index b376d16560b..6b1b0c12273 100644
--- a/rust/chains/tw_bitcoin/src/modules/psbt_request/utxo_psbt.rs
+++ b/rust/chains/tw_bitcoin/src/modules/psbt_request/utxo_psbt.rs
@@ -46,9 +46,31 @@ impl<'a> UtxoPsbt<'a> {
 
     pub fn build_non_witness_utxo(
         &self,
-        _non_witness_utxo: &bitcoin::Transaction,
+        non_witness_utxo: &bitcoin::Transaction,
     ) -> SigningResult<(TransactionInput, UtxoToSign)> {
-        todo!()
+        let prev_out_idx = self.utxo.previous_output.vout as usize;
+        let prev_out = non_witness_utxo
+            .output
+            .get(prev_out_idx)
+            .or_tw_err(SigningErrorType::Error_invalid_utxo)
+            .with_context(|| {
+                format!("'Psbt::non_witness_utxo' does not contain '{prev_out_idx}' output")
+            })?;
+
+        let script = Script::from(prev_out.script_pubkey.to_bytes());
+        let builder = self.prepare_builder(prev_out.value)?;
+
+        match ConditionScriptParser.parse(&script)? {
+            ConditionScript::P2PK(pubkey) => builder.p2pk(&pubkey),
+            ConditionScript::P2PKH(pubkey_hash) => {
+                let pubkey = self.public_keys.get_ecdsa_public_key(&pubkey_hash)?;
+                builder.p2pkh(&pubkey)
+            },
+            ConditionScript::P2WPKH(_) | ConditionScript::P2TR(_) => {
+                SigningError::err(SigningErrorType::Error_invalid_params)
+                    .context("P2WPKH and P2TR scripts should be specified in 'witness_utxo'")
+            },
+        }
     }
 
     pub fn build_witness_utxo(
@@ -77,10 +99,6 @@ impl<'a> UtxoPsbt<'a> {
         }
     }
 
-    fn has_tap_scripts(&self) -> bool {
-        !self.utxo_psbt.tap_scripts.is_empty()
-    }
-
     pub fn prepare_builder(&self, amount: u64) -> SigningResult<UtxoBuilder> {
         let prevout_hash = H256::from(self.utxo.previous_output.txid.to_raw_hash().into_32());
         let prevout_index = self.utxo.previous_output.vout;
@@ -103,4 +121,8 @@ impl<'a> UtxoPsbt<'a> {
             .sighash_type(sighash_ty)
             .amount(amount))
     }
+
+    fn has_tap_scripts(&self) -> bool {
+        !self.utxo_psbt.tap_scripts.is_empty()
+    }
 }
diff --git a/rust/tw_any_coin/tests/chains/bitcoin/bitcoin_sign/psbt.rs b/rust/tw_any_coin/tests/chains/bitcoin/bitcoin_sign/psbt.rs
index 180e362d1f8..141089020cf 100644
--- a/rust/tw_any_coin/tests/chains/bitcoin/bitcoin_sign/psbt.rs
+++ b/rust/tw_any_coin/tests/chains/bitcoin/bitcoin_sign/psbt.rs
@@ -8,24 +8,50 @@ use tw_encoding::hex::DecodeHex;
 use tw_proto::BitcoinV2::Proto;
 
 #[test]
-fn test_bitcoin_sign_psbt_thorchain_swap() {
+fn test_bitcoin_sign_psbt_thorchain_swap_witness() {
     let private_key = "f00ffbe44c5c2838c13d2778854ac66b75e04eb6054f0241989e223223ad5e55"
         .decode_hex()
         .unwrap();
 
-    let psbt = "70736274ff0100bc0200000001147010db5fbcf619067c1090fec65c131443fbc80fb4aaeebe940e44206098c60000000000ffffffff03409c0000000000001600143a4ca60b4321b354156ba16c2916502c8fab38a50000000000000000426a403d3a474149412e41544f4d3a636f736d6f7331737377797a666d743675396a373437773537753438746778646575393573757a666c6d7175753a303a743a35305e60000000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d000000000001011f6603010000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d00000000".decode_hex().unwrap();
+    let psbt = "70736274ff0100bc0200000001147010db5fbcf619067c1090fec65c131443fbc80fb4aaeebe940e44206098c60000000000ffffffff0360ea000000000000160014f22a703617035ef7f490743d50f26ae08c30d0a70000000000000000426a403d3a474149412e41544f4d3a636f736d6f7331737377797a666d743675396a373437773537753438746778646575393573757a666c6d7175753a303a743a35303e12000000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d000000000001011f6603010000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d00000000".decode_hex().unwrap();
     let input = Proto::PsbtSigningInput {
         psbt: psbt.into(),
         private_keys: vec![private_key.into()],
         ..Proto::PsbtSigningInput::default()
     };
 
+    // Successfully broadcasted: https://mempool.space/tx/634a416e82ac710166725f6a4090ac7b5db69687e86b2d2e38dcb3d91c956c32
     BitcoinPsbtSignHelper::new(&input).coin(CoinType::Bitcoin).sign_psbt(Expected {
-        psbt: "70736274ff0100bc0200000001147010db5fbcf619067c1090fec65c131443fbc80fb4aaeebe940e44206098c60000000000ffffffff03409c0000000000001600143a4ca60b4321b354156ba16c2916502c8fab38a50000000000000000426a403d3a474149412e41544f4d3a636f736d6f7331737377797a666d743675396a373437773537753438746778646575393573757a666c6d7175753a303a743a35305e60000000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d000000000001011f6603010000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d01086b024730440220415fb09a3fb0ae061d529de4d45dad70a2d0716b292ab4a6010a628c333155bd0220231e8c2ff668930e1e012009ac49eb05bdfd6b1bdfc9abea441306a446d8f3bd01210306d8c664ea8fd2683eebea1d3114d90e0a5429e5783ba49b80ddabce04ff28f300000000",
-        encoded: "02000000000101147010db5fbcf619067c1090fec65c131443fbc80fb4aaeebe940e44206098c60000000000ffffffff03409c0000000000001600143a4ca60b4321b354156ba16c2916502c8fab38a50000000000000000426a403d3a474149412e41544f4d3a636f736d6f7331737377797a666d743675396a373437773537753438746778646575393573757a666c6d7175753a303a743a35305e60000000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d024730440220415fb09a3fb0ae061d529de4d45dad70a2d0716b292ab4a6010a628c333155bd0220231e8c2ff668930e1e012009ac49eb05bdfd6b1bdfc9abea441306a446d8f3bd01210306d8c664ea8fd2683eebea1d3114d90e0a5429e5783ba49b80ddabce04ff28f300000000",
-        txid: "b508068c93bfc679863cbb3383c0e7fb5586b78f3ff7685c2a941110878634d0",
+        psbt: "70736274ff0100bc0200000001147010db5fbcf619067c1090fec65c131443fbc80fb4aaeebe940e44206098c60000000000ffffffff0360ea000000000000160014f22a703617035ef7f490743d50f26ae08c30d0a70000000000000000426a403d3a474149412e41544f4d3a636f736d6f7331737377797a666d743675396a373437773537753438746778646575393573757a666c6d7175753a303a743a35303e12000000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d000000000001011f6603010000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d01086c02483045022100b1229a008f20691639767bf925d6b8956ea957ccc633ad6b5de3618733a55e6b02205774d3320489b8a57a6f8de07f561de3e660ff8e587f6ac5422c49020cd4dc9101210306d8c664ea8fd2683eebea1d3114d90e0a5429e5783ba49b80ddabce04ff28f300000000",
+        encoded: "02000000000101147010db5fbcf619067c1090fec65c131443fbc80fb4aaeebe940e44206098c60000000000ffffffff0360ea000000000000160014f22a703617035ef7f490743d50f26ae08c30d0a70000000000000000426a403d3a474149412e41544f4d3a636f736d6f7331737377797a666d743675396a373437773537753438746778646575393573757a666c6d7175753a303a743a35303e12000000000000160014b139199ec796f36fc42e637f42da8e3e6720aa9d02483045022100b1229a008f20691639767bf925d6b8956ea957ccc633ad6b5de3618733a55e6b02205774d3320489b8a57a6f8de07f561de3e660ff8e587f6ac5422c49020cd4dc9101210306d8c664ea8fd2683eebea1d3114d90e0a5429e5783ba49b80ddabce04ff28f300000000",
+        txid: "634a416e82ac710166725f6a4090ac7b5db69687e86b2d2e38dcb3d91c956c32",
         vsize: 216,
-        weight: 861,
+        weight: 862,
         fee: 1736,
     });
 }
+
+#[test]
+fn test_bitcoin_sign_psbt_thorchain_swap_non_witness() {
+    // 1CKZYtNxAQnTbygz6vyhBYnwx4NvcxURMB
+    let private_key = "7a87cb2c9fa56f7a63dfc50659dca260473cb6bb0fd4d8a2beeaf5357d41de95"
+        .decode_hex()
+        .unwrap();
+
+    let psbt = "70736274ff01008202000000015c37bcf049b7e62dd5bfd707e0998ce86163b786e3cd45db2336cb794a8d8aa10000000000ffffffff03f82a000000000000160014bf5a13a26791a5db6406304a46952e264c2b28910000000000000000056a032b3a6291950000000000001976a9147c2c0ac72afbde13ecf52fca54368e7883b538b188ac000000000001007e0200000002714916920be4dbc87cbb8697ca9b1420d6b1e47e7d732e2d2e0e7a935087788d0000000000ffffffff326c951cd9b3dc382e2d6be88796b65d7bac90406a5f72660171ac826e414a630200000000ffffffff01efca0000000000001976a9147c2c0ac72afbde13ecf52fca54368e7883b538b188ac0000000000000000".decode_hex().unwrap();
+    let input = Proto::PsbtSigningInput {
+        psbt: psbt.into(),
+        private_keys: vec![private_key.into()],
+        ..Proto::PsbtSigningInput::default()
+    };
+
+    // Successfully broadcasted: https://mempool.space/tx/710e9270b57720f567ada156c6ac72177aa00a36789e2c6526fd80040fae3ce4
+    BitcoinPsbtSignHelper::new(&input).coin(CoinType::Bitcoin).sign_psbt(Expected {
+        psbt: "70736274ff01008202000000015c37bcf049b7e62dd5bfd707e0998ce86163b786e3cd45db2336cb794a8d8aa10000000000ffffffff03f82a000000000000160014bf5a13a26791a5db6406304a46952e264c2b28910000000000000000056a032b3a6291950000000000001976a9147c2c0ac72afbde13ecf52fca54368e7883b538b188ac000000000001007e0200000002714916920be4dbc87cbb8697ca9b1420d6b1e47e7d732e2d2e0e7a935087788d0000000000ffffffff326c951cd9b3dc382e2d6be88796b65d7bac90406a5f72660171ac826e414a630200000000ffffffff01efca0000000000001976a9147c2c0ac72afbde13ecf52fca54368e7883b538b188ac0000000001076a473044022057ce7a6147fd9e139df797adcec440bad60770f40cbd609363e3075b64d3eccd02200ae7dce5f7d1fa18c5e907a16c1b078fa90f537d36101447e53fbd058d2d950a0121036c3b7dfd678da989d91593e49918a6c9d8a1d37c7e9c0abeae2118c312e69b3100000000",
+        encoded: "02000000015c37bcf049b7e62dd5bfd707e0998ce86163b786e3cd45db2336cb794a8d8aa1000000006a473044022057ce7a6147fd9e139df797adcec440bad60770f40cbd609363e3075b64d3eccd02200ae7dce5f7d1fa18c5e907a16c1b078fa90f537d36101447e53fbd058d2d950a0121036c3b7dfd678da989d91593e49918a6c9d8a1d37c7e9c0abeae2118c312e69b31ffffffff03f82a000000000000160014bf5a13a26791a5db6406304a46952e264c2b28910000000000000000056a032b3a6291950000000000001976a9147c2c0ac72afbde13ecf52fca54368e7883b538b188ac00000000",
+        txid: "710e9270b57720f567ada156c6ac72177aa00a36789e2c6526fd80040fae3ce4",
+        vsize: 236,
+        weight: 944,
+        fee: 2662,
+    });
+}