Add support for merging unique psbts

This operation is needed for NS1R payjoins and in the future for the
more generalized NSMR case. Rust bitcoin does support combining PSBTs
that include the exact same unsigned tx; however, it does not have
support for merging PSBTs that are unique.

We add the additional merge operation as a new trait extension under the
merge submodule of a newly created psbt module.

payjoin/Cargo.toml

@@ -20,6 +20,7 @@ default = ["v2"]
 #[doc = "Core features for payjoin state machines"]
 _core = ["bitcoin/rand", "serde_json", "url", "bitcoin_uri"]
 directory = []
+psbt-merge = []
 v1 = ["_core"]
 v2 = ["_core", "bitcoin/serde", "hpke", "dep:http", "bhttp", "ohttp", "serde", "url/serde", "directory"]
 #[doc = "Functions to fetch OHTTP keys via CONNECT proxy using reqwest. Enables `v2` since only `v2` uses OHTTP."]

payjoin/src/psbt/merge.rs

@@ -0,0 +1,169 @@
+//! Utilities for merging unique v0 PSBTs
+use bitcoin::Psbt;
+
+#[allow(dead_code)]
+/// Try to merge two PSBTs
+/// PSBTs here should not have the same unsigned tx
+/// if you do have the same unsigned tx, use `combine` instead
+/// Note: this method does not merge non inputs or outputs
+/// Note: if there are duplicate inputs, the first input will be kept
+/// Note: if there are duplicate outputs, both outputs will be kept
+/// ```no_run
+/// let psbts = vec![psbt_1.clone(), psbt_2.clone(), ..., psbt_n.clone()];
+/// let merged_psbt = psbts.into_iter().reduce(merge_unsigned_tx).unwrap();
+/// ```
+pub(crate) fn merge_unsigned_tx(acc: Psbt, psbt: Psbt) -> Psbt {
+    let mut unsigned_tx = acc.unsigned_tx;
+    unsigned_tx.input.extend(psbt.unsigned_tx.input);
+    unsigned_tx.input.dedup_by_key(|input| input.previous_output);
+    unsigned_tx.output.extend(psbt.unsigned_tx.output);
+
+    Psbt::from_unsigned_tx(unsigned_tx).expect("")
+}
+
+#[cfg(test)]
+mod tests {
+    use bitcoin::absolute::LockTime;
+    use bitcoin::hashes::Hash;
+    use bitcoin::key::rand::Rng;
+    use bitcoin::secp256k1::rand::thread_rng;
+    use bitcoin::secp256k1::SECP256K1;
+    use bitcoin::{
+        Amount, Network, OutPoint, Psbt, ScriptBuf, Sequence, Transaction, TxIn, TxOut, Txid,
+        Witness,
+    };
+
+    use super::merge_unsigned_tx;
+
+    /// Create a random p2wpkh script
+    fn random_p2wpkh_script() -> ScriptBuf {
+        let sk = bitcoin::PrivateKey::generate(Network::Bitcoin);
+        let pk = sk.public_key(SECP256K1);
+
+        pk.p2wpkh_script_code().unwrap()
+    }
+
+    /// Create a random 32 byte txid
+    fn random_txid() -> Txid {
+        let mut rng = thread_rng();
+        let mut txid = [0u8; 32];
+        rng.try_fill(&mut txid).expect("should fill");
+        Txid::from_slice(&txid).unwrap()
+    }
+
+    /// Create a tx with random inputs and outputs
+    /// Note: all outputs have the same 1000 sat value
+    /// Transactions are created with version 2
+    fn create_tx(num_inputs: usize, num_outputs: usize) -> Transaction {
+        let txid = random_txid();
+
+        let mut inputs = vec![];
+        for i in 0..num_inputs {
+            let op = OutPoint::new(txid, i as u32);
+            inputs.push(TxIn {
+                previous_output: op,
+                script_sig: ScriptBuf::new(),
+                sequence: Sequence::MAX,
+                witness: Default::default(),
+            });
+        }
+
+        let mut outputs = vec![];
+        for _ in 0..num_outputs {
+            outputs.push(TxOut {
+                value: Amount::from_sat(1000),
+                script_pubkey: random_p2wpkh_script(),
+            });
+        }
+
+        Transaction {
+            version: bitcoin::transaction::Version(2),
+            lock_time: LockTime::ZERO,
+            input: inputs,
+            output: outputs,
+        }
+    }
+
+    /// Test that we can merge two psbts with unique unsigned txs
+    #[test]
+    fn test_merge_unsigned_txs() {
+        let txs = (0..10).map(|_| create_tx(2, 3)).collect::<Vec<_>>();
+        let psbts = txs.iter().map(|tx| Psbt::from_unsigned_tx(tx.clone()).unwrap());
+        let merged_psbt = psbts.reduce(merge_unsigned_tx).unwrap();
+        // let merged_psbt = merge_unsigned_tx(psbts).unwrap();
+
+        for tx in txs.iter() {
+            assert!(merged_psbt.unsigned_tx.input.contains(&tx.input[0]));
+            assert!(merged_psbt.unsigned_tx.input.contains(&tx.input[1]));
+            assert!(merged_psbt.unsigned_tx.output.contains(&tx.output[0]));
+            assert!(merged_psbt.unsigned_tx.output.contains(&tx.output[1]));
+            assert!(merged_psbt.unsigned_tx.output.contains(&tx.output[2]));
+        }
+    }
+
+    /// Test merging empty PSBTs
+    #[test]
+    fn test_merge_empty_psbts() {
+        let tx_1 = create_tx(0, 0);
+        let tx_2 = create_tx(0, 0);
+        let psbts =
+            vec![Psbt::from_unsigned_tx(tx_1).unwrap(), Psbt::from_unsigned_tx(tx_2).unwrap()];
+
+        let merged_psbt = psbts.into_iter().reduce(merge_unsigned_tx).unwrap();
+
+        assert_eq!(merged_psbt.inputs.len(), 0);
+        assert_eq!(merged_psbt.outputs.len(), 0);
+    }
+
+    /// Test that we cannot merge two psbts if psbts share inputs
+    #[test]
+    fn should_not_merge_if_psbt_share_inputs() {
+        let tx = create_tx(1, 1);
+        let psbt = Psbt::from_unsigned_tx(tx.clone()).unwrap();
+        let psbts = vec![psbt.clone(), psbt.clone()];
+
+        let res = psbts.into_iter().reduce(merge_unsigned_tx).unwrap();
+        let unsigned_tx = res.unsigned_tx;
+
+        assert_eq!(unsigned_tx.input.len(), 1);
+        assert_eq!(unsigned_tx.input[0].previous_output, tx.input[0].previous_output);
+        assert_eq!(unsigned_tx.output.len(), 2);
+        assert_eq!(unsigned_tx.output[0], tx.output[0]);
+        assert_eq!(unsigned_tx.output[1], tx.output[0]);
+    }
+
+    /// Test that we cannot merge two psbts if psbts have inputs with witness data
+    #[test]
+    fn should_not_merge_signed_psbt() {
+        let tx_1 = create_tx(1, 1);
+        let tx_2 = create_tx(1, 1);
+        let mut original_psbt = Psbt::from_unsigned_tx(tx_1.clone()).unwrap();
+        let mut other = Psbt::from_unsigned_tx(tx_2.clone()).unwrap();
+
+        original_psbt.inputs[0].final_script_witness = Some(Witness::new());
+        original_psbt.unsigned_tx.input[0].witness = Witness::new();
+        other.inputs[0].final_script_witness = Some(Witness::new());
+        let psbts = vec![original_psbt.clone(), other.clone()];
+        let merged_psbt = psbts.into_iter().reduce(merge_unsigned_tx).unwrap();
+
+        assert_eq!(merged_psbt.unsigned_tx.input[0], original_psbt.unsigned_tx.input[0]);
+        assert_eq!(merged_psbt.unsigned_tx.input[1], other.unsigned_tx.input[0]);
+        assert_eq!(merged_psbt.unsigned_tx.output[0], original_psbt.unsigned_tx.output[0]);
+        assert_eq!(merged_psbt.unsigned_tx.output[1], other.unsigned_tx.output[0]);
+    }
+
+    /// Test merging PSBTs with only inputs or only outputs
+    #[test]
+    fn test_merge_inputs_or_outputs_only() {
+        let tx_1 = create_tx(2, 0);
+        let tx_2 = create_tx(0, 3);
+
+        let psbts =
+            vec![Psbt::from_unsigned_tx(tx_1).unwrap(), Psbt::from_unsigned_tx(tx_2).unwrap()];
+
+        let merged_psbt = psbts.into_iter().reduce(merge_unsigned_tx).unwrap();
+
+        assert_eq!(merged_psbt.inputs.len(), 2);
+        assert_eq!(merged_psbt.outputs.len(), 3);
+    }
+}

payjoin/src/psbt.rs → payjoin/src/psbt/mod.rs

@@ -1,5 +1,8 @@
 //! Utilities to make work with PSBTs easier
 
+#[cfg(feature = "psbt-merge")]
+pub(crate) mod merge;
+
 use std::collections::BTreeMap;
 use std::fmt;