From 54935e5d8c36b70f0b74b82dbc964c403270d1d3 Mon Sep 17 00:00:00 2001 From: spacebear <144076611+grizznaut@users.noreply.github.com> Date: Sun, 2 Jun 2024 17:15:34 -0400 Subject: [PATCH] Refactor output substitution Always error if `disable_output_substitution` is true. --- payjoin-cli/src/app/v1.rs | 17 +++++++++-------- payjoin-cli/src/app/v2.rs | 17 +++++++++-------- payjoin/src/receive/mod.rs | 16 ++++++++++------ payjoin/src/receive/v2.rs | 13 ++++++------- payjoin/tests/integration.rs | 12 ++++++------ 5 files changed, 40 insertions(+), 35 deletions(-) diff --git a/payjoin-cli/src/app/v1.rs b/payjoin-cli/src/app/v1.rs index 0ec7a30c..6d94061d 100644 --- a/payjoin-cli/src/app/v1.rs +++ b/payjoin-cli/src/app/v1.rs @@ -305,14 +305,15 @@ impl App { _ = try_contributing_inputs(&mut provisional_payjoin, &bitcoind) .map_err(|e| log::warn!("Failed to contribute inputs: {}", e)); - if !provisional_payjoin.is_output_substitution_disabled() { - // Substitute the receiver output address. - let receiver_substitute_address = bitcoind - .get_new_address(None, None) - .map_err(|e| Error::Server(e.into()))? - .assume_checked(); - provisional_payjoin.substitute_output_address(receiver_substitute_address); - } + _ = provisional_payjoin + .try_substituting_output_address(|| { + bitcoind + .get_new_address(None, None) + .map_err(|e| Error::Server(e.into()))? + .require_network(network) + .map_err(|e| Error::Server(e.into())) + }) + .map_err(|e| log::warn!("Failed to substitute output address: {}", e)); let payjoin_proposal = provisional_payjoin.finalize_proposal( |psbt: &Psbt| { diff --git a/payjoin-cli/src/app/v2.rs b/payjoin-cli/src/app/v2.rs index 889e54aa..ca08c368 100644 --- a/payjoin-cli/src/app/v2.rs +++ b/payjoin-cli/src/app/v2.rs @@ -275,14 +275,15 @@ impl App { _ = try_contributing_inputs(&mut provisional_payjoin.inner, &bitcoind) .map_err(|e| log::warn!("Failed to contribute inputs: {}", e)); - if !provisional_payjoin.is_output_substitution_disabled() { - // Substitute the receiver output address. - let receiver_substitute_address = bitcoind - .get_new_address(None, None) - .map_err(|e| Error::Server(e.into()))? - .assume_checked(); - provisional_payjoin.substitute_output_address(receiver_substitute_address); - } + _ = provisional_payjoin + .try_substituting_output_address(|| { + bitcoind + .get_new_address(None, None) + .map_err(|e| Error::Server(e.into()))? + .require_network(network) + .map_err(|e| Error::Server(e.into())) + }) + .map_err(|e| log::warn!("Failed to substitute output address: {}", e)); let payjoin_proposal = provisional_payjoin.finalize_proposal( |psbt: &Psbt| { diff --git a/payjoin/src/receive/mod.rs b/payjoin/src/receive/mod.rs index 2ea09ee7..ac584d69 100644 --- a/payjoin/src/receive/mod.rs +++ b/payjoin/src/receive/mod.rs @@ -473,14 +473,18 @@ impl ProvisionalProposal { ); } - pub fn is_output_substitution_disabled(&self) -> bool { - self.params.disable_output_substitution - } - - /// Just replace an output address with - pub fn substitute_output_address(&mut self, substitute_address: bitcoin::Address) { + /// If output substitution is enabled, replace the receiver's output address with a new one. + pub fn try_substituting_output_address( + &mut self, + generate_address: impl Fn() -> Result, + ) -> Result<(), Error> { + if self.params.disable_output_substitution { + return Err(Error::Server("Output substitution is disabled.".into())); + } + let substitute_address = generate_address()?; self.payjoin_psbt.unsigned_tx.output[self.owned_vouts[0]].script_pubkey = substitute_address.script_pubkey(); + Ok(()) } /// Apply additional fee contribution now that the receiver has contributed input diff --git a/payjoin/src/receive/v2.rs b/payjoin/src/receive/v2.rs index 25a8bb6e..5c3a4750 100644 --- a/payjoin/src/receive/v2.rs +++ b/payjoin/src/receive/v2.rs @@ -460,13 +460,12 @@ impl ProvisionalProposal { self.inner.contribute_non_witness_input(tx, outpoint) } - pub fn is_output_substitution_disabled(&self) -> bool { - self.inner.is_output_substitution_disabled() - } - - /// Just replace an output address with - pub fn substitute_output_address(&mut self, substitute_address: bitcoin::Address) { - self.inner.substitute_output_address(substitute_address) + /// If output substitution is enabled, replace the receiver's output address with a new one. + pub fn try_substituting_output_address( + &mut self, + generate_address: impl Fn() -> Result, + ) -> Result<(), Error> { + self.inner.try_substituting_output_address(generate_address) } pub fn finalize_proposal( diff --git a/payjoin/tests/integration.rs b/payjoin/tests/integration.rs index 1549a02f..97c6e323 100644 --- a/payjoin/tests/integration.rs +++ b/payjoin/tests/integration.rs @@ -166,9 +166,9 @@ mod integration { bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout }; payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute); - let receiver_substitute_address = - receiver.get_new_address(None, None).unwrap().assume_checked(); - payjoin.substitute_output_address(receiver_substitute_address); + _ = payjoin.try_substituting_output_address(|| { + Ok(receiver.get_new_address(None, None).unwrap().assume_checked()) + }); let payjoin_proposal = payjoin .finalize_proposal( |psbt: &Psbt| { @@ -692,9 +692,9 @@ mod integration { bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout }; payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute); - let receiver_substitute_address = - receiver.get_new_address(None, None).unwrap().assume_checked(); - payjoin.substitute_output_address(receiver_substitute_address); + _ = payjoin.try_substituting_output_address(|| { + Ok(receiver.get_new_address(None, None).unwrap().assume_checked()) + }); let payjoin_proposal = payjoin .finalize_proposal( |psbt: &Psbt| {