Refactor output substitution

Always error if `disable_output_substitution` is true.

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| {

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| {

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<bitcoin::Address, Error>,
+    ) -> 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

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<bitcoin::Address, Error>,
+    ) -> Result<(), Error> {
+        self.inner.try_substituting_output_address(generate_address)
     }
 
     pub fn finalize_proposal(

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| {