From 4bc8fe362226b07becd5ea8edd839bc5a4ff7a44 Mon Sep 17 00:00:00 2001
From: spacebear <144076611+grizznaut@users.noreply.github.com>
Date: Tue, 25 Jun 2024 12:00:06 -0400
Subject: [PATCH] Allow any script instead of just known address types

---
 payjoin-cli/src/app/v1.rs    |  9 +++++----
 payjoin-cli/src/app/v2.rs    |  9 +++++----
 payjoin/src/receive/mod.rs   | 11 +++++------
 payjoin/src/receive/v2.rs    |  8 ++++----
 payjoin/tests/integration.rs |  8 ++++----
 5 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/payjoin-cli/src/app/v1.rs b/payjoin-cli/src/app/v1.rs
index 88c75b10..ed9400a3 100644
--- a/payjoin-cli/src/app/v1.rs
+++ b/payjoin-cli/src/app/v1.rs
@@ -310,14 +310,15 @@ impl App {
             .map_err(|e| log::warn!("Failed to contribute inputs: {}", e));
 
         _ = provisional_payjoin
-            .try_substituting_output_address(|| {
-                bitcoind
+            .try_substitute_receiver_output(|| {
+                Ok(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| Error::Server(e.into()))?
+                    .script_pubkey())
             })
-            .map_err(|e| log::warn!("Failed to substitute output address: {}", e));
+            .map_err(|e| log::warn!("Failed to substitute output: {}", 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 db50248d..8b0b6a9b 100644
--- a/payjoin-cli/src/app/v2.rs
+++ b/payjoin-cli/src/app/v2.rs
@@ -306,14 +306,15 @@ impl App {
             .map_err(|e| log::warn!("Failed to contribute inputs: {}", e));
 
         _ = provisional_payjoin
-            .try_substituting_output_address(|| {
-                bitcoind
+            .try_substitute_receiver_output(|| {
+                Ok(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| Error::Server(e.into()))?
+                    .script_pubkey())
             })
-            .map_err(|e| log::warn!("Failed to substitute output address: {}", e));
+            .map_err(|e| log::warn!("Failed to substitute output: {}", 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 b33c2f38..58acd60e 100644
--- a/payjoin/src/receive/mod.rs
+++ b/payjoin/src/receive/mod.rs
@@ -477,17 +477,16 @@ impl ProvisionalProposal {
         self.params.disable_output_substitution
     }
 
-    /// If output substitution is enabled, replace the receiver's output address with a new one.
-    pub fn try_substituting_output_address(
+    /// If output substitution is enabled, replace the receiver's output script with a new one.
+    pub fn try_substitute_receiver_output(
         &mut self,
-        generate_address: impl Fn() -> Result<bitcoin::Address, Error>,
+        generate_script: impl Fn() -> Result<bitcoin::ScriptBuf, 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();
+        let substitute_script = generate_script()?;
+        self.payjoin_psbt.unsigned_tx.output[self.owned_vouts[0]].script_pubkey = substitute_script;
         Ok(())
     }
 
diff --git a/payjoin/src/receive/v2.rs b/payjoin/src/receive/v2.rs
index 1ba0fac5..99cb0071 100644
--- a/payjoin/src/receive/v2.rs
+++ b/payjoin/src/receive/v2.rs
@@ -387,12 +387,12 @@ impl ProvisionalProposal {
         self.inner.is_output_substitution_disabled()
     }
 
-    /// If output substitution is enabled, replace the receiver's output address with a new one.
-    pub fn try_substituting_output_address(
+    /// If output substitution is enabled, replace the receiver's output script with a new one.
+    pub fn try_substitute_receiver_output(
         &mut self,
-        generate_address: impl Fn() -> Result<bitcoin::Address, Error>,
+        generate_script: impl Fn() -> Result<bitcoin::ScriptBuf, Error>,
     ) -> Result<(), Error> {
-        self.inner.try_substituting_output_address(generate_address)
+        self.inner.try_substitute_receiver_output(generate_script)
     }
 
     pub fn finalize_proposal(
diff --git a/payjoin/tests/integration.rs b/payjoin/tests/integration.rs
index b1559c5d..119135b6 100644
--- a/payjoin/tests/integration.rs
+++ b/payjoin/tests/integration.rs
@@ -170,8 +170,8 @@ mod integration {
                 bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
             payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
 
-            _ = payjoin.try_substituting_output_address(|| {
-                Ok(receiver.get_new_address(None, None).unwrap().assume_checked())
+            _ = payjoin.try_substitute_receiver_output(|| {
+                Ok(receiver.get_new_address(None, None).unwrap().assume_checked().script_pubkey())
             });
             let payjoin_proposal = payjoin
                 .finalize_proposal(
@@ -699,8 +699,8 @@ mod integration {
                 bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
             payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
 
-            _ = payjoin.try_substituting_output_address(|| {
-                Ok(receiver.get_new_address(None, None).unwrap().assume_checked())
+            _ = payjoin.try_substitute_receiver_output(|| {
+                Ok(receiver.get_new_address(None, None).unwrap().assume_checked().script_pubkey())
             });
             let payjoin_proposal = payjoin
                 .finalize_proposal(