Add RemoveProposal to the Removed case of CommitEffect

awslabs · Jul 19, 2024 · db66744 · db66744
1 parent fc5ea84
commit db66744
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 18 deletions.
diff --git a/mls-rs-uniffi/src/lib.rs b/mls-rs-uniffi/src/lib.rs
@@ -227,7 +227,7 @@ impl From<mls_rs::group::CommitEffect> for CommitEffect {
                     .map(|p| Arc::new(p.proposal.into()))
                     .collect(),
             },
-            group::CommitEffect::Removed => CommitEffect::Removed,
+            group::CommitEffect::Removed(_) => CommitEffect::Removed,
             group::CommitEffect::ReInit(_) => CommitEffect::ReInit,
         }
     }

diff --git a/mls-rs/src/client.rs b/mls-rs/src/client.rs
@@ -968,7 +968,7 @@ mod tests {
             assert_matches!(
                 message,
                 ReceivedMessage::Commit(CommitMessageDescription {
-                    effect: CommitEffect::Removed,
+                    effect: CommitEffect::Removed(_),
                     ..
                 })
             );

diff --git a/mls-rs/src/external_client/group.rs b/mls-rs/src/external_client/group.rs
@@ -21,6 +21,7 @@ use crate::{
             ApplicationMessageDescription, CommitMessageDescription, EventOrContent,
             MessageProcessor, ProposalMessageDescription, ProvisionalState,
         },
+        proposal_filter::ProposalInfo,
         snapshot::RawGroupState,
         state::GroupState,
         transcript_hash::InterimTranscriptHash,
@@ -638,8 +639,11 @@ where
         &mut self.state
     }
 
-    fn can_continue_processing(&self, _provisional_state: &ProvisionalState) -> bool {
-        true
+    fn removal_proposal(
+        &self,
+        _provisional_state: &ProvisionalState,
+    ) -> Option<ProposalInfo<RemoveProposal>> {
+        None
     }
 
     #[cfg(feature = "private_message")]

diff --git a/mls-rs/src/group/message_processor.rs b/mls-rs/src/group/message_processor.rs
@@ -14,7 +14,7 @@ use super::{
     state::GroupState,
     transcript_hash::InterimTranscriptHash,
     transcript_hashes, validate_group_info_member, GroupContext, GroupInfo, ReInitProposal,
-    Welcome,
+    RemoveProposal, Welcome,
 };
 use crate::{
     client::MlsError,
@@ -121,7 +121,7 @@ impl NewEpoch {
 #[derive(Clone, Debug, PartialEq)]
 pub enum CommitEffect {
     NewEpoch(Box<NewEpoch>),
-    Removed,
+    Removed(ProposalInfo<RemoveProposal>),
     ReInit(ProposalInfo<ReInitProposal>),
 }
 
@@ -636,12 +636,12 @@ pub(crate) trait MessageProcessor: Send + Sync {
             return Err(MlsError::CommitMissingPath);
         }
 
-        if !self.can_continue_processing(&provisional_state) {
+        if let Some(removal) = self.removal_proposal(&provisional_state) {
             return Ok(CommitMessageDescription {
                 is_external: matches!(auth_content.content.sender, Sender::NewMemberCommit),
                 authenticated_data: auth_content.content.authenticated_data,
                 committer: *sender,
-                effect: CommitEffect::Removed,
+                effect: CommitEffect::Removed(removal),
             });
         }
 
@@ -727,7 +727,11 @@ pub(crate) trait MessageProcessor: Send + Sync {
     fn identity_provider(&self) -> Self::IdentityProvider;
     fn cipher_suite_provider(&self) -> &Self::CipherSuiteProvider;
     fn psk_storage(&self) -> Self::PreSharedKeyStorage;
-    fn can_continue_processing(&self, provisional_state: &ProvisionalState) -> bool;
+
+    fn removal_proposal(
+        &self,
+        provisional_state: &ProvisionalState,
+    ) -> Option<ProposalInfo<RemoveProposal>>;
 
     #[cfg(feature = "private_message")]
     fn min_epoch_available(&self) -> Option<u64>;

diff --git a/mls-rs/src/group/mod.rs b/mls-rs/src/group/mod.rs
@@ -1827,13 +1827,19 @@ where
         &mut self.state
     }
 
-    fn can_continue_processing(&self, provisional_state: &ProvisionalState) -> bool {
-        !(provisional_state
-            .applied_proposals
-            .removals
-            .iter()
-            .any(|p| p.proposal.to_remove == self.private_tree.self_index)
-            && self.pending_commit.is_none())
+    fn removal_proposal(
+        &self,
+        provisional_state: &ProvisionalState,
+    ) -> Option<ProposalInfo<RemoveProposal>> {
+        match self.pending_commit {
+            Some(_) => None,
+            None => provisional_state
+                .applied_proposals
+                .removals
+                .iter()
+                .find(|p| p.proposal.to_remove == self.private_tree.self_index)
+                .cloned(),
+        }
     }
 
     #[cfg(feature = "private_message")]

diff --git a/mls-rs/src/group/test_utils.rs b/mls-rs/src/group/test_utils.rs
@@ -461,8 +461,11 @@ impl MessageProcessor for GroupWithoutKeySchedule {
         self.inner.psk_storage()
     }
 
-    fn can_continue_processing(&self, provisional_state: &ProvisionalState) -> bool {
-        self.inner.can_continue_processing(provisional_state)
+    fn removal_proposal(
+        &self,
+        provisional_state: &ProvisionalState,
+    ) -> Option<ProposalInfo<RemoveProposal>> {
+        self.inner.removal_proposal(provisional_state)
     }
 
     #[cfg(feature = "private_message")]