From 1ea6ae9276b04c43d7e3e9a95bb60cee1140fcd3 Mon Sep 17 00:00:00 2001
From: vmammal <valuedmammal@protonmail.com>
Date: Sun, 11 Feb 2024 15:32:36 -0500
Subject: [PATCH] test(psbt): fixup test_psbt_multiple_internalkey_signers

to verify the signature of the input and ensure the right internal
key is used to sign. Get the signature directly from the witness
field, that way we can use the default SignOptions, allowing extra
taproot fields to be reset. This fixes a shortcoming of a previous
version of the test that relied on the result of `finalize_psbt`
but would have erroneously allowed signing with the wrong key.
---
 crates/bdk/tests/psbt.rs | 103 ++++++++++++++++++++++++---------------
 1 file changed, 65 insertions(+), 38 deletions(-)

diff --git a/crates/bdk/tests/psbt.rs b/crates/bdk/tests/psbt.rs
index 1f24864d63..b4d054c115 100644
--- a/crates/bdk/tests/psbt.rs
+++ b/crates/bdk/tests/psbt.rs
@@ -157,41 +157,68 @@ fn test_psbt_fee_rate_with_missing_txout() {
     assert!(pkh_psbt.fee_rate().is_none());
 }
 
-// This is supposed to check that we can always finalize a signed Psbt and broadcast without
-// issue, even in the presence of multiple internal key signers, but we're not yet confident
-// that this test is effective. Ideally, we don't rely on `finalize_mut` to finalize.
-// See https://github.com/bitcoindevkit/bdk/pull/1310#discussion_r1473516585
-// for the discussion that motivated this change.
-// #[test]
-// fn test_psbt_multiple_internalkey_signers() {
-//     use bdk::signer::{SignerContext, SignerOrdering, SignerWrapper};
-//     use bdk::KeychainKind;
-//     use bitcoin::{secp256k1::Secp256k1, PrivateKey};
-//     use miniscript::psbt::PsbtExt;
-//     use std::sync::Arc;
-
-//     let secp = Secp256k1::new();
-//     let (mut wallet, _) = get_funded_wallet(get_test_tr_single_sig());
-//     let send_to = wallet.get_address(AddressIndex::New);
-//     let mut builder = wallet.build_tx();
-//     builder.add_recipient(send_to.script_pubkey(), 10_000);
-//     let mut psbt = builder.finish().unwrap();
-//     // Adds a signer for the wrong internal key, bdk should not use this key to sign
-//     wallet.add_signer(
-//         KeychainKind::External,
-//         // A signerordering lower than 100, bdk will use this signer first
-//         SignerOrdering(0),
-//         Arc::new(SignerWrapper::new(
-//             PrivateKey::from_wif("5J5PZqvCe1uThJ3FZeUUFLCh2FuK9pZhtEK4MzhNmugqTmxCdwE").unwrap(),
-//             SignerContext::Tap {
-//                 is_internal_key: true,
-//             },
-//         )),
-//     );
-//     let _ = wallet.sign(&mut psbt, SignOptions::default()).unwrap();
-//     // Checks that we signed using the right key
-//     assert!(
-//         psbt.finalize_mut(&secp).is_ok(),
-//         "The wrong internal key was used"
-//     );
-// }
+#[test]
+fn test_psbt_multiple_internalkey_signers() {
+    use bdk::signer::{SignerContext, SignerOrdering, SignerWrapper};
+    use bdk::KeychainKind;
+    use bitcoin::key::TapTweak;
+    use bitcoin::secp256k1::{schnorr, KeyPair, Message, Secp256k1, XOnlyPublicKey};
+    use bitcoin::sighash::{Prevouts, SighashCache, TapSighashType};
+    use bitcoin::{PrivateKey, TxOut};
+    use std::sync::Arc;
+
+    let secp = Secp256k1::new();
+    let wif = "cNJmN3fH9DDbDt131fQNkVakkpzawJBSeybCUNmP1BovpmGQ45xG";
+    let desc = format!("tr({})", wif);
+    let prv = PrivateKey::from_wif(wif).unwrap();
+    let keypair = KeyPair::from_secret_key(&secp, &prv.inner);
+
+    let (mut wallet, _) = get_funded_wallet(&desc);
+    let to_spend = wallet.get_balance().total();
+    let send_to = wallet.get_address(AddressIndex::New);
+    let mut builder = wallet.build_tx();
+    builder.drain_to(send_to.script_pubkey()).drain_wallet();
+    let mut psbt = builder.finish().unwrap();
+    let unsigned_tx = psbt.unsigned_tx.clone();
+
+    // Adds a signer for the wrong internal key, bdk should not use this key to sign
+    wallet.add_signer(
+        KeychainKind::External,
+        // A signerordering lower than 100, bdk will use this signer first
+        SignerOrdering(0),
+        Arc::new(SignerWrapper::new(
+            PrivateKey::from_wif("5J5PZqvCe1uThJ3FZeUUFLCh2FuK9pZhtEK4MzhNmugqTmxCdwE").unwrap(),
+            SignerContext::Tap {
+                is_internal_key: true,
+            },
+        )),
+    );
+    let finalized = wallet.sign(&mut psbt, SignOptions::default()).unwrap();
+    assert!(finalized);
+
+    // To verify, we need the signature, message, and pubkey
+    let witness = psbt.inputs[0].final_script_witness.as_ref().unwrap();
+    assert!(!witness.is_empty());
+    let signature = schnorr::Signature::from_slice(witness.iter().next().unwrap()).unwrap();
+
+    // the prevout we're spending
+    let prevouts = &[TxOut {
+        script_pubkey: send_to.script_pubkey(),
+        value: to_spend,
+    }];
+    let prevouts = Prevouts::All(prevouts);
+    let input_index = 0;
+    let mut sighash_cache = SighashCache::new(unsigned_tx);
+    let sighash = sighash_cache
+        .taproot_key_spend_signature_hash(input_index, &prevouts, TapSighashType::Default)
+        .unwrap();
+    let message = Message::from(sighash);
+
+    // add tweak. this was taken from `signer::sign_psbt_schnorr`
+    let keypair = keypair.tap_tweak(&secp, None).to_inner();
+    let (xonlykey, _parity) = XOnlyPublicKey::from_keypair(&keypair);
+
+    // Must verify if we used the correct key to sign
+    let verify_res = secp.verify_schnorr(&signature, &message, &xonlykey);
+    assert!(verify_res.is_ok(), "The wrong internal key was used");
+}