Merge remote-tracking branch 'upstream/franziskus/readd-api' into 10-08-temporary-readd-api

richardhuaaa · richardhuaaa · commit 0cf20a903ffc · 2025-10-08T15:31:00.000-07:00
diff --git a/openmls/src/group/mls_group/commit_builder.rs b/openmls/src/group/mls_group/commit_builder.rs
@@ -824,6 +824,36 @@ pub struct CommitMessageBundle {
     group_info: Option<GroupInfo>,
 }
 
+/// The result of a commit with an add proposal. This includes
+/// - The Commit as an [`MlsMessageOut`]
+/// - The [`Welcome`] as an [`MlsMessageOut`]
+/// - Optionally a [`GroupInfo`] as an [`MlsMessageOut`]
+pub struct WelcomeCommitMessages {
+    /// The Commit as an [`MlsMessageOut`].
+    pub commit: MlsMessageOut,
+
+    /// The [`Welcome`] as an [`MlsMessageOut`].
+    pub welcome: MlsMessageOut,
+
+    /// Optionally a [`GroupInfo`] as an [`MlsMessageOut`].
+    pub group_info: Option<MlsMessageOut>,
+}
+
+impl TryFrom<CommitMessageBundle> for WelcomeCommitMessages {
+    type Error = LibraryError;
+
+    fn try_from(value: CommitMessageBundle) -> Result<Self, Self::Error> {
+        let (commit, welcome_opt, group_info) = value.into_messages();
+        Ok(Self {
+            commit,
+            welcome: welcome_opt.ok_or(LibraryError::custom(
+                "WelcomeCommitMessages must only be used with commits that produce a welcome.",
+            ))?,
+            group_info,
+        })
+    }
+}
+
 #[cfg(test)]
 impl CommitMessageBundle {
     pub fn new(
diff --git a/openmls/src/group/mls_group/errors.rs b/openmls/src/group/mls_group/errors.rs
@@ -187,6 +187,42 @@ pub enum AddMembersError<StorageError> {
     StorageError(StorageError),
 }
 
+/// Add members error
+#[derive(Error, Debug, PartialEq, Clone)]
+pub enum ReAddMembersError<StorageError> {
+    /// Unable to map the key packages to the given leaf indices.
+    #[error("Unable to map the key packages to the given leaf indices.")]
+    InvalidInput,
+
+    /// See [`EmptyInputError`] for more details.
+    #[error(transparent)]
+    EmptyInput(#[from] EmptyInputError),
+
+    /// See [`MlsGroupStateError`] for more details.
+    #[error(transparent)]
+    GroupStateError(#[from] MlsGroupStateError),
+
+    /// See [`LibraryError`] for more details.
+    #[error(transparent)]
+    LibraryError(#[from] LibraryError),
+
+    /// The member that should be removed can not be found.
+    #[error("The member that should be removed can not be found.")]
+    UnknownMember,
+
+    /// Error writing to storage
+    #[error("Error writing to storage: {0}")]
+    StorageError(StorageError),
+
+    /// See [`CommitBuilderStageError`] for more details.
+    #[error(transparent)]
+    CommitBuilderStageError(#[from] CommitBuilderStageError<StorageError>),
+
+    /// See [`CreateCommitError`] for more details.
+    #[error(transparent)]
+    CreateCommitError(#[from] CreateCommitError),
+}
+
 /// Propose add members error
 #[derive(Error, Debug, PartialEq, Clone)]
 pub enum ProposeAddMemberError<StorageError> {
diff --git a/openmls/src/group/mls_group/membership.rs b/openmls/src/group/mls_group/membership.rs
@@ -11,8 +11,12 @@ use super::{
     *,
 };
 use crate::{
-    binary_tree::array_representation::LeafNodeIndex, key_packages::KeyPackage,
-    messages::group_info::GroupInfo, storage::OpenMlsProvider, treesync::LeafNode,
+    binary_tree::array_representation::LeafNodeIndex,
+    group::{ReAddMembersError, WelcomeCommitMessages},
+    key_packages::KeyPackage,
+    messages::group_info::GroupInfo,
+    storage::OpenMlsProvider,
+    treesync::LeafNode,
 };
 
 type UpdateResult<Provider> = Result<
@@ -84,6 +88,94 @@ impl MlsGroup {
         self.add_members_internal(provider, signer, key_packages, true)
     }
 
+    /// Re-add members.
+    ///
+    /// If there are issues with a set of members and they are suspected to be
+    /// forked, this functions allows to remove and re-add them.
+    /// This is a convenience function that first removes, and then re-adds the
+    /// list of provided users.
+    ///
+    /// Note that this functions _does not_ enforce that the removed `members`
+    /// and new members in the `key_packages` correspond.
+    ///
+    /// For the best performance it is important to sort the key packages in the
+    /// order they appear in the tree (left to right). If `sort` is set to
+    /// `true`, this function sorts the indices internally, based on the
+    /// credential. Note that this only works if the credential doesn't change
+    /// and the credential allows uniquely identifying a leaf node.
+    pub fn readd_members<Provider: OpenMlsProvider>(
+        &mut self,
+        provider: &Provider,
+        signer: &impl Signer,
+        members: &[LeafNodeIndex],
+        key_packages: &[KeyPackage],
+        sort: bool,
+    ) -> Result<WelcomeCommitMessages, ReAddMembersError<Provider::StorageError>> {
+        self.is_operational()?;
+
+        if members.is_empty() {
+            return Err(EmptyInputError::RemoveMembers.into());
+        }
+
+        if key_packages.is_empty() {
+            return Err(EmptyInputError::AddMembers.into());
+        }
+
+        let key_packages = if sort {
+            // Sort the key packages to get the order from left to right.
+            // This reduces churn by putting everyone back into their place.
+            // This is not the most efficient, but we don't expect these lists
+            // to be too long.
+            let member_credentials_opt = members
+                .iter()
+                .map(|&leaf_index| self.member(leaf_index).map(|c| (leaf_index, c)));
+            let member_credentials = member_credentials_opt.clone().flatten();
+
+            if member_credentials_opt.len() != member_credentials.clone().count() {
+                // We couldn't find an index. The remove is not valid.
+                return Err(ReAddMembersError::InvalidInput);
+            }
+
+            // Get leaf indices for all key packages.
+            let kp_indices = key_packages.iter().cloned().map(|kp| {
+                self.member_leaf_index(kp.leaf_node().credential())
+                    .map(|i| (i, kp))
+            });
+            let mut indices_flat: Vec<_> = kp_indices.clone().flatten().collect();
+
+            if kp_indices.len() != indices_flat.len() {
+                // We couldn't find a credential. The re-add is not valid.
+                return Err(ReAddMembersError::InvalidInput);
+            }
+
+            // Sort the key packages by index.
+            indices_flat.sort_by(|(ai, _), (bi, _)| ai.cmp(bi));
+            indices_flat
+                .into_iter()
+                .map(|(_, kp)| kp)
+                .collect::<Vec<_>>()
+        } else {
+            key_packages.to_vec()
+        };
+
+        let bundle = self
+            .commit_builder()
+            .propose_removals(members.iter().cloned())
+            .propose_adds(key_packages.into_iter())
+            .load_psks(provider.storage())?
+            .build(provider.rand(), provider.crypto(), signer, |_| true)?
+            .stage_commit(provider)?;
+
+        provider
+            .storage()
+            .write_group_state(self.group_id(), &self.group_state)
+            .map_err(ReAddMembersError::StorageError)?;
+
+        self.reset_aad();
+
+        Ok(bundle.try_into()?)
+    }
+
     /// Adds members to the group.
     ///
     /// New members are added by providing a `KeyPackage` for each member.
@@ -303,6 +395,14 @@ impl MlsGroup {
         self.public_group().members()
     }
 
+    /// Returns the [`LeafNodeIndex`] of a member corresponding to the given
+    /// credential. Returns `None` if the member can not be found in this group.
+    pub fn member_leaf_index(&self, credential: &Credential) -> Option<LeafNodeIndex> {
+        self.members()
+            .find(|m| &m.credential == credential)
+            .map(|m| m.index)
+    }
+
     /// Returns the [`Credential`] of a member corresponding to the given
     /// leaf index. Returns `None` if the member can not be found in this group.
     pub fn member(&self, leaf_index: LeafNodeIndex) -> Option<&Credential> {
diff --git a/openmls/tests/readd.rs b/openmls/tests/readd.rs
@@ -0,0 +1,159 @@
+use openmls::{prelude::*, test_utils::single_group_test_framework::*};
+use openmls_test::openmls_test;
+
+#[openmls_test]
+fn readd() {
+    let alice_party = CorePartyState::<Provider>::new("alice");
+    let bob_party = CorePartyState::<Provider>::new("bob");
+    let charlie_party = CorePartyState::<Provider>::new("charlie");
+    let yuk_party = CorePartyState::<Provider>::new("yuk");
+
+    let alice_pre_group = alice_party.generate_pre_group(ciphersuite);
+    let bob_pre_group = bob_party.generate_pre_group(ciphersuite);
+    let charlie_pre_group = charlie_party.generate_pre_group(ciphersuite);
+    let yuk_pre_group = yuk_party.generate_pre_group(ciphersuite);
+
+    let group_id = GroupId::from_slice(b"Test Group");
+
+    // Define the MlsGroup configuration
+    let group_config = MlsGroupCreateConfig::builder()
+        .ciphersuite(ciphersuite)
+        .use_ratchet_tree_extension(true)
+        .build();
+
+    let mut group_state =
+        GroupState::new_from_party(group_id.clone(), alice_pre_group, group_config.clone())
+            .unwrap();
+
+    // Generate KeyPackages
+    let bob_key_package = bob_pre_group.key_package_bundle.key_package().clone();
+    let charlie_key_package = charlie_pre_group.key_package_bundle.key_package().clone();
+    let yuk_key_package = yuk_pre_group.key_package_bundle.key_package().clone();
+
+    let [alice] = group_state.members_mut(&["alice"]);
+
+    // === Alice adds Bob ===
+    let (_commit, welcome, _group_info) = alice
+        .group
+        .add_members(
+            &alice_party.provider,
+            &alice.party.signer,
+            &[bob_key_package, charlie_key_package, yuk_key_package],
+        )
+        .expect("Could not add folks");
+
+    let welcome: MlsMessageIn = welcome.into();
+    let welcome = welcome.into_welcome().unwrap();
+
+    // Bob
+    let staged_join = StagedWelcome::new_from_welcome(
+        &bob_party.provider,
+        group_config.join_config(),
+        welcome.clone(),
+        None,
+    )
+    .expect("Error constructing staged join");
+    let mut bob_group = staged_join
+        .into_group(&bob_party.provider)
+        .expect("Error joining group from StagedWelcome");
+
+    // Charlie
+    let staged_join = StagedWelcome::new_from_welcome(
+        &charlie_party.provider,
+        group_config.join_config(),
+        welcome.clone(),
+        None,
+    )
+    .expect("Error constructing staged join");
+
+    let _charlie_group = staged_join
+        .into_group(&charlie_party.provider)
+        .expect("Error joining group from StagedWelcome");
+
+    // Yuk
+    let staged_join = StagedWelcome::new_from_welcome(
+        &yuk_party.provider,
+        group_config.join_config(),
+        welcome.clone(),
+        None,
+    )
+    .expect("Error constructing staged join");
+
+    let _yuk_group = staged_join
+        .into_group(&yuk_party.provider)
+        .expect("Error joining group from StagedWelcome");
+
+    // Bob re-adds alice and yuk
+    let alice_key_package = KeyPackage::builder()
+        .build(
+            ciphersuite,
+            &alice_party.provider,
+            &alice.party.signer,
+            alice.party.credential_with_key.clone(),
+        )
+        .unwrap();
+    let yuk_key_package = KeyPackage::builder()
+        .build(
+            ciphersuite,
+            &yuk_party.provider,
+            &yuk_pre_group.signer,
+            yuk_pre_group.credential_with_key.clone(),
+        )
+        .unwrap();
+
+    let commit_messages = bob_group
+        .readd_members(
+            &bob_party.provider,
+            &bob_pre_group.signer,
+            &[
+                LeafNodeIndex::new(0), /*Alice */
+                LeafNodeIndex::new(3), /*Yuk */
+            ],
+            &[
+                // This is the wrong order, but we ask the function to order.
+                yuk_key_package.key_package().clone(),
+                alice_key_package.key_package().clone(),
+            ],
+            // If we make this false (no sorting), the check for indices below will fail.
+            true,
+        )
+        .unwrap();
+    bob_group.merge_pending_commit(&bob_party.provider).unwrap();
+
+    let welcome: MlsMessageIn = commit_messages.welcome.into();
+    let welcome = welcome.into_welcome().unwrap();
+
+    // New Yuk
+    let staged_join = StagedWelcome::new_from_welcome(
+        &yuk_party.provider,
+        group_config.join_config(),
+        welcome.clone(),
+        None,
+    )
+    .expect("Error constructing staged join");
+
+    let yuk_group = staged_join
+        .into_group(&yuk_party.provider)
+        .expect("Error joining group from StagedWelcome");
+
+    // New Alice
+    let staged_join = StagedWelcome::new_from_welcome(
+        &alice_party.provider,
+        group_config.join_config(),
+        welcome.clone(),
+        None,
+    )
+    .expect("Error constructing staged join");
+
+    let alice_group = staged_join
+        .into_group(&alice_party.provider)
+        .expect("Error joining group from StagedWelcome");
+
+    // Yuk and Alice are back in.
+    assert_eq!(alice_group.confirmation_tag(), bob_group.confirmation_tag());
+    assert_eq!(yuk_group.confirmation_tag(), bob_group.confirmation_tag());
+
+    // They are in the same place
+    assert_eq!(alice_group.own_leaf_index().u32(), 0);
+    assert_eq!(yuk_group.own_leaf_index().u32(), 3);
+}
