Allow contributing multiple inputs

`contribute_witness_input` is now `contribute_witness_inputs`. It takes an arbitrary number of inputs as a HashMap and returns a `ProvisionalProposal`.
payjoin · Aug 9, 2024 · 24c0bee · 24c0bee
1 parent 5cfa490
commit 24c0bee
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 52 deletions.
diff --git a/payjoin-cli/src/app/v1.rs b/payjoin-cli/src/app/v1.rs
@@ -332,7 +332,7 @@ impl App {
 }
 
 fn try_contributing_inputs(
-    payjoin: payjoin::receive::WantsInputs,
+    mut payjoin: payjoin::receive::WantsInputs,
     bitcoind: &bitcoincore_rpc::Client,
 ) -> Result<payjoin::receive::ProvisionalProposal> {
     use bitcoin::OutPoint;
@@ -351,11 +351,10 @@ fn try_contributing_inputs(
         .find(|i| i.txid == selected_outpoint.txid && i.vout == selected_outpoint.vout)
         .context("This shouldn't happen. Failed to retrieve the privacy preserving utxo from those we provided to the seclector.")?;
     log::debug!("selected utxo: {:#?}", selected_utxo);
-
-    //  calculate receiver payjoin outputs given receiver payjoin inputs and original_psbt,
     let txo_to_contribute = bitcoin::TxOut {
         value: selected_utxo.amount,
         script_pubkey: selected_utxo.script_pub_key.clone(),
     };
-    Ok(payjoin.contribute_witness_input(txo_to_contribute, selected_outpoint))
+
+    Ok(payjoin.contribute_witness_inputs(HashMap::from([(selected_outpoint, txo_to_contribute)])))
 }
diff --git a/payjoin-cli/src/app/v2.rs b/payjoin-cli/src/app/v2.rs
@@ -348,7 +348,7 @@ impl App {
 }
 
 fn try_contributing_inputs(
-    payjoin: payjoin::receive::v2::WantsInputs,
+    mut payjoin: payjoin::receive::v2::WantsInputs,
     bitcoind: &bitcoincore_rpc::Client,
 ) -> Result<payjoin::receive::v2::ProvisionalProposal> {
     use bitcoin::OutPoint;
@@ -367,13 +367,12 @@ fn try_contributing_inputs(
         .find(|i| i.txid == selected_outpoint.txid && i.vout == selected_outpoint.vout)
         .context("This shouldn't happen. Failed to retrieve the privacy preserving utxo from those we provided to the seclector.")?;
     log::debug!("selected utxo: {:#?}", selected_utxo);
-
-    //  calculate receiver payjoin outputs given receiver payjoin inputs and original_psbt,
     let txo_to_contribute = bitcoin::TxOut {
         value: selected_utxo.amount,
         script_pubkey: selected_utxo.script_pub_key.clone(),
     };
-    Ok(payjoin.contribute_witness_input(txo_to_contribute, selected_outpoint))
+
+    Ok(payjoin.contribute_witness_inputs(HashMap::from([(selected_outpoint, txo_to_contribute)])))
 }
 
 async fn unwrap_ohttp_keys_or_else_fetch(config: &AppConfig) -> Result<payjoin::OhttpKeys> {

diff --git a/payjoin/src/receive/mod.rs b/payjoin/src/receive/mod.rs
@@ -397,6 +397,7 @@ impl WantsOutputs {
             payjoin_psbt,
             params: self.params,
             owned_vouts: self.owned_vouts,
+            change_amount: None,
         })
     }
 }
@@ -408,6 +409,8 @@ pub struct WantsInputs {
     payjoin_psbt: Psbt,
     params: Params,
     owned_vouts: Vec<usize>,
+    // Input excess value to be added back as change to a receiver output
+    change_amount: Option<Amount>,
 }
 
 impl WantsInputs {
@@ -420,7 +423,7 @@ impl WantsInputs {
     /// UIH "Unnecessary input heuristic" is avoided for multi-output transactions.
     /// A simple consolidation is otherwise chosen if available.
     pub fn try_preserving_privacy(
-        &self,
+        &mut self,
         candidate_inputs: HashMap<Amount, OutPoint>,
     ) -> Result<OutPoint, SelectionError> {
         if candidate_inputs.is_empty() {
@@ -446,7 +449,7 @@ impl WantsInputs {
     // if min(in) > min(out) then UIH1 else UIH2
     // https://eprint.iacr.org/2022/589.pdf
     fn avoid_uih(
-        &self,
+        &mut self,
         candidate_inputs: HashMap<Amount, OutPoint>,
     ) -> Result<OutPoint, SelectionError> {
         let min_original_out_sats = self
@@ -475,6 +478,7 @@ impl WantsInputs {
             if candidate_min_in > candidate_min_out {
                 // The candidate avoids UIH2 but conforms to UIH1: Optimal change heuristic.
                 // It implies the smallest output is the sender's change address.
+                self.change_amount = Some(candidate_sats);
                 return Ok(candidate.1);
             }
         }
@@ -484,18 +488,24 @@ impl WantsInputs {
     }
 
     fn select_first_candidate(
-        &self,
+        &mut self,
         candidate_inputs: HashMap<Amount, OutPoint>,
     ) -> Result<OutPoint, SelectionError> {
-        candidate_inputs
-            .values()
-            .next()
-            .cloned()
-            .ok_or_else(|| SelectionError::from(InternalSelectionError::NotFound))
+        match candidate_inputs.into_iter().next() {
+            Some((amount, outpoint)) => {
+                self.change_amount = Some(amount);
+                Ok(outpoint)
+            }
+            None => Err(SelectionError::from(InternalSelectionError::NotFound)),
+        }
     }
 
-    pub fn contribute_witness_input(self, txo: TxOut, outpoint: OutPoint) -> ProvisionalProposal {
+    pub fn contribute_witness_inputs(
+        self,
+        inputs: HashMap<OutPoint, TxOut>,
+    ) -> ProvisionalProposal {
         let mut payjoin_psbt = self.payjoin_psbt.clone();
+
         // The payjoin proposal must not introduce mixed input sequence numbers
         let original_sequence = self
             .payjoin_psbt
@@ -505,26 +515,33 @@ impl WantsInputs {
             .map(|input| input.sequence)
             .unwrap_or_default();
 
-        // Add the value of new receiver input to receiver output
-        let txo_value = txo.value;
-        let vout_to_augment =
-            self.owned_vouts.choose(&mut rand::thread_rng()).expect("owned_vouts is empty");
-        payjoin_psbt.unsigned_tx.output[*vout_to_augment].value += txo_value;
+        // Add the receiver change amount to the receiver output, if applicable
+        if let Some(txo_value) = self.change_amount {
+            // TODO: ensure that owned_vouts only refers to outpoints actually owned by the
+            // receiver (e.g. not a forwarded payment)
+            let vout_to_augment =
+                self.owned_vouts.choose(&mut rand::thread_rng()).expect("owned_vouts is empty");
+            payjoin_psbt.unsigned_tx.output[*vout_to_augment].value += txo_value;
+        }
 
-        // Insert contribution at random index for privacy
+        // Insert contributions at random indices for privacy
         let mut rng = rand::thread_rng();
-        let index = rng.gen_range(0..=self.payjoin_psbt.unsigned_tx.input.len());
-        payjoin_psbt
-            .inputs
-            .insert(index, bitcoin::psbt::Input { witness_utxo: Some(txo), ..Default::default() });
-        payjoin_psbt.unsigned_tx.input.insert(
-            index,
-            bitcoin::TxIn {
-                previous_output: outpoint,
-                sequence: original_sequence,
-                ..Default::default()
-            },
-        );
+        for (outpoint, txo) in inputs.into_iter() {
+            let index = rng.gen_range(0..=self.payjoin_psbt.unsigned_tx.input.len());
+            payjoin_psbt.inputs.insert(
+                index,
+                bitcoin::psbt::Input { witness_utxo: Some(txo), ..Default::default() },
+            );
+            payjoin_psbt.unsigned_tx.input.insert(
+                index,
+                bitcoin::TxIn {
+                    previous_output: outpoint,
+                    sequence: original_sequence,
+                    ..Default::default()
+                },
+            );
+        }
+
         ProvisionalProposal {
             original_psbt: self.original_psbt,
             payjoin_psbt,

diff --git a/payjoin/src/receive/v2/mod.rs b/payjoin/src/receive/v2/mod.rs
@@ -428,14 +428,17 @@ impl WantsInputs {
     // if min(in) > min(out) then UIH1 else UIH2
     // https://eprint.iacr.org/2022/589.pdf
     pub fn try_preserving_privacy(
-        &self,
+        &mut self,
         candidate_inputs: HashMap<Amount, OutPoint>,
     ) -> Result<OutPoint, SelectionError> {
         self.inner.try_preserving_privacy(candidate_inputs)
     }
 
-    pub fn contribute_witness_input(self, txo: TxOut, outpoint: OutPoint) -> ProvisionalProposal {
-        let inner = self.inner.contribute_witness_input(txo, outpoint);
+    pub fn contribute_witness_inputs(
+        self,
+        inputs: HashMap<OutPoint, TxOut>,
+    ) -> ProvisionalProposal {
+        let inner = self.inner.contribute_witness_inputs(inputs);
         ProvisionalProposal { inner, context: self.context }
     }
 

diff --git a/payjoin/tests/integration.rs b/payjoin/tests/integration.rs
@@ -652,7 +652,7 @@ mod integration {
                 })
                 .expect("Receiver should have at least one output");
 
-            let payjoin = payjoin
+            let mut payjoin = payjoin
                 .try_substitute_receiver_outputs(None)
                 .expect("Could not substitute outputs");
 
@@ -668,17 +668,15 @@ mod integration {
                 .iter()
                 .find(|i| i.txid == selected_outpoint.txid && i.vout == selected_outpoint.vout)
                 .unwrap();
-
-            //  calculate receiver payjoin outputs given receiver payjoin inputs and original_psbt,
             let txo_to_contribute = bitcoin::TxOut {
                 value: selected_utxo.amount,
                 script_pubkey: selected_utxo.script_pub_key.clone(),
             };
-            let outpoint_to_contribute =
-                bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
-            let payjoin =
-                payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
 
+            let payjoin = payjoin
+                .contribute_witness_inputs(HashMap::from([(selected_outpoint, txo_to_contribute)]));
+
+            // Sign and finalize the proposal PSBT
             let payjoin_proposal = payjoin
                 .finalize_proposal(
                     |psbt: &Psbt| {
@@ -899,7 +897,7 @@ mod integration {
         let proposal = proposal.check_no_mixed_input_scripts().unwrap();
 
         // Receive Check 4: have we seen this input before? More of a check for non-interactive i.e. payment processor receivers.
-        let mut payjoin = proposal
+        let payjoin = proposal
             .check_no_inputs_seen_before(|_| Ok(false))
             .unwrap()
             .identify_receiver_outputs(|output_script| {
@@ -910,7 +908,7 @@ mod integration {
             })
             .expect("Receiver should have at least one output");
 
-        let payjoin = payjoin
+        let mut payjoin = payjoin
             .try_substitute_receiver_output(|| {
                 Ok(receiver.get_new_address(None, None).unwrap().assume_checked().script_pubkey())
             })
@@ -928,15 +926,13 @@ mod integration {
             .iter()
             .find(|i| i.txid == selected_outpoint.txid && i.vout == selected_outpoint.vout)
             .unwrap();
-
-        //  calculate receiver payjoin outputs given receiver payjoin inputs and original_psbt,
         let txo_to_contribute = bitcoin::TxOut {
             value: selected_utxo.amount,
             script_pubkey: selected_utxo.script_pub_key.clone(),
         };
-        let outpoint_to_contribute =
-            bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
-        let payjoin = payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
+
+        let payjoin = payjoin
+            .contribute_witness_inputs(HashMap::from([(selected_outpoint, txo_to_contribute)]));
 
         let payjoin_proposal = payjoin
             .finalize_proposal(