Use FnMut closures for receiver typestate checks

payjoin-cli/src/app/v1.rs

@@ -321,17 +321,18 @@ impl App {
         let _to_broadcast_in_failure_case = proposal.extract_tx_to_schedule_broadcast();
 
         // Receive Check 2: receiver can't sign for proposal inputs
-        let proposal = proposal.check_inputs_not_owned(|input| {
+        let proposal = proposal.check_inputs_not_owned(&mut |input| {
             wallet.is_mine(input).map_err(|e| ImplementationError::from(e.into_boxed_dyn_error()))
         })?;
         log::trace!("check2");
 
         // Receive Check 3: have we seen this input before? More of a check for non-interactive i.e. payment processor receivers.
-        let payjoin = proposal
-            .check_no_inputs_seen_before(|input| Ok(self.db.insert_input_seen_before(*input)?))?;
+        let payjoin = proposal.check_no_inputs_seen_before(&mut |input| {
+            Ok(self.db.insert_input_seen_before(*input)?)
+        })?;
         log::trace!("check3");
 
-        let payjoin = payjoin.identify_receiver_outputs(|output_script| {
+        let payjoin = payjoin.identify_receiver_outputs(&mut |output_script| {
             wallet
                 .is_mine(output_script)
                 .map_err(|e| ImplementationError::from(e.into_boxed_dyn_error()))

payjoin-cli/src/app/v2/mod.rs

@@ -374,7 +374,7 @@ impl App {
     ) -> Result<()> {
         let wallet = self.wallet();
         let proposal = proposal
-            .check_inputs_not_owned(|input| {
+            .check_inputs_not_owned(&mut |input| {
                 wallet
                     .is_mine(input)
                     .map_err(|e| ImplementationError::from(e.into_boxed_dyn_error()))
@@ -389,7 +389,9 @@ impl App {
         persister: &ReceiverPersister,
     ) -> Result<()> {
         let proposal = proposal
-            .check_no_inputs_seen_before(|input| Ok(self.db.insert_input_seen_before(*input)?))
+            .check_no_inputs_seen_before(&mut |input| {
+                Ok(self.db.insert_input_seen_before(*input)?)
+            })
             .save(persister)?;
         self.identify_receiver_outputs(proposal, persister).await
     }
@@ -401,7 +403,7 @@ impl App {
     ) -> Result<()> {
         let wallet = self.wallet();
         let proposal = proposal
-            .identify_receiver_outputs(|output_script| {
+            .identify_receiver_outputs(&mut |output_script| {
                 wallet
                     .is_mine(output_script)
                     .map_err(|e| ImplementationError::from(e.into_boxed_dyn_error()))

payjoin-ffi/src/receive/mod.rs

@@ -417,7 +417,7 @@ impl MaybeInputsOwned {
         is_owned: impl Fn(&Vec<u8>) -> Result<bool, ImplementationError>,
     ) -> MaybeInputsOwnedTransition {
         MaybeInputsOwnedTransition(Arc::new(RwLock::new(Some(
-            self.0.clone().check_inputs_not_owned(|input| Ok(is_owned(&input.to_bytes())?)),
+            self.0.clone().check_inputs_not_owned(&mut |input| Ok(is_owned(&input.to_bytes())?)),
         ))))
     }
 }
@@ -473,7 +473,7 @@ impl MaybeInputsSeen {
         MaybeInputsSeenTransition(Arc::new(RwLock::new(Some(
             self.0
                 .clone()
-                .check_no_inputs_seen_before(|outpoint| Ok(is_known(&(*outpoint).into())?)),
+                .check_no_inputs_seen_before(&mut |outpoint| Ok(is_known(&(*outpoint).into())?)),
         ))))
     }
 }
@@ -532,7 +532,7 @@ impl OutputsUnknown {
         OutputsUnknownTransition(Arc::new(RwLock::new(Some(
             self.0
                 .clone()
-                .identify_receiver_outputs(|input| Ok(is_receiver_output(&input.to_bytes())?)),
+                .identify_receiver_outputs(&mut |input| Ok(is_receiver_output(&input.to_bytes())?)),
         ))))
     }
 }

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

@@ -112,7 +112,7 @@ pub struct MaybeInputsOwned {
 impl MaybeInputsOwned {
     pub fn check_inputs_not_owned(
         self,
-        is_owned: impl Fn(&bitcoin::Script) -> Result<bool, ImplementationError>,
+        is_owned: &mut impl FnMut(&bitcoin::Script) -> Result<bool, ImplementationError>,
     ) -> Result<MaybeInputsSeen, Error> {
         let inner = self.v1.check_inputs_not_owned(is_owned)?;
         Ok(MaybeInputsSeen { v1: inner, contexts: self.contexts })
@@ -127,7 +127,7 @@ pub struct MaybeInputsSeen {
 impl MaybeInputsSeen {
     pub fn check_no_inputs_seen_before(
         self,
-        is_seen: impl Fn(&bitcoin::OutPoint) -> Result<bool, ImplementationError>,
+        is_seen: &mut impl FnMut(&bitcoin::OutPoint) -> Result<bool, ImplementationError>,
     ) -> Result<OutputsUnknown, Error> {
         let inner = self.v1.check_no_inputs_seen_before(is_seen)?;
         Ok(OutputsUnknown { v1: inner, contexts: self.contexts })
@@ -142,7 +142,7 @@ pub struct OutputsUnknown {
 impl OutputsUnknown {
     pub fn identify_receiver_outputs(
         self,
-        is_receiver_output: impl Fn(&bitcoin::Script) -> Result<bool, ImplementationError>,
+        is_receiver_output: &mut impl FnMut(&bitcoin::Script) -> Result<bool, ImplementationError>,
     ) -> Result<WantsOutputs, Error> {
         let inner = self.v1.identify_receiver_outputs(is_receiver_output)?;
         Ok(WantsOutputs { v1: inner, contexts: self.contexts })

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

@@ -141,8 +141,8 @@ impl UncheckedProposal {
 /// Call [`Self::check_inputs_not_owned`] to proceed.
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct MaybeInputsOwned {
-    psbt: Psbt,
-    params: Params,
+    pub(crate) psbt: Psbt,
+    pub(crate) params: Params,
 }
 
 impl MaybeInputsOwned {
@@ -160,7 +160,7 @@ impl MaybeInputsOwned {
     /// An attacker can try to spend the receiver's own inputs. This check prevents that.
     pub fn check_inputs_not_owned(
         self,
-        is_owned: impl Fn(&Script) -> Result<bool, ImplementationError>,
+        is_owned: &mut impl FnMut(&Script) -> Result<bool, ImplementationError>,
     ) -> Result<MaybeInputsSeen, ReplyableError> {
         let mut err: Result<(), ReplyableError> = Ok(());
         if let Some(e) = self
@@ -206,7 +206,7 @@ impl MaybeInputsSeen {
     ///    original proposal PSBT of the current, new payjoin.
     pub fn check_no_inputs_seen_before(
         self,
-        is_known: impl Fn(&OutPoint) -> Result<bool, ImplementationError>,
+        is_known: &mut impl FnMut(&OutPoint) -> Result<bool, ImplementationError>,
     ) -> Result<OutputsUnknown, ReplyableError> {
         self.psbt.input_pairs().try_for_each(|input| {
             match is_known(&input.txin.previous_output) {
@@ -248,7 +248,7 @@ impl OutputsUnknown {
     /// outputs.
     pub fn identify_receiver_outputs(
         self,
-        is_receiver_output: impl Fn(&Script) -> Result<bool, ImplementationError>,
+        is_receiver_output: &mut impl FnMut(&Script) -> Result<bool, ImplementationError>,
     ) -> Result<WantsOutputs, ReplyableError> {
         let owned_vouts: Vec<usize> = self
             .psbt
@@ -899,14 +899,21 @@ pub(crate) mod test {
         UncheckedProposal { psbt: PARSED_ORIGINAL_PSBT.clone(), params }
     }
 
+    pub(crate) fn maybe_inputs_owned_from_test_vector() -> MaybeInputsOwned {
+        let pairs = url::form_urlencoded::parse(QUERY_PARAMS.as_bytes());
+        let params = Params::from_query_pairs(pairs, &[Version::One])
+            .expect("Could not parse params from query pairs");
+        MaybeInputsOwned { psbt: PARSED_ORIGINAL_PSBT.clone(), params }
+    }
+
     fn wants_outputs_from_test_vector(proposal: UncheckedProposal) -> WantsOutputs {
         proposal
             .assume_interactive_receiver()
-            .check_inputs_not_owned(|_| Ok(false))
+            .check_inputs_not_owned(&mut |_| Ok(false))
             .expect("No inputs should be owned")
-            .check_no_inputs_seen_before(|_| Ok(false))
+            .check_no_inputs_seen_before(&mut |_| Ok(false))
             .expect("No inputs should be seen before")
-            .identify_receiver_outputs(|script| {
+            .identify_receiver_outputs(&mut |script| {
                 let network = Network::Bitcoin;
                 Ok(Address::from_script(script, network).unwrap()
                     == Address::from_str("3CZZi7aWFugaCdUCS15dgrUUViupmB8bVM")
@@ -925,6 +932,35 @@ pub(crate) mod test {
             .expect("Contributed inputs should allow for valid fee contributions")
     }
 
+    #[test]
+    fn test_mutable_receiver_state_closures() {
+        let mut call_count = 0;
+        let maybe_inputs_owned = maybe_inputs_owned_from_test_vector();
+
+        fn mock_callback(call_count: &mut usize, ret: bool) -> Result<bool, ImplementationError> {
+            *call_count += 1;
+            Ok(ret)
+        }
+
+        let maybe_inputs_seen = maybe_inputs_owned
+            .check_inputs_not_owned(&mut |_| mock_callback(&mut call_count, false));
+        assert_eq!(call_count, 1);
+
+        let outputs_unknown = maybe_inputs_seen
+            .map_err(|_| "Check inputs owned closure failed".to_string())
+            .expect("Next receiver state should be accessible")
+            .check_no_inputs_seen_before(&mut |_| mock_callback(&mut call_count, false));
+        assert_eq!(call_count, 2);
+
+        let _wants_outputs = outputs_unknown
+            .map_err(|_| "Check no inputs seen closure failed".to_string())
+            .expect("Next receiver state should be accessible")
+            .identify_receiver_outputs(&mut |_| mock_callback(&mut call_count, true));
+        // there are 2 receiver outputs so we should expect this callback to run twice incrementing
+        // call count twice
+        assert_eq!(call_count, 4);
+    }
+
     #[test]
     fn is_output_substitution_disabled() {
         let mut proposal = unchecked_proposal_from_test_vector();
@@ -975,11 +1011,11 @@ pub(crate) mod test {
         let proposal = unchecked_proposal_from_test_vector();
         let wants_inputs = proposal
             .assume_interactive_receiver()
-            .check_inputs_not_owned(|_| Ok(false))
+            .check_inputs_not_owned(&mut |_| Ok(false))
             .expect("No inputs should be owned")
-            .check_no_inputs_seen_before(|_| Ok(false))
+            .check_no_inputs_seen_before(&mut |_| Ok(false))
             .expect("No inputs should be seen before")
-            .identify_receiver_outputs(|script| {
+            .identify_receiver_outputs(&mut |script| {
                 let network = Network::Bitcoin;
                 let target_address = Address::from_str("3CZZi7aWFugaCdUCS15dgrUUViupmB8bVM")
                     .map_err(ImplementationError::new)?

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

@@ -568,7 +568,7 @@ impl Receiver<MaybeInputsOwned> {
     /// An attacker can try to spend the receiver's own inputs. This check prevents that.
     pub fn check_inputs_not_owned(
         self,
-        is_owned: impl Fn(&Script) -> Result<bool, ImplementationError>,
+        is_owned: &mut impl FnMut(&Script) -> Result<bool, ImplementationError>,
     ) -> MaybeFatalTransition<SessionEvent, Receiver<MaybeInputsSeen>, ReplyableError> {
         let inner = match self.state.v1.clone().check_inputs_not_owned(is_owned) {
             Ok(inner) => inner,
@@ -618,7 +618,7 @@ impl Receiver<MaybeInputsSeen> {
     ///    original proposal PSBT of the current, new payjoin.
     pub fn check_no_inputs_seen_before(
         self,
-        is_known: impl Fn(&OutPoint) -> Result<bool, ImplementationError>,
+        is_known: &mut impl FnMut(&OutPoint) -> Result<bool, ImplementationError>,
     ) -> MaybeFatalTransition<SessionEvent, Receiver<OutputsUnknown>, ReplyableError> {
         let inner = match self.state.v1.clone().check_no_inputs_seen_before(is_known) {
             Ok(inner) => inner,
@@ -673,7 +673,7 @@ impl Receiver<OutputsUnknown> {
     /// outputs.
     pub fn identify_receiver_outputs(
         self,
-        is_receiver_output: impl Fn(&Script) -> Result<bool, ImplementationError>,
+        is_receiver_output: &mut impl FnMut(&Script) -> Result<bool, ImplementationError>,
     ) -> MaybeFatalTransition<SessionEvent, Receiver<WantsOutputs>, ReplyableError> {
         let inner = match self.state.inner.clone().identify_receiver_outputs(is_receiver_output) {
             Ok(inner) => inner,
@@ -1099,6 +1099,50 @@ pub mod test {
         }
     }
 
+    pub(crate) fn maybe_inputs_owned_v2_from_test_vector() -> MaybeInputsOwned {
+        let pairs = url::form_urlencoded::parse(QUERY_PARAMS.as_bytes());
+        let params = Params::from_query_pairs(pairs, &[Version::Two])
+            .expect("Test utils query params should not fail");
+        MaybeInputsOwned {
+            v1: v1::MaybeInputsOwned { psbt: PARSED_ORIGINAL_PSBT.clone(), params },
+            context: SHARED_CONTEXT.clone(),
+        }
+    }
+
+    #[test]
+    fn test_v2_mutable_receiver_state_closures() {
+        let mut call_count = 0;
+        let maybe_inputs_owned = maybe_inputs_owned_v2_from_test_vector();
+        let receiver = v2::Receiver { state: maybe_inputs_owned };
+
+        fn mock_callback(call_count: &mut usize, ret: bool) -> Result<bool, ImplementationError> {
+            *call_count += 1;
+            Ok(ret)
+        }
+
+        let maybe_inputs_seen =
+            receiver.check_inputs_not_owned(&mut |_| mock_callback(&mut call_count, false));
+        assert_eq!(call_count, 1);
+
+        let outputs_unknown = maybe_inputs_seen
+            .0
+            .map_err(|_| "Check inputs owned closure failed".to_string())
+            .expect("Next receiver state should be accessible")
+            .1
+            .check_no_inputs_seen_before(&mut |_| mock_callback(&mut call_count, false));
+        assert_eq!(call_count, 2);
+
+        let _wants_outputs = outputs_unknown
+            .0
+            .map_err(|_| "Check no inputs seen closure failed".to_string())
+            .expect("Next receiver state should be accessible")
+            .1
+            .identify_receiver_outputs(&mut |_| mock_callback(&mut call_count, true));
+        // there are 2 receiver outputs so we should expect this callback to run twice incrementing
+        // call count twice
+        assert_eq!(call_count, 4);
+    }
+
     #[test]
     fn test_unchecked_proposal_transient_error() -> Result<(), BoxError> {
         let unchecked_proposal = unchecked_proposal_v2_from_test_vector();
@@ -1127,7 +1171,7 @@ pub mod test {
         let receiver = v2::Receiver { state: unchecked_proposal };
 
         let maybe_inputs_owned = receiver.assume_interactive_receiver();
-        let maybe_inputs_seen = maybe_inputs_owned.0 .1.check_inputs_not_owned(|_| {
+        let maybe_inputs_seen = maybe_inputs_owned.0 .1.check_inputs_not_owned(&mut |_| {
             Err(ImplementationError::new(ReplyableError::Implementation("mock error".into())))
         });
 
@@ -1150,9 +1194,9 @@ pub mod test {
         let receiver = v2::Receiver { state: unchecked_proposal };
 
         let maybe_inputs_owned = receiver.assume_interactive_receiver();
-        let maybe_inputs_seen = maybe_inputs_owned.0 .1.check_inputs_not_owned(|_| Ok(false));
+        let maybe_inputs_seen = maybe_inputs_owned.0 .1.check_inputs_not_owned(&mut |_| Ok(false));
         let outputs_unknown = match maybe_inputs_seen.0 {
-            Ok(state) => state.1.check_no_inputs_seen_before(|_| {
+            Ok(state) => state.1.check_no_inputs_seen_before(&mut |_| {
                 Err(ImplementationError::new(ReplyableError::Implementation("mock error".into())))
             }),
             Err(_) => panic!("Expected Ok, got Err"),
@@ -1177,13 +1221,13 @@ pub mod test {
         let receiver = v2::Receiver { state: unchecked_proposal };
 
         let maybe_inputs_owned = receiver.assume_interactive_receiver();
-        let maybe_inputs_seen = maybe_inputs_owned.0 .1.check_inputs_not_owned(|_| Ok(false));
+        let maybe_inputs_seen = maybe_inputs_owned.0 .1.check_inputs_not_owned(&mut |_| Ok(false));
         let outputs_unknown = match maybe_inputs_seen.0 {
-            Ok(state) => state.1.check_no_inputs_seen_before(|_| Ok(false)),
+            Ok(state) => state.1.check_no_inputs_seen_before(&mut |_| Ok(false)),
             Err(_) => panic!("Expected Ok, got Err"),
         };
         let wants_outputs = match outputs_unknown.0 {
-            Ok(state) => state.1.identify_receiver_outputs(|_| {
+            Ok(state) => state.1.identify_receiver_outputs(&mut |_| {
                 Err(ImplementationError::new(ReplyableError::Implementation("mock error".into())))
             }),
             Err(_) => panic!("Expected Ok, got Err"),

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

@@ -178,15 +178,15 @@ mod tests {
         let maybe_inputs_owned = unchecked_proposal.clone().assume_interactive_receiver();
         let maybe_inputs_seen = maybe_inputs_owned
             .clone()
-            .check_inputs_not_owned(|_| Ok(false))
+            .check_inputs_not_owned(&mut |_| Ok(false))
             .expect("No inputs should be owned");
         let outputs_unknown = maybe_inputs_seen
             .clone()
-            .check_no_inputs_seen_before(|_| Ok(false))
+            .check_no_inputs_seen_before(&mut |_| Ok(false))
             .expect("No inputs should be seen before");
         let wants_outputs = outputs_unknown
             .clone()
-            .identify_receiver_outputs(|_| Ok(true))
+            .identify_receiver_outputs(&mut |_| Ok(true))
             .expect("Outputs should be identified");
         let wants_inputs = wants_outputs.clone().commit_outputs();
         let wants_fee_range = wants_inputs.clone().commit_inputs();
@@ -375,15 +375,15 @@ mod tests {
         let maybe_inputs_owned = unchecked_proposal.clone().assume_interactive_receiver();
         let maybe_inputs_seen = maybe_inputs_owned
             .clone()
-            .check_inputs_not_owned(|_| Ok(false))
+            .check_inputs_not_owned(&mut |_| Ok(false))
             .expect("No inputs should be owned");
         let outputs_unknown = maybe_inputs_seen
             .clone()
-            .check_no_inputs_seen_before(|_| Ok(false))
+            .check_no_inputs_seen_before(&mut |_| Ok(false))
             .expect("No inputs should be seen before");
         let wants_outputs = outputs_unknown
             .clone()
-            .identify_receiver_outputs(|_| Ok(true))
+            .identify_receiver_outputs(&mut |_| Ok(true))
             .expect("Outputs should be identified");
         let wants_inputs = wants_outputs.clone().commit_outputs();
         let wants_fee_range = wants_inputs.clone().commit_inputs();
@@ -425,15 +425,15 @@ mod tests {
         let maybe_inputs_owned = unchecked_proposal.clone().assume_interactive_receiver();
         let maybe_inputs_seen = maybe_inputs_owned
             .clone()
-            .check_inputs_not_owned(|_| Ok(false))
+            .check_inputs_not_owned(&mut |_| Ok(false))
             .expect("No inputs should be owned");
         let outputs_unknown = maybe_inputs_seen
             .clone()
-            .check_no_inputs_seen_before(|_| Ok(false))
+            .check_no_inputs_seen_before(&mut |_| Ok(false))
             .expect("No inputs should be seen before");
         let wants_outputs = outputs_unknown
             .clone()
-            .identify_receiver_outputs(|_| Ok(true))
+            .identify_receiver_outputs(&mut |_| Ok(true))
             .expect("Outputs should be identified");
         let wants_inputs = wants_outputs.clone().commit_outputs();
         let wants_fee_range = wants_inputs.clone().commit_inputs();