Add custom error types for output and input contributions

spacebear21 · spacebear21 · commit 455909ef54f4 · 2024-08-28T18:41:54.000-04:00
Partially addresses payjoin#312
diff --git a/payjoin-cli/src/app/v1.rs b/payjoin-cli/src/app/v1.rs
@@ -348,7 +348,8 @@ impl App {
                     .require_network(network)
                     .map_err(|e| Error::Server(e.into()))?
                     .script_pubkey(),
-            )?
+            )
+            .map_err(|e| Error::Server(e.into()))?
             .commit_outputs();
 
         let provisional_payjoin = try_contributing_inputs(payjoin.clone(), &bitcoind)
@@ -397,6 +398,7 @@ fn try_contributing_inputs(
 
     Ok(payjoin
         .contribute_witness_inputs(vec![(selected_outpoint, txo_to_contribute)])
+        .expect("This shouldn't happen. Failed to contribute inputs.")
         .commit_inputs())
 }
 
diff --git a/payjoin-cli/src/app/v2.rs b/payjoin-cli/src/app/v2.rs
@@ -377,6 +377,7 @@ fn try_contributing_inputs(
 
     Ok(payjoin
         .contribute_witness_inputs(vec![(selected_outpoint, txo_to_contribute)])
+        .expect("This shouldn't happen. Failed to contribute inputs.")
         .commit_inputs())
 }
 
diff --git a/payjoin/src/receive/error.rs b/payjoin/src/receive/error.rs
@@ -196,6 +196,51 @@ impl std::error::Error for RequestError {
     }
 }
 
+/// Error that may occur when output substitution fails.
+///
+/// This is currently opaque type because we aren't sure which variants will stay.
+/// You can only display it.
+#[derive(Debug)]
+pub struct OutputSubstitutionError(InternalOutputSubstitutionError);
+
+#[derive(Debug)]
+pub(crate) enum InternalOutputSubstitutionError {
+    /// Output substitution is disabled
+    OutputSubstitutionDisabled(&'static str),
+    /// Current output substitution implementation doesn't support reducing the number of outputs
+    NotEnoughOutputs,
+    /// The provided drain script could not be identified in the provided replacement outputs
+    InvalidDrainScript,
+}
+
+impl fmt::Display for OutputSubstitutionError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match &self.0 {
+            InternalOutputSubstitutionError::OutputSubstitutionDisabled(reason) => write!(f, "{}", &format!("Output substitution is disabled: {}", reason)),
+            InternalOutputSubstitutionError::NotEnoughOutputs => write!(
+                f,
+                "Current output substitution implementation doesn't support reducing the number of outputs"
+            ),
+            InternalOutputSubstitutionError::InvalidDrainScript =>
+                write!(f, "The provided drain script could not be identified in the provided replacement outputs"),
+        }
+    }
+}
+
+impl From<InternalOutputSubstitutionError> for OutputSubstitutionError {
+    fn from(value: InternalOutputSubstitutionError) -> Self { OutputSubstitutionError(value) }
+}
+
+impl std::error::Error for OutputSubstitutionError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match &self.0 {
+            InternalOutputSubstitutionError::OutputSubstitutionDisabled(_) => None,
+            InternalOutputSubstitutionError::NotEnoughOutputs => None,
+            InternalOutputSubstitutionError::InvalidDrainScript => None,
+        }
+    }
+}
+
 /// Error that may occur when coin selection fails.
 ///
 /// This is currently opaque type because we aren't sure which variants will stay.
@@ -230,3 +275,29 @@ impl fmt::Display for SelectionError {
 impl From<InternalSelectionError> for SelectionError {
     fn from(value: InternalSelectionError) -> Self { SelectionError(value) }
 }
+
+/// Error that may occur when input contribution fails.
+///
+/// This is currently opaque type because we aren't sure which variants will stay.
+/// You can only display it.
+#[derive(Debug)]
+pub struct InputContributionError(InternalInputContributionError);
+
+#[derive(Debug)]
+pub(crate) enum InternalInputContributionError {
+    /// Total input value is not enough to cover additional output value
+    ValueTooLow,
+}
+
+impl fmt::Display for InputContributionError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match &self.0 {
+            InternalInputContributionError::ValueTooLow =>
+                write!(f, "Total input value is not enough to cover additional output value"),
+        }
+    }
+}
+
+impl From<InternalInputContributionError> for InputContributionError {
+    fn from(value: InternalInputContributionError) -> Self { InputContributionError(value) }
+}
diff --git a/payjoin/src/receive/mod.rs b/payjoin/src/receive/mod.rs
@@ -38,8 +38,11 @@ mod optional_parameters;
 pub mod v2;
 
 use bitcoin::secp256k1::rand::{self, Rng};
-pub use error::{Error, RequestError, SelectionError};
-use error::{InternalRequestError, InternalSelectionError};
+pub use error::{Error, OutputSubstitutionError, RequestError, SelectionError};
+use error::{
+    InputContributionError, InternalInputContributionError, InternalOutputSubstitutionError,
+    InternalRequestError, InternalSelectionError,
+};
 use optional_parameters::Params;
 
 use crate::input_type::InputType;
@@ -359,7 +362,10 @@ impl WantsOutputs {
     }
 
     /// Substitute the receiver output script with the provided script.
-    pub fn substitute_receiver_script(self, output_script: &Script) -> Result<WantsOutputs, Error> {
+    pub fn substitute_receiver_script(
+        self,
+        output_script: &Script,
+    ) -> Result<WantsOutputs, OutputSubstitutionError> {
         let output_value = self.original_psbt.unsigned_tx.output[self.change_vout].value;
         let outputs = vec![TxOut { value: output_value, script_pubkey: output_script.into() }];
         self.replace_receiver_outputs(outputs, output_script)
@@ -374,7 +380,7 @@ impl WantsOutputs {
         self,
         replacement_outputs: Vec<TxOut>,
         drain_script: &Script,
-    ) -> Result<WantsOutputs, Error> {
+    ) -> Result<WantsOutputs, OutputSubstitutionError> {
         let mut payjoin_psbt = self.original_psbt.clone();
         let mut change_vout: Option<usize> = None;
         if self.params.disable_output_substitution {
@@ -387,14 +393,16 @@ impl WantsOutputs {
                 .find(|txo| txo.script_pubkey == original_output.script_pubkey)
             {
                 Some(txo) if txo.value < original_output.value => {
-                    return Err(Error::Server(
-                        "Decreasing the receiver output value is not allowed".into(),
-                    ));
+                    return Err(InternalOutputSubstitutionError::OutputSubstitutionDisabled(
+                        "Decreasing the receiver output value is not allowed",
+                    )
+                    .into());
                 }
                 None =>
-                    return Err(Error::Server(
-                        "Changing the receiver output script pubkey is not allowed".into(),
-                    )),
+                    return Err(InternalOutputSubstitutionError::OutputSubstitutionDisabled(
+                        "Changing the receiver output script pubkey is not allowed",
+                    )
+                    .into()),
                 _ => log::info!("Receiver is augmenting outputs"),
             }
         }
@@ -406,7 +414,7 @@ impl WantsOutputs {
             if self.owned_vouts.contains(&i) {
                 // Receiver output: substitute with a provided output picked randomly
                 if replacement_outputs.is_empty() {
-                    return Err(Error::Server("Not enough outputs".into()));
+                    return Err(InternalOutputSubstitutionError::NotEnoughOutputs.into());
                 }
                 let index = rng.gen_range(0..replacement_outputs.len());
                 let txo = replacement_outputs.swap_remove(index);
@@ -435,7 +443,7 @@ impl WantsOutputs {
             original_psbt: self.original_psbt,
             payjoin_psbt,
             params: self.params,
-            change_vout: change_vout.ok_or(Error::Server("Invalid drain script".into()))?,
+            change_vout: change_vout.ok_or(InternalOutputSubstitutionError::InvalidDrainScript)?,
             owned_vouts: self.owned_vouts,
         })
     }
@@ -475,13 +483,13 @@ impl WantsInputs {
         candidate_inputs: HashMap<Amount, OutPoint>,
     ) -> Result<OutPoint, SelectionError> {
         if candidate_inputs.is_empty() {
-            return Err(SelectionError::from(InternalSelectionError::Empty));
+            return Err(InternalSelectionError::Empty.into());
         }
 
         if self.payjoin_psbt.outputs.len() > 2 {
             // This UIH avoidance function supports only
             // many-input, n-output transactions such that n <= 2 for now
-            return Err(SelectionError::from(InternalSelectionError::TooManyOutputs));
+            return Err(InternalSelectionError::TooManyOutputs.into());
         }
 
         if self.payjoin_psbt.outputs.len() == 2 {
@@ -531,26 +539,22 @@ impl WantsInputs {
         }
 
         // No suitable privacy preserving selection found
-        Err(SelectionError::from(InternalSelectionError::NotFound))
+        Err(InternalSelectionError::NotFound.into())
     }
 
     fn select_first_candidate(
         &self,
         candidate_inputs: HashMap<Amount, OutPoint>,
     ) -> Result<OutPoint, SelectionError> {
-        candidate_inputs
-            .values()
-            .next()
-            .cloned()
-            .ok_or(SelectionError::from(InternalSelectionError::NotFound))
+        candidate_inputs.values().next().cloned().ok_or(InternalSelectionError::NotFound.into())
     }
 
     /// Add the provided list of inputs to the transaction.
     /// Any excess input amount is added to the change_vout output indicated previously.
     pub fn contribute_witness_inputs(
         self,
         inputs: impl IntoIterator<Item = (OutPoint, TxOut)>,
-    ) -> WantsInputs {
+    ) -> Result<WantsInputs, InputContributionError> {
         let mut payjoin_psbt = self.payjoin_psbt.clone();
         // The payjoin proposal must not introduce mixed input sequence numbers
         let original_sequence = self
@@ -587,15 +591,15 @@ impl WantsInputs {
             let change_amount = receiver_input_amount - receiver_min_input_amount;
             payjoin_psbt.unsigned_tx.output[self.change_vout].value += change_amount;
         } else {
-            todo!("Input amount is not enough to cover additional output value");
+            return Err(InternalInputContributionError::ValueTooLow.into());
         }
 
-        WantsInputs {
+        Ok(WantsInputs {
             original_psbt: self.original_psbt,
             payjoin_psbt,
             params: self.params,
             change_vout: self.change_vout,
-        }
+        })
     }
 
     // Compute the minimum amount that the receiver must contribute to the transaction as input
diff --git a/payjoin/src/receive/v2/mod.rs b/payjoin/src/receive/v2/mod.rs
@@ -15,7 +15,10 @@ use serde::{Deserialize, Serialize, Serializer};
 use url::Url;
 
 use super::v2::error::{InternalSessionError, SessionError};
-use super::{Error, InternalRequestError, RequestError, SelectionError};
+use super::{
+    Error, InputContributionError, InternalRequestError, OutputSubstitutionError, RequestError,
+    SelectionError,
+};
 use crate::psbt::PsbtExt;
 use crate::receive::optional_parameters::Params;
 use crate::v2::OhttpEncapsulationError;
@@ -402,7 +405,10 @@ impl WantsOutputs {
     }
 
     /// Substitute the receiver output script with the provided script.
-    pub fn substitute_receiver_script(self, output_script: &Script) -> Result<WantsOutputs, Error> {
+    pub fn substitute_receiver_script(
+        self,
+        output_script: &Script,
+    ) -> Result<WantsOutputs, OutputSubstitutionError> {
         let inner = self.inner.substitute_receiver_script(output_script)?;
         Ok(WantsOutputs { inner, context: self.context })
     }
@@ -416,7 +422,7 @@ impl WantsOutputs {
         self,
         replacement_outputs: Vec<TxOut>,
         drain_script: &Script,
-    ) -> Result<WantsOutputs, Error> {
+    ) -> Result<WantsOutputs, OutputSubstitutionError> {
         let inner = self.inner.replace_receiver_outputs(replacement_outputs, drain_script)?;
         Ok(WantsOutputs { inner, context: self.context })
     }
@@ -460,9 +466,9 @@ impl WantsInputs {
     pub fn contribute_witness_inputs(
         self,
         inputs: impl IntoIterator<Item = (OutPoint, TxOut)>,
-    ) -> WantsInputs {
-        let inner = self.inner.contribute_witness_inputs(inputs);
-        WantsInputs { inner, context: self.context }
+    ) -> Result<WantsInputs, InputContributionError> {
+        let inner = self.inner.contribute_witness_inputs(inputs)?;
+        Ok(WantsInputs { inner, context: self.context })
     }
 
     /// Proceed to the proposal finalization step.
diff --git a/payjoin/tests/integration.rs b/payjoin/tests/integration.rs
@@ -644,6 +644,7 @@ mod integration {
 
             let payjoin = payjoin
                 .contribute_witness_inputs(vec![(selected_outpoint, txo_to_contribute)])
+                .unwrap()
                 .commit_inputs();
 
             // Sign and finalize the proposal PSBT
@@ -1103,7 +1104,7 @@ mod integration {
             }
         };
 
-        let payjoin = payjoin.contribute_witness_inputs(inputs).commit_inputs();
+        let payjoin = payjoin.contribute_witness_inputs(inputs).unwrap().commit_inputs();
 
         let payjoin_proposal = payjoin
             .finalize_proposal(

Original file line number	Diff line number	Diff line change
`@@ -377,6 +377,7 @@ fn try_contributing_inputs(`
`377`	`377`
`378`	`378`	`Ok(payjoin`
`379`	`379`	`.contribute_witness_inputs(vec![(selected_outpoint, txo_to_contribute)])`
	`380`	`+ .expect("This shouldn't happen. Failed to contribute inputs.")`
`380`	`381`	`.commit_inputs())`
`381`	`382`	`}`
`382`	`383`