[Temporary] Remove Vec from MaybeInputsOwned non-blocking interface

Author: xstoicunicornx

payjoin/src/core/receive/mod.rs

@@ -228,6 +228,67 @@ impl<'a> From<&'a InputPair> for InternalInputPair<'a> {
     fn from(pair: &'a InputPair) -> Self { Self { psbtin: &pair.psbtin, txin: &pair.txin } }
 }
 
+pub struct Validator<T> {
+    items: Vec<T>,
+    positives: Vec<T>,
+    negatives: Vec<T>,
+    finalized: bool,
+}
+
+impl<T> Validator<T>
+where
+    T: Eq + Clone,
+{
+    fn new(items: Vec<T>) -> Self {
+        Validator { items, positives: vec![], negatives: vec![], finalized: false }
+    }
+
+    pub fn is_finalized(&self) -> bool { self.finalized }
+
+    pub fn get_reference(&self) -> UntaggedReference<T> { UntaggedReference(self.items[0].clone()) }
+
+    pub fn mark_reference(&mut self, tagged_reference: TaggedReference<T>) -> Result<bool, Error> {
+        if self.finalized {
+            return Err(ImplementationError::from(
+                "Validation state machine already validated, no references left to mark",
+            )
+            .into());
+        } else if self.items[0] != tagged_reference.0 {
+            return Err(ImplementationError::from(
+                "Incorrect reference returned for current validation state",
+            )
+            .into());
+        }
+
+        if tagged_reference.1 {
+            self.positives.push(self.items.remove(0));
+        } else {
+            self.negatives.push(self.items.remove(0));
+        }
+        if self.items.is_empty() {
+            self.finalized = true
+        }
+        Ok(self.finalized)
+    }
+
+    fn get_positives(&self) -> Vec<T> { self.positives.clone() }
+}
+
+pub struct UntaggedReference<T>(T);
+
+pub struct TaggedReference<T>(T, bool);
+
+impl<T> UntaggedReference<T>
+where
+    T: Clone,
+{
+    pub fn value(&self) -> T { self.0.clone() }
+
+    pub fn mark(&self, is_owned_seen: bool) -> TaggedReference<T> {
+        TaggedReference(self.0.clone(), is_owned_seen)
+    }
+}
+
 /// Validate the payload of a Payjoin request for PSBT and Params sanity
 pub(crate) fn parse_payload(
     base64: &str,
@@ -424,6 +485,41 @@ impl OriginalPayload {
         Ok(scripts)
     }
 
+    /// Utility function to check inputs owned using Validator
+    fn validate_inputs_not_owned(
+        &self,
+        is_owned: &mut impl FnMut(&Script) -> Result<bool, ImplementationError>,
+    ) -> Result<Validator<ScriptBuf>, Error> {
+        let input_scripts = self.extract_input_scripts()?;
+        let mut validator = Validator::new(input_scripts);
+        while !validator.is_finalized() {
+            let untagged_reference = validator.get_reference();
+            let input_script = untagged_reference.value();
+            let input_script_owned = is_owned(&input_script)?;
+            let tagged_reference = untagged_reference.mark(input_script_owned);
+            validator.mark_reference(tagged_reference)?;
+        }
+        Ok(validator)
+    }
+
+    /// Process inputs owned validator
+    ///
+    /// An attacker can try to spend the receiver's own inputs. This check prevents that.
+    #[warn(private_interfaces)]
+    pub fn process_inputs_owned_validator(
+        &self,
+        validator: Validator<ScriptBuf>,
+    ) -> Result<(), Error> {
+        if !validator.is_finalized() {
+            return Err(ImplementationError::from("Validator has not finished validation").into());
+        }
+        let owned_inputs = validator.get_positives();
+        if !owned_inputs.is_empty() {
+            return Err(InternalPayloadError::InputOwned(owned_inputs[0].clone()).into());
+        }
+        Ok(())
+    }
+
     /// Utility function to run callback to check if inputs are owned and gather the result
     pub fn gather_inputs_owned_result(
         &self,

payjoin/src/core/receive/v1/mod.rs

@@ -153,8 +153,9 @@ impl UncheckedOriginalPayload {
 /// to extract the signed original PSBT to schedule a fallback in case the Payjoin process fails.
 ///
 /// Call [`Self::check_inputs_not_owned`] to proceed. If caller needs to use non-blocking calls
-/// for checking inputs are not owned call [`Self::extract_input_scripts`] to extract the scripts
-/// to be checked and [`Self::process_inputs_owned_result`] to return the result and proceed to the
+/// for checking inputs are not owned call [`Self::get_inputs_owned_validator`] to get
+/// a Validator instance that will provide the input scripts to be checked
+/// and [`Self::process_inputs_owned_validator`] to return a finalized Validator and proceed to the
 /// next state.
 #[derive(Debug, Clone)]
 pub struct MaybeInputsOwned {
@@ -178,25 +179,33 @@ impl MaybeInputsOwned {
         self,
         is_owned: &mut impl FnMut(&Script) -> Result<bool, ImplementationError>,
     ) -> Result<MaybeInputsSeen, Error> {
-        let inputs_owned_result = self.original.gather_inputs_owned_result(is_owned)?;
-        self.process_inputs_owned_result(inputs_owned_result)
+        let validator = self.original.validate_inputs_not_owned(is_owned)?;
+        self.process_inputs_owned_validator(validator)
     }
 
-    /// Extracts the inputs txout script pubkeys
+    /// Utility function for extracting the inputs txout script pubkeys
     ///
     /// Use this for using non-blocking calls to check whether inputs are owned
     pub fn extract_input_scripts(&self) -> Result<Vec<ScriptBuf>, Error> {
         self.original.extract_input_scripts()
     }
 
-    /// Process result of whether the original PSBT has no receiver-owned inputs.
+    /// Provides Validator for checking whether input is owned
     ///
     /// Use this for using non-blocking calls to check whether inputs are owned
-    pub fn process_inputs_owned_result(
+    pub fn get_inputs_owned_validator(&self) -> Result<Validator<ScriptBuf>, Error> {
+        let input_scripts = self.extract_input_scripts()?;
+        Ok(Validator::new(input_scripts))
+    }
+
+    /// Provides Validator for checking whether input is owned
+    ///
+    /// Use this for using non-blocking calls to check whether inputs are owned
+    pub fn process_inputs_owned_validator(
         self,
-        inputs_owned_result: Vec<(ScriptBuf, bool)>,
+        validator: Validator<ScriptBuf>,
     ) -> Result<MaybeInputsSeen, Error> {
-        self.original.process_inputs_owned_result(inputs_owned_result)?;
+        self.original.process_inputs_owned_validator(validator)?;
         Ok(MaybeInputsSeen { original: self.original })
     }
 }

payjoin/src/core/receive/v2/mod.rs

@@ -57,7 +57,7 @@ use crate::persist::{
     MaybeFatalOrSuccessTransition, MaybeFatalTransition, MaybeFatalTransitionWithNoResults,
     MaybeSuccessTransition, MaybeTransientTransition, NextStateTransition,
 };
-use crate::receive::{parse_payload, InputPair, OriginalPayload, PsbtContext};
+use crate::receive::{parse_payload, InputPair, OriginalPayload, PsbtContext, Validator};
 use crate::time::Time;
 use crate::uri::ShortId;
 use crate::{ImplementationError, IntoUrl, IntoUrlError, Request, Version};
@@ -665,9 +665,10 @@ pub struct MaybeInputsOwned {
 /// to extract the signed original PSBT to schedule a fallback in case the Payjoin process fails.
 ///
 /// Call [`Receiver<MaybeInputsOwned>::check_inputs_not_owned`] to proceed. If caller needs to use
-/// non-blocking calls for checking inputs are not owned call [`Receiver<MaybeInputsOwned>::extract_input_scripts`]
-/// to extract the scripts to be checked and [`Receiver<MaybeInputsOwned>::process_inputs_owned_result`]
-/// to return the result and proceed to the next state.
+/// non-blocking calls for checking inputs are not owned call
+/// [`Receiver<MaybeInputsOwned>::get_inputs_owned_validator`] to get a Validator instance that
+/// will provide the input scripts to be checked and [`Receiver<MaybeInputsOwned>::process_inputs_owned_validator`]
+/// to return a finalized Validator and proceed to the next state.
 impl Receiver<MaybeInputsOwned> {
     /// Extracts the original transaction received from the sender.
     ///
@@ -690,9 +691,9 @@ impl Receiver<MaybeInputsOwned> {
         Error,
         Receiver<HasReplyableError>,
     > {
-        let inputs_owned_result = self.original.gather_inputs_owned_result(is_owned);
-        match inputs_owned_result {
-            Ok(inputs_owned_result) => self.process_inputs_owned_result(inputs_owned_result),
+        let validator_result = self.original.validate_inputs_not_owned(is_owned);
+        match validator_result {
+            Ok(validator) => self.process_inputs_owned_validator(validator),
             Err(e) => match e {
                 Error::Implementation(_) => MaybeFatalTransition::transient(e),
                 _ => MaybeFatalTransition::replyable_error(
@@ -707,26 +708,34 @@ impl Receiver<MaybeInputsOwned> {
         }
     }
 
-    /// Extracts the inputs txout script pubkeys
+    /// Utility function for extracting the inputs txout script pubkeys
     ///
     /// Use this for using non-blocking calls to check whether inputs are owned
     pub fn extract_input_scripts(&self) -> Result<Vec<ScriptBuf>, Error> {
         self.original.extract_input_scripts()
     }
 
-    /// Process result of whether the original PSBT inputs are owned.
+    /// Provides Validator for checking whether input is owned
     ///
     /// Use this for using non-blocking calls to check whether inputs are owned
-    pub fn process_inputs_owned_result(
+    pub fn get_inputs_owned_validator(&self) -> Result<Validator<ScriptBuf>, Error> {
+        let input_scripts = self.extract_input_scripts()?;
+        Ok(Validator::new(input_scripts))
+    }
+
+    /// Process Validator for whether the original PSBT inputs are owned.
+    ///
+    /// Use this for using non-blocking calls to check whether inputs are owned
+    pub fn process_inputs_owned_validator(
         self,
-        inputs_owned_result: Vec<(ScriptBuf, bool)>,
+        validator: Validator<ScriptBuf>,
     ) -> MaybeFatalTransition<
         SessionEvent,
         Receiver<MaybeInputsSeen>,
         Error,
         Receiver<HasReplyableError>,
     > {
-        match self.state.original.process_inputs_owned_result(inputs_owned_result) {
+        match self.state.original.process_inputs_owned_validator(validator) {
             Ok(inner) => inner,
             Err(e) => match e {
                 Error::Implementation(_) => {