From 65720a7471bd65f025473afed9bd3067a3306688 Mon Sep 17 00:00:00 2001
From: spacebear <git@spacebear.dev>
Date: Tue, 30 Jul 2024 15:54:21 -0400
Subject: [PATCH] Add new typestates for output and input contributions

This effectively splits the `ProvisionalProposal` state into three
states that must be completed in order:
- `WantsOutputs` allows the receiver to substitute or add output(s) and
produces a WantsInputs
- `WantsInputs` allows the receiver to contribute input(s) as needed to
fund their outputs and produces a ProvisionalProposal
- `ProvisionalProposal` is only responsible for finalizing the payjoin
proposal and producing a PSBT that will be acceptable to the sender
---
 payjoin-cli/src/app/mod.rs    |  32 --------
 payjoin-cli/src/app/v1.rs     |  57 +++++++++----
 payjoin-cli/src/app/v2.rs     |  42 ++++++++--
 payjoin/src/receive/mod.rs    | 150 +++++++++++++++++++++++++++-------
 payjoin/src/receive/v2/mod.rs |  76 ++++++++++++-----
 payjoin/tests/integration.rs  |  30 ++++---
 6 files changed, 277 insertions(+), 110 deletions(-)

diff --git a/payjoin-cli/src/app/mod.rs b/payjoin-cli/src/app/mod.rs
index 2bc54f6e..35c33e9e 100644
--- a/payjoin-cli/src/app/mod.rs
+++ b/payjoin-cli/src/app/mod.rs
@@ -94,38 +94,6 @@ pub trait App {
     }
 }
 
-fn try_contributing_inputs(
-    payjoin: &mut payjoin::receive::ProvisionalProposal,
-    bitcoind: &bitcoincore_rpc::Client,
-) -> Result<()> {
-    use bitcoin::OutPoint;
-
-    let available_inputs = bitcoind
-        .list_unspent(None, None, None, None, None)
-        .context("Failed to list unspent from bitcoind")?;
-    let candidate_inputs: HashMap<Amount, OutPoint> = available_inputs
-        .iter()
-        .map(|i| (i.amount, OutPoint { txid: i.txid, vout: i.vout }))
-        .collect();
-
-    let selected_outpoint = payjoin.try_preserving_privacy(candidate_inputs).expect("gg");
-    let selected_utxo = available_inputs
-        .iter()
-        .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.to_sat(),
-        script_pubkey: selected_utxo.script_pub_key.clone(),
-    };
-    let outpoint_to_contribute =
-        bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
-    payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
-    Ok(())
-}
-
 struct Headers<'a>(&'a hyper::HeaderMap);
 impl payjoin::receive::Headers for Headers<'_> {
     fn get_header(&self, key: &str) -> Option<&str> {
diff --git a/payjoin-cli/src/app/v1.rs b/payjoin-cli/src/app/v1.rs
index 37a5eed0..cd9fc0ac 100644
--- a/payjoin-cli/src/app/v1.rs
+++ b/payjoin-cli/src/app/v1.rs
@@ -15,7 +15,7 @@ use payjoin::{Error, PjUriBuilder, Uri, UriExt};
 
 use super::config::AppConfig;
 use super::App as AppTrait;
-use crate::app::{http_agent, try_contributing_inputs, Headers};
+use crate::app::{http_agent, Headers};
 use crate::db::Database;
 #[cfg(feature = "danger-local-https")]
 pub const LOCAL_CERT_FILE: &str = "localhost.der";
@@ -297,7 +297,7 @@ impl App {
         })?;
         log::trace!("check4");
 
-        let mut provisional_payjoin = payjoin.identify_receiver_outputs(|output_script| {
+        let provisional_payjoin = payjoin.identify_receiver_outputs(|output_script| {
             if let Ok(address) = bitcoin::Address::from_script(output_script, network) {
                 bitcoind
                     .get_address_info(&address)
@@ -308,19 +308,17 @@ impl App {
             }
         })?;
 
-        _ = try_contributing_inputs(&mut provisional_payjoin, &bitcoind)
-            .map_err(|e| log::warn!("Failed to contribute inputs: {}", e));
+        let provisional_payjoin = provisional_payjoin.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()))?
+                .script_pubkey())
+        })?;
 
-        _ = provisional_payjoin
-            .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()))?
-                    .script_pubkey())
-            })
-            .map_err(|e| log::warn!("Failed to substitute output: {}", e));
+        let provisional_payjoin = try_contributing_inputs(provisional_payjoin, &bitcoind)
+            .map_err(|e| Error::Server(e.into()))?;
 
         let payjoin_proposal = provisional_payjoin.finalize_proposal(
             |psbt: &Psbt| {
@@ -339,3 +337,34 @@ impl App {
         Ok(payjoin_proposal)
     }
 }
+
+fn try_contributing_inputs(
+    payjoin: payjoin::receive::WantsInputs,
+    bitcoind: &bitcoincore_rpc::Client,
+) -> Result<payjoin::receive::ProvisionalProposal> {
+    use bitcoin::OutPoint;
+
+    let available_inputs = bitcoind
+        .list_unspent(None, None, None, None, None)
+        .context("Failed to list unspent from bitcoind")?;
+    let candidate_inputs: HashMap<Amount, OutPoint> = available_inputs
+        .iter()
+        .map(|i| (i.amount, OutPoint { txid: i.txid, vout: i.vout }))
+        .collect();
+
+    let selected_outpoint = payjoin.try_preserving_privacy(candidate_inputs).expect("gg");
+    let selected_utxo = available_inputs
+        .iter()
+        .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.to_sat(),
+        script_pubkey: selected_utxo.script_pub_key.clone(),
+    };
+    let outpoint_to_contribute =
+        bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
+    Ok(payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute))
+}
diff --git a/payjoin-cli/src/app/v2.rs b/payjoin-cli/src/app/v2.rs
index 8beb504c..f6791e8a 100644
--- a/payjoin-cli/src/app/v2.rs
+++ b/payjoin-cli/src/app/v2.rs
@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::str::FromStr;
 use std::sync::Arc;
 
@@ -272,8 +273,6 @@ impl App {
         &self,
         proposal: payjoin::receive::v2::UncheckedProposal,
     ) -> Result<payjoin::receive::v2::PayjoinProposal, Error> {
-        use crate::app::try_contributing_inputs;
-
         let bitcoind = self.bitcoind().map_err(|e| Error::Server(e.into()))?;
 
         // in a payment processor where the sender could go offline, this is where you schedule to broadcast the original_tx
@@ -323,7 +322,7 @@ impl App {
         })?;
         log::trace!("check4");
 
-        let mut provisional_payjoin = payjoin.identify_receiver_outputs(|output_script| {
+        let provisional_payjoin = payjoin.identify_receiver_outputs(|output_script| {
             if let Ok(address) = bitcoin::Address::from_script(output_script, network) {
                 bitcoind
                     .get_address_info(&address)
@@ -334,8 +333,10 @@ impl App {
             }
         })?;
 
-        _ = try_contributing_inputs(&mut provisional_payjoin.inner, &bitcoind)
-            .map_err(|e| log::warn!("Failed to contribute inputs: {}", e));
+        let provisional_payjoin = provisional_payjoin.try_substitute_receiver_outputs(None)?;
+
+        let provisional_payjoin = try_contributing_inputs(provisional_payjoin, &bitcoind)
+            .map_err(|e| Error::Server(e.into()))?;
 
         let payjoin_proposal = provisional_payjoin.finalize_proposal(
             |psbt: &Psbt| {
@@ -352,6 +353,37 @@ impl App {
     }
 }
 
+fn try_contributing_inputs(
+    payjoin: payjoin::receive::v2::WantsInputs,
+    bitcoind: &bitcoincore_rpc::Client,
+) -> Result<payjoin::receive::v2::ProvisionalProposal> {
+    use bitcoin::OutPoint;
+
+    let available_inputs = bitcoind
+        .list_unspent(None, None, None, None, None)
+        .context("Failed to list unspent from bitcoind")?;
+    let candidate_inputs: HashMap<Amount, OutPoint> = available_inputs
+        .iter()
+        .map(|i| (i.amount, OutPoint { txid: i.txid, vout: i.vout }))
+        .collect();
+
+    let selected_outpoint = payjoin.try_preserving_privacy(candidate_inputs).expect("gg");
+    let selected_utxo = available_inputs
+        .iter()
+        .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.to_sat(),
+        script_pubkey: selected_utxo.script_pub_key.clone(),
+    };
+    let outpoint_to_contribute =
+        bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
+    Ok(payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute))
+}
+
 async fn unwrap_ohttp_keys_or_else_fetch(config: &AppConfig) -> Result<payjoin::OhttpKeys> {
     if let Some(keys) = config.ohttp_keys.clone() {
         println!("Using OHTTP Keys from config");
diff --git a/payjoin/src/receive/mod.rs b/payjoin/src/receive/mod.rs
index 794fa793..ea05e234 100644
--- a/payjoin/src/receive/mod.rs
+++ b/payjoin/src/receive/mod.rs
@@ -305,7 +305,7 @@ impl OutputsUnknown {
     pub fn identify_receiver_outputs(
         self,
         is_receiver_output: impl Fn(&Script) -> Result<bool, Error>,
-    ) -> Result<ProvisionalProposal, Error> {
+    ) -> Result<WantsOutputs, Error> {
         let owned_vouts: Vec<usize> = self
             .psbt
             .unsigned_tx
@@ -323,7 +323,7 @@ impl OutputsUnknown {
             return Err(Error::BadRequest(InternalRequestError::MissingPayment.into()));
         }
 
-        Ok(ProvisionalProposal {
+        Ok(WantsOutputs {
             original_psbt: self.psbt.clone(),
             payjoin_psbt: self.psbt,
             params: self.params,
@@ -332,16 +332,83 @@ impl OutputsUnknown {
     }
 }
 
-/// A mutable checked proposal that the receiver may contribute inputs to to make a payjoin.
+/// A checked proposal that the receiver may substitute or add outputs to
 #[derive(Debug, Clone)]
-pub struct ProvisionalProposal {
+pub struct WantsOutputs {
     original_psbt: Psbt,
     payjoin_psbt: Psbt,
     params: Params,
     owned_vouts: Vec<usize>,
 }
 
-impl ProvisionalProposal {
+impl WantsOutputs {
+    pub fn is_output_substitution_disabled(&self) -> bool {
+        self.params.disable_output_substitution
+    }
+
+    /// If output substitution is enabled, replace the receiver's output script with a new one.
+    pub fn try_substitute_receiver_output(
+        self,
+        generate_script: impl Fn() -> Result<bitcoin::ScriptBuf, Error>,
+    ) -> Result<WantsInputs, Error> {
+        let output_value = self.payjoin_psbt.unsigned_tx.output[self.owned_vouts[0]].value;
+        let outputs = vec![TxOut { value: output_value, script_pubkey: generate_script()? }];
+        self.try_substitute_receiver_outputs(Some(outputs))
+    }
+
+    pub fn try_substitute_receiver_outputs(
+        self,
+        outputs: Option<Vec<TxOut>>,
+    ) -> Result<WantsInputs, Error> {
+        let mut payjoin_psbt = self.payjoin_psbt.clone();
+        match outputs {
+            Some(o) => {
+                if self.params.disable_output_substitution {
+                    return Err(Error::Server("Output substitution is disabled.".into()));
+                }
+                let mut replacement_outputs = o.into_iter();
+                let mut outputs = vec![];
+                for (i, output) in self.payjoin_psbt.unsigned_tx.output.iter().enumerate() {
+                    if self.owned_vouts.contains(&i) {
+                        // Receiver output: substitute with a provided output
+                        // TODO: pick from outputs in random order?
+                        outputs.push(
+                            replacement_outputs
+                                .next()
+                                .ok_or(Error::Server("Not enough outputs".into()))?,
+                        );
+                    } else {
+                        // Sender output: leave it as is
+                        outputs.push(output.clone());
+                    }
+                }
+                // Append all remaining outputs
+                outputs.extend(replacement_outputs);
+                payjoin_psbt.unsigned_tx.output = outputs;
+                // TODO: update self.owned_vouts?
+                // TODO: is tx funded or does it need more inputs?
+            }
+            None => log::info!("No outputs provided: skipping output substitution."),
+        }
+        Ok(WantsInputs {
+            original_psbt: self.original_psbt,
+            payjoin_psbt,
+            params: self.params,
+            owned_vouts: self.owned_vouts,
+        })
+    }
+}
+
+/// A checked proposal that the receiver may contribute inputs to to make a payjoin
+#[derive(Debug, Clone)]
+pub struct WantsInputs {
+    original_psbt: Psbt,
+    payjoin_psbt: Psbt,
+    params: Params,
+    owned_vouts: Vec<usize>,
+}
+
+impl WantsInputs {
     /// Select receiver input such that the payjoin avoids surveillance.
     /// Return the input chosen that has been applied to the Proposal.
     ///
@@ -425,7 +492,8 @@ impl ProvisionalProposal {
             .ok_or_else(|| SelectionError::from(InternalSelectionError::NotFound))
     }
 
-    pub fn contribute_witness_input(&mut self, txo: TxOut, outpoint: OutPoint) {
+    pub fn contribute_witness_input(self, txo: TxOut, outpoint: OutPoint) -> 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
@@ -439,15 +507,15 @@ impl ProvisionalProposal {
         let txo_value = txo.value;
         let vout_to_augment =
             self.owned_vouts.choose(&mut rand::thread_rng()).expect("owned_vouts is empty");
-        self.payjoin_psbt.unsigned_tx.output[*vout_to_augment].value += txo_value;
+        payjoin_psbt.unsigned_tx.output[*vout_to_augment].value += txo_value;
 
         // Insert contribution at random index for privacy
         let mut rng = rand::thread_rng();
         let index = rng.gen_range(0..=self.payjoin_psbt.unsigned_tx.input.len());
-        self.payjoin_psbt
+        payjoin_psbt
             .inputs
             .insert(index, bitcoin::psbt::Input { witness_utxo: Some(txo), ..Default::default() });
-        self.payjoin_psbt.unsigned_tx.input.insert(
+        payjoin_psbt.unsigned_tx.input.insert(
             index,
             bitcoin::TxIn {
                 previous_output: outpoint,
@@ -455,9 +523,20 @@ impl ProvisionalProposal {
                 ..Default::default()
             },
         );
+        ProvisionalProposal {
+            original_psbt: self.original_psbt,
+            payjoin_psbt,
+            params: self.params,
+            owned_vouts: self.owned_vouts,
+        }
     }
 
-    pub fn contribute_non_witness_input(&mut self, tx: bitcoin::Transaction, outpoint: OutPoint) {
+    pub fn contribute_non_witness_input(
+        self,
+        tx: bitcoin::Transaction,
+        outpoint: OutPoint,
+    ) -> 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
@@ -471,18 +550,18 @@ impl ProvisionalProposal {
         let txo_value = tx.output[outpoint.vout as usize].value;
         let vout_to_augment =
             self.owned_vouts.choose(&mut rand::thread_rng()).expect("owned_vouts is empty");
-        self.payjoin_psbt.unsigned_tx.output[*vout_to_augment].value += txo_value;
+        payjoin_psbt.unsigned_tx.output[*vout_to_augment].value += txo_value;
 
         // Insert contribution at random index for privacy
         let mut rng = rand::thread_rng();
         let index = rng.gen_range(0..=self.payjoin_psbt.unsigned_tx.input.len());
 
         // Add the new input to the PSBT
-        self.payjoin_psbt.inputs.insert(
+        payjoin_psbt.inputs.insert(
             index,
             bitcoin::psbt::Input { non_witness_utxo: Some(tx), ..Default::default() },
         );
-        self.payjoin_psbt.unsigned_tx.input.insert(
+        payjoin_psbt.unsigned_tx.input.insert(
             index,
             bitcoin::TxIn {
                 previous_output: outpoint,
@@ -490,25 +569,36 @@ impl ProvisionalProposal {
                 ..Default::default()
             },
         );
+        ProvisionalProposal {
+            original_psbt: self.original_psbt,
+            payjoin_psbt,
+            params: self.params,
+            owned_vouts: self.owned_vouts,
+        }
     }
 
-    pub fn is_output_substitution_disabled(&self) -> bool {
-        self.params.disable_output_substitution
-    }
-
-    /// If output substitution is enabled, replace the receiver's output script with a new one.
-    pub fn try_substitute_receiver_output(
-        &mut self,
-        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()));
+    // TODO: temporary workaround
+    fn skip_contribute_inputs(self) -> ProvisionalProposal {
+        ProvisionalProposal {
+            original_psbt: self.original_psbt,
+            payjoin_psbt: self.payjoin_psbt,
+            params: self.params,
+            owned_vouts: self.owned_vouts,
         }
-        let substitute_script = generate_script()?;
-        self.payjoin_psbt.unsigned_tx.output[self.owned_vouts[0]].script_pubkey = substitute_script;
-        Ok(())
     }
+}
 
+/// A checked proposal that the receiver may sign and finalize to make a proposal PSBT that the
+/// sender will accept.
+#[derive(Debug, Clone)]
+pub struct ProvisionalProposal {
+    original_psbt: Psbt,
+    payjoin_psbt: Psbt,
+    params: Params,
+    owned_vouts: Vec<usize>,
+}
+
+impl ProvisionalProposal {
     /// Apply additional fee contribution now that the receiver has contributed input
     /// this is kind of a "build_proposal" step before we sign and finalize and extract
     ///
@@ -720,7 +810,11 @@ mod test {
                         .unwrap()
                         .require_network(network))
             })
-            .expect("Receiver output should be identified");
+            .expect("Receiver output should be identified")
+            .try_substitute_receiver_outputs(None)
+            .expect("Substitute outputs should do nothing")
+            .skip_contribute_inputs(); // TODO: temporary workaround
+
         let payjoin = payjoin.apply_fee(None);
 
         assert!(payjoin.is_ok(), "Payjoin should be a valid PSBT");
diff --git a/payjoin/src/receive/v2/mod.rs b/payjoin/src/receive/v2/mod.rs
index c040b244..1401f8e4 100644
--- a/payjoin/src/receive/v2/mod.rs
+++ b/payjoin/src/receive/v2/mod.rs
@@ -370,20 +370,50 @@ impl OutputsUnknown {
     pub fn identify_receiver_outputs(
         self,
         is_receiver_output: impl Fn(&Script) -> Result<bool, Error>,
-    ) -> Result<ProvisionalProposal, Error> {
+    ) -> Result<WantsOutputs, Error> {
         let inner = self.inner.identify_receiver_outputs(is_receiver_output)?;
-        Ok(ProvisionalProposal { inner, context: self.context })
+        Ok(WantsOutputs { inner, context: self.context })
     }
 }
 
-/// A mutable checked proposal that the receiver may contribute inputs to to make a payjoin.
+/// A checked proposal that the receiver may substitute or add outputs to
 #[derive(Debug, Clone)]
-pub struct ProvisionalProposal {
-    pub inner: super::ProvisionalProposal,
+pub struct WantsOutputs {
+    inner: super::WantsOutputs,
     context: SessionContext,
 }
 
-impl ProvisionalProposal {
+impl WantsOutputs {
+    pub fn is_output_substitution_disabled(&self) -> bool {
+        self.inner.is_output_substitution_disabled()
+    }
+
+    /// If output substitution is enabled, replace the receiver's output script with a new one.
+    pub fn try_substitute_receiver_output(
+        self,
+        generate_script: impl Fn() -> Result<bitcoin::ScriptBuf, Error>,
+    ) -> Result<WantsInputs, Error> {
+        let inner = self.inner.try_substitute_receiver_output(generate_script)?;
+        Ok(WantsInputs { inner, context: self.context })
+    }
+
+    pub fn try_substitute_receiver_outputs(
+        self,
+        generate_outputs: Option<Vec<TxOut>>,
+    ) -> Result<WantsInputs, Error> {
+        let inner = self.inner.try_substitute_receiver_outputs(generate_outputs)?;
+        Ok(WantsInputs { inner, context: self.context })
+    }
+}
+
+/// A checked proposal that the receiver may contribute inputs to to make a payjoin
+#[derive(Debug, Clone)]
+pub struct WantsInputs {
+    inner: super::WantsInputs,
+    context: SessionContext,
+}
+
+impl WantsInputs {
     /// Select receiver input such that the payjoin avoids surveillance.
     /// Return the input chosen that has been applied to the Proposal.
     ///
@@ -402,26 +432,30 @@ impl ProvisionalProposal {
         self.inner.try_preserving_privacy(candidate_inputs)
     }
 
-    pub fn contribute_witness_input(&mut self, txo: TxOut, outpoint: OutPoint) {
-        self.inner.contribute_witness_input(txo, outpoint)
+    pub fn contribute_witness_input(self, txo: TxOut, outpoint: OutPoint) -> ProvisionalProposal {
+        let inner = self.inner.contribute_witness_input(txo, outpoint);
+        ProvisionalProposal { inner, context: self.context }
     }
 
-    pub fn contribute_non_witness_input(&mut self, tx: bitcoin::Transaction, outpoint: OutPoint) {
-        self.inner.contribute_non_witness_input(tx, outpoint)
-    }
-
-    pub fn is_output_substitution_disabled(&self) -> bool {
-        self.inner.is_output_substitution_disabled()
+    pub fn contribute_non_witness_input(
+        self,
+        tx: bitcoin::Transaction,
+        outpoint: OutPoint,
+    ) -> ProvisionalProposal {
+        let inner = self.inner.contribute_non_witness_input(tx, outpoint);
+        ProvisionalProposal { inner, context: self.context }
     }
+}
 
-    /// If output substitution is enabled, replace the receiver's output script with a new one.
-    pub fn try_substitute_receiver_output(
-        &mut self,
-        generate_script: impl Fn() -> Result<bitcoin::ScriptBuf, Error>,
-    ) -> Result<(), Error> {
-        self.inner.try_substitute_receiver_output(generate_script)
-    }
+/// A checked proposal that the receiver may sign and finalize to make a proposal PSBT that the
+/// sender will accept.
+#[derive(Debug, Clone)]
+pub struct ProvisionalProposal {
+    inner: super::ProvisionalProposal,
+    context: SessionContext,
+}
 
+impl ProvisionalProposal {
     pub fn finalize_proposal(
         self,
         wallet_process_psbt: impl Fn(&Psbt) -> Result<Psbt, Error>,
diff --git a/payjoin/tests/integration.rs b/payjoin/tests/integration.rs
index 62385acf..ddabb044 100644
--- a/payjoin/tests/integration.rs
+++ b/payjoin/tests/integration.rs
@@ -138,7 +138,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| {
@@ -149,6 +149,16 @@ mod integration {
                 })
                 .expect("Receiver should have at least one output");
 
+            let payjoin = payjoin
+                .try_substitute_receiver_output(|| {
+                    Ok(receiver
+                        .get_new_address(None, None)
+                        .unwrap()
+                        .assume_checked()
+                        .script_pubkey())
+                })
+                .expect("Could not substitute outputs");
+
             // Select receiver payjoin inputs. TODO Lock them.
             let available_inputs = receiver.list_unspent(None, None, None, None, None).unwrap();
             let candidate_inputs: HashMap<Amount, OutPoint> = available_inputs
@@ -169,11 +179,9 @@ mod integration {
             };
             let outpoint_to_contribute =
                 bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
-            payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
+            let payjoin =
+                payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
 
-            _ = payjoin.try_substitute_receiver_output(|| {
-                Ok(receiver.get_new_address(None, None).unwrap().assume_checked().script_pubkey())
-            });
             let payjoin_proposal = payjoin
                 .finalize_proposal(
                     |psbt: &Psbt| {
@@ -737,7 +745,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| {
@@ -748,6 +756,10 @@ mod integration {
                 })
                 .expect("Receiver should have at least one output");
 
+            let payjoin = payjoin
+                .try_substitute_receiver_outputs(None)
+                .expect("Could not substitute outputs");
+
             // Select receiver payjoin inputs. TODO Lock them.
             let available_inputs = receiver.list_unspent(None, None, None, None, None).unwrap();
             let candidate_inputs: HashMap<Amount, OutPoint> = available_inputs
@@ -768,11 +780,9 @@ mod integration {
             };
             let outpoint_to_contribute =
                 bitcoin::OutPoint { txid: selected_utxo.txid, vout: selected_utxo.vout };
-            payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
+            let payjoin =
+                payjoin.contribute_witness_input(txo_to_contribute, outpoint_to_contribute);
 
-            _ = payjoin.try_substitute_receiver_output(|| {
-                Ok(receiver.get_new_address(None, None).unwrap().assume_checked().script_pubkey())
-            });
             let payjoin_proposal = payjoin
                 .finalize_proposal(
                     |psbt: &Psbt| {