Modify state storage to handle multiple members in the same group

mls-rs-uniffi/src/lib.rs

@@ -586,9 +586,9 @@ async fn signing_identity_to_identifier(
 impl Group {
     /// Write the current state of the group to storage defined by
     /// [`ClientConfig::group_state_storage`]
-    pub async fn write_to_storage(&self) -> Result<(), Error> {
+    pub async fn write_to_storage(&self) -> Result<Vec<u8>, Error> {
         let mut group = self.inner().await;
-        group.write_to_storage().await.map_err(Into::into)
+        Ok(group.write_to_storage().await?.to_vec())
     }
 
     /// Perform a commit of received proposals (or an empty commit).

mls-rs/src/client.rs

@@ -606,11 +606,11 @@ where
     /// this client was configured to use.
     #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
     #[inline(never)]
-    pub async fn load_group(&self, group_id: &[u8]) -> Result<Group<C>, MlsError> {
+    pub async fn load_group(&self, group_state_id: &[u8]) -> Result<Group<C>, MlsError> {
         let snapshot = self
             .config
             .group_state_storage()
-            .state(group_id)
+            .state(group_state_id)
             .await
             .map_err(|e| MlsError::GroupStorageError(e.into_any_error()))?
             .ok_or(MlsError::GroupNotFound)?;

mls-rs/src/group/external_commit.rs

@@ -259,6 +259,13 @@ impl<C: ClientConfig> ExternalCommitBuilder<C> {
             )
             .await?;
 
+        #[cfg(feature = "prior_epoch")]
+        {
+            let repo = &mut group.state_repo;
+            let id = &group.state.context.group_id;
+            repo.set_group_state_id(id, group.private_tree.self_index)?;
+        }
+
         group.apply_pending_commit().await?;
 
         Ok((group, commit_output.commit_message))

mls-rs/src/group/mod.rs

@@ -341,12 +341,13 @@ where
 
         let state_repo = GroupStateRepository::new(
             #[cfg(feature = "prior_epoch")]
-            context.group_id.clone(),
+            &context.group_id,
+            #[cfg(feature = "prior_epoch")]
+            private_tree.self_index,
             config.group_state_storage(),
             config.key_package_repo(),
             None,
-        )
-        .await?;
+        )?;
 
         let key_schedule_result = KeySchedule::from_random_epoch_secret(
             &cipher_suite_provider,
@@ -604,12 +605,13 @@ where
 
         let state_repo = GroupStateRepository::new(
             #[cfg(feature = "prior_epoch")]
-            group_info.group_context.group_id.clone(),
+            &group_info.group_context.group_id,
+            #[cfg(feature = "prior_epoch")]
+            private_tree.self_index,
             config.group_state_storage(),
             config.key_package_repo(),
             used_key_package_ref,
-        )
-        .await?;
+        )?;
 
         let group = Group {
             config,
@@ -1838,8 +1840,8 @@ pub(crate) mod test_utils;
 mod tests {
     use crate::{
         client::test_utils::{
-            test_client_with_key_pkg, TestClientBuilder, TEST_CIPHER_SUITE,
-            TEST_CUSTOM_PROPOSAL_TYPE, TEST_PROTOCOL_VERSION,
+            test_client_with_key_pkg, test_client_with_key_pkg_custom, TestClientBuilder,
+            TEST_CIPHER_SUITE, TEST_CUSTOM_PROPOSAL_TYPE, TEST_PROTOCOL_VERSION,
         },
         client_builder::{test_utils::TestClientConfig, ClientBuilder, MlsConfig},
         crypto::test_utils::TestCryptoProvider,
@@ -3910,7 +3912,7 @@ mod tests {
             .unwrap()
             .0;
 
-        bob.write_to_storage().await.unwrap();
+        let state_id = bob.write_to_storage().await.unwrap().to_vec();
 
         // Bob reloads his group data, but with parameters that will cause his generated leaves to
         // not support the mandatory extension.
@@ -3919,7 +3921,7 @@ mod tests {
             .key_package_repo(bob.config.key_package_repo())
             .group_state_storage(bob.config.group_state_storage())
             .build()
-            .load_group(alice.group_id())
+            .load_group(&state_id)
             .await
             .unwrap();
 
@@ -4225,4 +4227,42 @@ mod tests {
 
         assert_eq!(update.committer, *group.private_tree.self_index);
     }
+
+    #[maybe_async::test(not(mls_build_async), async(mls_build_async, crate::futures_test))]
+    async fn can_have_two_members_in_one_group() {
+        let mut group1 = test_group(TEST_PROTOCOL_VERSION, TEST_CIPHER_SUITE)
+            .await
+            .group;
+
+        let (client2, _) = test_client_with_key_pkg_custom(
+            TEST_PROTOCOL_VERSION,
+            TEST_CIPHER_SUITE,
+            "arnold",
+            |c| {
+                c.0.group_state_storage = group1.config.group_state_storage();
+                c.0.key_package_repo = group1.config.key_package_repo();
+            },
+        )
+        .await;
+
+        let commit = group1
+            .commit_builder()
+            .add_member(client2.generate_key_package_message().await.unwrap())
+            .unwrap()
+            .build()
+            .await
+            .unwrap();
+
+        let (mut group2, _) = client2
+            .join_group(None, &commit.welcome_messages[0])
+            .await
+            .unwrap();
+
+        group1.apply_pending_commit().await.unwrap();
+
+        group1
+            .process_incoming_message(group2.commit(vec![]).await.unwrap().commit_message)
+            .await
+            .unwrap();
+    }
 }

mls-rs/src/group/snapshot.rs

@@ -40,7 +40,7 @@ use super::{cipher_suite_provider, epoch::EpochSecrets, state_repo::GroupStateRe
 pub(crate) struct Snapshot {
     version: u16,
     pub(crate) state: RawGroupState,
-    private_tree: TreeKemPrivate,
+    pub(crate) private_tree: TreeKemPrivate,
     epoch_secrets: EpochSecrets,
     key_schedule: KeySchedule,
     #[cfg(all(feature = "std", feature = "by_ref_proposal"))]
@@ -51,6 +51,12 @@ pub(crate) struct Snapshot {
     signer: SignatureSecretKey,
 }
 
+impl Snapshot {
+    pub(crate) fn group_state_id(&self) -> Result<Vec<u8>, mls_rs_codec::Error> {
+        (&self.state.context.group_id, self.private_tree.self_index).mls_encode_to_vec()
+    }
+}
+
 #[derive(Debug, MlsEncode, MlsDecode, MlsSize, PartialEq, Clone)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub(crate) struct RawGroupState {
@@ -150,10 +156,17 @@ where
 {
     /// Write the current state of the group to the
     /// [`GroupStorageProvider`](crate::GroupStateStorage)
-    /// that is currently in use by the group.
+    /// that is currently in use by the group. Returns an identifier
+    /// of the stored state that can be later used to load the group
+    /// using [`load_group`](crate::Client::load_group).
     #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
-    pub async fn write_to_storage(&mut self) -> Result<(), MlsError> {
-        self.state_repo.write_to_storage(self.snapshot()).await
+    pub async fn write_to_storage(&mut self) -> Result<Vec<u8>, MlsError> {
+        let snapshot = self.snapshot();
+        let state_id = snapshot.group_state_id()?;
+
+        self.state_repo.write_to_storage(snapshot).await?;
+
+        Ok(state_id)
     }
 
     pub(crate) fn snapshot(&self) -> Snapshot {
@@ -182,12 +195,13 @@ where
 
         let state_repo = GroupStateRepository::new(
             #[cfg(feature = "prior_epoch")]
-            snapshot.state.context.group_id.clone(),
+            &snapshot.state.context.group_id,
+            #[cfg(feature = "prior_epoch")]
+            snapshot.private_tree.self_index,
             config.group_state_storage(),
             config.key_package_repo(),
             None,
-        )
-        .await?;
+        )?;
 
         Ok(Group {
             config,