1use crate::storage::OpenMlsProvider;
30use crate::test_utils::OpenMlsRustCrypto;
31use crate::treesync::LeafNodeParameters;
32use crate::{
33 binary_tree::array_representation::LeafNodeIndex,
34 ciphersuite::{hash_ref::KeyPackageRef, *},
35 credentials::*,
36 framing::*,
37 group::*,
38 key_packages::*,
39 messages::*,
40 treesync::{node::Node, LeafNode, RatchetTree, RatchetTreeIn},
41};
42use openmls_basic_credential::SignatureKeyPair;
43use openmls_traits::{
44 crypto::OpenMlsCrypto,
45 types::{Ciphersuite, HpkeKeyPair, SignatureScheme},
46 OpenMlsProvider as _,
47};
48use rand::Rng;
49
50use std::{collections::HashMap, sync::RwLock};
51use tls_codec::{Deserialize, Serialize};
52
53pub mod client;
54pub mod errors;
55
56use self::client::*;
57use self::errors::*;
58
59#[derive(Clone)]
60pub struct Group {
66 pub group_id: GroupId,
67 pub members: Vec<(usize, Vec<u8>)>,
68 pub ciphersuite: Ciphersuite,
69 pub group_config: MlsGroupJoinConfig,
70 pub public_tree: RatchetTree,
71 pub exporter_secret: Vec<u8>,
72}
73
74impl Group {
75 pub fn random_group_member(&self) -> (u32, Vec<u8>) {
77 let index = (rand::rng().next_u32() as usize) % self.members.len();
78 let (i, identity) = self.members[index].clone();
79 (i as u32, identity)
80 }
81 pub fn group_id(&self) -> &GroupId {
82 &self.group_id
83 }
84
85 pub fn members(&self) -> impl Iterator<Item = (u32, Vec<u8>)> + '_ {
86 self.members
87 .clone()
88 .into_iter()
89 .map(|(index, id)| (index as u32, id))
90 }
91}
92
93#[derive(Debug)]
94pub enum ActionType {
95 Commit,
96 Proposal,
97}
98
99#[derive(Debug, PartialEq, Eq)]
100pub enum CodecUse {
101 SerializedMessages,
102 StructMessages,
103}
104
105pub struct MlsGroupTestSetup<Provider: OpenMlsProvider> {
112 pub clients: RwLock<HashMap<Vec<u8>, RwLock<Client<Provider>>>>,
114 pub groups: RwLock<HashMap<GroupId, Group>>,
115 pub waiting_for_welcome: RwLock<HashMap<Vec<u8>, Vec<u8>>>,
117 pub default_mgp: MlsGroupCreateConfig,
118 pub use_codec: CodecUse,
121}
122
123#[allow(clippy::await_holding_lock)]
142impl<Provider: OpenMlsProvider + Default> MlsGroupTestSetup<Provider> {
143 #[openmls_traits::maybe_async]
148 pub async fn new(
149 default_mgp: MlsGroupCreateConfig,
150 number_of_clients: usize,
151 use_codec: CodecUse,
152 ) -> Self {
153 let mut clients = HashMap::new();
154 for i in 0..number_of_clients {
155 let identity = i.to_be_bytes().to_vec();
156 let provider = Provider::default();
157 let mut credentials = HashMap::new();
158 for ciphersuite in provider.crypto().supported_ciphersuites().iter() {
159 let credential = BasicCredential::new(identity.clone());
160 let signature_keys =
161 SignatureKeyPair::new(ciphersuite.signature_algorithm()).unwrap();
162 signature_keys.store(provider.storage()).await.unwrap();
163 let signature_key = OpenMlsSignaturePublicKey::new(
164 signature_keys.public().into(),
165 signature_keys.signature_scheme(),
166 )
167 .unwrap();
168
169 credentials.insert(
170 *ciphersuite,
171 CredentialWithKey {
172 credential: credential.into(),
173 signature_key: signature_key.into(),
174 },
175 );
176 }
177 let client = Client {
178 identity: identity.clone(),
179 credentials,
180 provider,
181 groups: RwLock::new(HashMap::new()),
182 };
183 clients.insert(identity, RwLock::new(client));
184 }
185 let groups = RwLock::new(HashMap::new());
186 let waiting_for_welcome = RwLock::new(HashMap::new());
187 MlsGroupTestSetup {
188 clients: RwLock::new(clients),
189 groups,
190 waiting_for_welcome,
191 default_mgp,
192 use_codec,
193 }
194 }
195
196 #[openmls_traits::maybe_async]
201 pub async fn get_fresh_key_package(
202 &self,
203 client: &Client<Provider>,
204 ciphersuite: Ciphersuite,
205 ) -> Result<KeyPackage, SetupError<Provider::StorageError>> {
206 let key_package = client.get_fresh_key_package(ciphersuite).await?;
207 self.waiting_for_welcome
208 .write()
209 .expect("An unexpected error occurred.")
210 .insert(
211 key_package
212 .hash_ref(client.provider.crypto())?
213 .as_slice()
214 .to_vec(),
215 client.identity.clone(),
216 );
217 Ok(key_package)
218 }
219
220 pub fn identity_by_index(&self, index: usize, group: &Group) -> Option<Vec<u8>> {
222 let (_, id) = group
223 .members
224 .iter()
225 .find(|(leaf_index, _)| index == *leaf_index)
226 .expect("Couldn't find member at leaf index");
227 let clients = self.clients.read().expect("An unexpected error occurred.");
228 let client = clients
229 .get(id)
230 .expect("An unexpected error occurred.")
231 .read()
232 .expect("An unexpected error occurred.");
233 client.identity(&group.group_id)
234 }
235
236 pub fn identity_by_id(&self, id: &[u8], group: &Group) -> Option<Vec<u8>> {
238 let (_, id) = group
239 .members
240 .iter()
241 .find(|(_, leaf_id)| id == leaf_id)
242 .expect("Couldn't find member at leaf index");
243 let clients = self.clients.read().expect("An unexpected error occurred.");
244 let client = clients
245 .get(id)
246 .expect("An unexpected error occurred.")
247 .read()
248 .expect("An unexpected error occurred.");
249 client.identity(&group.group_id)
250 }
251
252 #[openmls_traits::maybe_async]
259 pub async fn deliver_welcome(
260 &self,
261 welcome: Welcome,
262 group: &Group,
263 ) -> Result<(), SetupError<Provider::StorageError>> {
264 let welcome = match self.use_codec {
266 CodecUse::SerializedMessages => {
267 let serialized_welcome = welcome
268 .tls_serialize_detached()
269 .map_err(ClientError::TlsCodecError)?;
270 Welcome::tls_deserialize(&mut serialized_welcome.as_slice())
271 .map_err(ClientError::TlsCodecError)?
272 }
273 CodecUse::StructMessages => welcome,
274 };
275 let clients = self.clients.read().expect("An unexpected error occurred.");
276 for egs in welcome.secrets() {
277 let client_id = self
278 .waiting_for_welcome
279 .write()
280 .expect("An unexpected error occurred.")
281 .remove(egs.new_member().as_slice())
282 .ok_or(SetupError::NoFreshKeyPackage)?;
283 let client = clients
284 .get(&client_id)
285 .expect("An unexpected error occurred.")
286 .read()
287 .expect("An unexpected error occurred.");
288 client
289 .join_group(
290 group.group_config.clone(),
291 welcome.clone(),
292 Some(group.public_tree.clone().into()),
293 )
294 .await?;
295 }
296 Ok(())
297 }
298
299 #[openmls_traits::maybe_async]
303 pub async fn distribute_to_members<AS: Fn(&Credential) -> bool>(
304 &self,
305 sender_id: &[u8],
308 group: &mut Group,
309 message: &MlsMessageIn,
310 authentication_service: &AS,
311 ) -> Result<(), ClientError<Provider::StorageError>> {
312 let message: ProtocolMessage = match self.use_codec {
314 CodecUse::SerializedMessages => {
315 let mls_message_out: MlsMessageOut = message.clone().into();
316 let serialized_message = mls_message_out
317 .tls_serialize_detached()
318 .map_err(ClientError::TlsCodecError)?;
319
320 MlsMessageIn::tls_deserialize(&mut serialized_message.as_slice())
321 .map_err(ClientError::TlsCodecError)?
322 }
323 CodecUse::StructMessages => message.clone(),
324 }
325 .into_protocol_message()
326 .expect("Unexptected message type.");
327 let clients = self.clients.read().expect("An unexpected error occurred.");
328 for (_index, member_id) in &group.members {
330 if message.content_type() == ContentType::Application && member_id == sender_id {
331 continue;
332 }
333 let member = clients
334 .get(member_id)
335 .expect("An unexpected error occurred.")
336 .read()
337 .expect("An unexpected error occurred.");
338 member
339 .receive_messages_for_group(&message, sender_id, &authentication_service)
340 .await?;
341 }
342 let sender = clients
344 .get(sender_id)
345 .expect("An unexpected error occurred.")
346 .read()
347 .expect("An unexpected error occurred.");
348 let sender_groups = sender.groups.read().expect("An unexpected error occurred.");
349 let sender_group = sender_groups
350 .get(&group.group_id)
351 .expect("An unexpected error occurred.");
352 group.members = sender
353 .get_members_of_group(&group.group_id)?
354 .iter()
355 .map(
356 |Member {
357 index, credential, ..
358 }| { (index.usize(), credential.serialized_content().to_vec()) },
359 )
360 .collect();
361 group.public_tree = sender_group.export_ratchet_tree();
362 group.exporter_secret = sender_group
363 .export_secret(sender.provider.crypto(), "test", &[], 32)
364 .map_err(ClientError::ExportSecretError)?
365 .as_slice()
366 .to_vec();
367 Ok(())
368 }
369
370 #[openmls_traits::maybe_async]
376 pub async fn check_group_states<AS: Fn(&Credential) -> bool>(
377 &self,
378 group: &mut Group,
379 authentication_service: AS,
380 ) {
381 let clients = self.clients.read().expect("An unexpected error occurred.");
382
383 let mut messages: Vec<(Vec<u8>, MlsMessageOut)> = Vec::new();
384 for (_, m_id) in &group.members {
385 let m = clients
386 .get(m_id)
387 .expect("An unexpected error occurred.")
388 .read()
389 .expect("An unexpected error occurred.");
390 let mut group_states = m.groups.write().expect("An unexpected error occurred.");
391 if let Some(group_state) = group_states.get_mut(&group.group_id) {
393 assert_eq!(group_state.export_ratchet_tree(), group.public_tree);
394 assert_eq!(
395 group_state
396 .export_secret(m.provider.crypto(), "test", &[], 32)
397 .expect("An unexpected error occurred.")
398 .as_slice(),
399 group.exporter_secret
400 );
401 let signature_pk = group_state.own_leaf().unwrap().signature_key();
404 let signer = SignatureKeyPair::read(
405 m.provider.storage(),
406 signature_pk.as_slice(),
407 group_state.ciphersuite().signature_algorithm(),
408 )
409 .await
410 .unwrap();
411 let message = group_state
412 .create_message(&m.provider, &signer, "Hello World!".as_bytes())
413 .await
414 .expect("Error composing message while checking group states.");
415 messages.push((m_id.to_vec(), message));
416 }
417 }
418 drop(clients);
419 for (sender_id, message) in messages {
420 self.distribute_to_members(&sender_id, group, &message.into(), &authentication_service)
421 .await
422 .expect("Error sending messages to clients while checking group states.");
423 }
424 }
425
426 pub fn random_new_members_for_group(
432 &self,
433 group: &Group,
434 number_of_members: usize,
435 ) -> Result<Vec<Vec<u8>>, SetupError<Provider::StorageError>> {
436 let clients = self.clients.read().expect("An unexpected error occurred.");
437 if number_of_members + group.members.len() > clients.len() {
438 return Err(SetupError::NotEnoughClients);
439 }
440 let mut new_member_ids: Vec<Vec<u8>> = Vec::new();
441
442 for _ in 0..number_of_members {
443 let is_in_new_members = |client_id| {
444 new_member_ids
445 .iter()
446 .any(|new_member_id| client_id == new_member_id)
447 };
448 let is_in_group = |client_id| {
449 group
450 .members
451 .iter()
452 .any(|(_, member_id)| client_id == member_id)
453 };
454 let new_member_id = clients
455 .keys()
456 .find(|&client_id| !is_in_group(client_id) && !is_in_new_members(client_id))
457 .expect("An unexpected error occurred.");
458 new_member_ids.push(new_member_id.clone());
459 }
460 Ok(new_member_ids)
461 }
462
463 #[openmls_traits::maybe_async]
470 pub async fn create_group(
471 &self,
472 ciphersuite: Ciphersuite,
473 ) -> Result<GroupId, SetupError<Provider::StorageError>> {
474 let clients = self.clients.read().expect("An unexpected error occurred.");
476 let group_creator_id = ((rand::rng().next_u32() as usize) % clients.len())
477 .to_be_bytes()
478 .to_vec();
479 let group_creator = clients
480 .get(&group_creator_id)
481 .expect("An unexpected error occurred.")
482 .read()
483 .expect("An unexpected error occurred.");
484 let mut groups = self.groups.write().expect("An unexpected error occurred.");
485 let group_id = group_creator
486 .create_group(self.default_mgp.clone(), ciphersuite)
487 .await?;
488 let creator_groups = group_creator
489 .groups
490 .read()
491 .expect("An unexpected error occurred.");
492 let group = creator_groups
493 .get(&group_id)
494 .expect("An unexpected error occurred.");
495 let public_tree = group.export_ratchet_tree();
496 let exporter_secret = group
497 .export_secret(group_creator.provider.crypto(), "test", &[], 32)?
498 .as_slice()
499 .to_vec();
500 let member_ids = vec![(0, group_creator_id)];
501 let group = Group {
502 group_id: group_id.clone(),
503 members: member_ids,
504 ciphersuite,
505 group_config: self.default_mgp.join_config.clone(),
506 public_tree,
507 exporter_secret,
508 };
509 groups.insert(group_id.clone(), group);
510 Ok(group_id)
511 }
512
513 #[openmls_traits::maybe_async]
515 pub async fn create_random_group<AS: Fn(&Credential) -> bool>(
516 &self,
517 target_group_size: usize,
518 ciphersuite: Ciphersuite,
519 authentication_service: AS,
520 ) -> Result<GroupId, SetupError<Provider::StorageError>> {
521 let group_id = self.create_group(ciphersuite).await?;
523
524 let mut groups = self.groups.write().expect("An unexpected error occurred.");
525 let group = groups
526 .get_mut(&group_id)
527 .expect("An unexpected error occurred.");
528
529 let mut new_members = self.random_new_members_for_group(group, target_group_size - 1)?;
531
532 while !new_members.is_empty() {
534 let adder_id = group.random_group_member();
536 let number_of_adds = ((rand::rng().next_u32() as usize) % 5 % new_members.len()) + 1;
538 let members_to_add = new_members.drain(0..number_of_adds).collect();
539 self.add_clients(
540 ActionType::Commit,
541 group,
542 &adder_id.1,
543 members_to_add,
544 &authentication_service,
545 )
546 .await?;
547 }
548 Ok(group_id)
549 }
550
551 #[openmls_traits::maybe_async]
555 pub async fn self_update<AS: Fn(&Credential) -> bool>(
556 &self,
557 action_type: ActionType,
558 group: &mut Group,
559 client_id: &[u8],
560 leaf_node_parameters: LeafNodeParameters,
561 authentication_service: &AS,
562 ) -> Result<(), SetupError<Provider::StorageError>> {
563 let clients = self.clients.read().expect("An unexpected error occurred.");
564 let client = clients
565 .get(client_id)
566 .ok_or(SetupError::UnknownClientId)?
567 .read()
568 .expect("An unexpected error occurred.");
569 let (messages, welcome_option, _) = client
570 .self_update(action_type, &group.group_id, leaf_node_parameters)
571 .await?;
572 self.distribute_to_members(
573 &client.identity,
574 group,
575 &messages.into(),
576 authentication_service,
577 )
578 .await?;
579 if let Some(welcome) = welcome_option {
580 self.deliver_welcome(welcome, group).await?;
581 }
582 Ok(())
583 }
584
585 #[openmls_traits::maybe_async]
592 pub async fn add_clients<AS: Fn(&Credential) -> bool>(
593 &self,
594 action_type: ActionType,
595 group: &mut Group,
596 adder_id: &[u8],
597 addees: Vec<Vec<u8>>,
598 authentication_service: &AS,
599 ) -> Result<(), SetupError<Provider::StorageError>> {
600 let clients = self.clients.read().expect("An unexpected error occurred.");
601 let adder = clients
602 .get(adder_id)
603 .ok_or(SetupError::UnknownClientId)?
604 .read()
605 .expect("An unexpected error occurred.");
606 if group
607 .members
608 .iter()
609 .any(|(_, id)| addees.iter().any(|client_id| client_id == id))
610 {
611 return Err(SetupError::ClientAlreadyInGroup);
612 }
613 let mut key_packages = Vec::new();
614 for addee_id in &addees {
615 let addee = clients
616 .get(addee_id)
617 .ok_or(SetupError::UnknownClientId)?
618 .read()
619 .expect("An unexpected error occurred.");
620 let key_package = self
621 .get_fresh_key_package(&addee, group.ciphersuite)
622 .await?;
623 key_packages.push(key_package);
624 }
625 let (messages, welcome_option, _) = adder
626 .add_members(action_type, &group.group_id, &key_packages)
627 .await?;
628 for message in messages {
629 self.distribute_to_members(adder_id, group, &message.into(), authentication_service)
630 .await?;
631 }
632 if let Some(welcome) = welcome_option {
633 self.deliver_welcome(welcome, group).await?;
634 }
635 Ok(())
636 }
637
638 #[openmls_traits::maybe_async]
643 pub async fn remove_clients<AS: Fn(&Credential) -> bool>(
644 &self,
645 action_type: ActionType,
646 group: &mut Group,
647 remover_id: &[u8],
648 target_members: &[LeafNodeIndex],
649 authentication_service: AS,
650 ) -> Result<(), SetupError<Provider::StorageError>> {
651 let clients = self.clients.read().expect("An unexpected error occurred.");
652 let remover = clients
653 .get(remover_id)
654 .ok_or(SetupError::UnknownClientId)?
655 .read()
656 .expect("An unexpected error occurred.");
657 let (messages, welcome_option, _) = remover
658 .remove_members(action_type, &group.group_id, target_members)
659 .await?;
660 for message in messages {
661 self.distribute_to_members(remover_id, group, &message.into(), &authentication_service)
662 .await?;
663 }
664 if let Some(welcome) = welcome_option {
665 self.deliver_welcome(welcome, group).await?;
666 }
667 Ok(())
668 }
669
670 #[openmls_traits::maybe_async]
674 pub async fn perform_random_operation<AS: Fn(&Credential) -> bool>(
675 &self,
676 group: &mut Group,
677 authentication_service: &AS,
678 ) -> Result<(), SetupError<Provider::StorageError>> {
679 let mut rng = rand::rng();
680
681 let member_id = group.random_group_member();
683 println!("Member performing the operation: {member_id:?}");
684
685 let action_type = match (rng.next_u32() as usize) % 2 {
687 0 => ActionType::Proposal,
688 1 => ActionType::Commit,
689 _ => return Err(SetupError::Unknown),
690 };
691
692 let operation_type = (rng.next_u32() as usize) % 3;
694 match operation_type {
695 0 => {
696 println!("Performing a self-update with action type: {action_type:?}");
697 self.self_update(
698 action_type,
699 group,
700 &member_id.1,
701 LeafNodeParameters::default(),
702 authentication_service,
703 )
704 .await?;
705 }
706 1 => {
707 if group.members.len() > 1 {
709 let number_of_removals =
711 (((rng.next_u32() as usize) % group.members.len()) % 5) + 1;
712
713 let (own_index, _) = group
714 .members
715 .iter()
716 .find(|(_, identity)| identity == &member_id.1)
717 .expect("An unexpected error occurred.")
718 .clone();
719 println!("Index of the member performing the {action_type:?}: {own_index:?}");
720
721 let mut target_member_leaf_indices = Vec::new();
722 let mut target_member_identities = Vec::new();
723 let clients = self.clients.read().expect("An unexpected error occurred.");
724 println!("Removing members:");
726 for _ in 0..number_of_removals {
727 let mut member_list_index = (rng.next_u32() as usize) % group.members.len();
729 let (mut leaf_index, mut identity) =
732 group.members[member_list_index].clone();
733 while leaf_index == own_index
734 || target_member_identities.contains(&identity)
735 {
736 member_list_index = (rng.next_u32() as usize) % group.members.len();
737 let (new_leaf_index, new_identity) =
738 group.members[member_list_index].clone();
739 leaf_index = new_leaf_index;
740 identity = new_identity;
741 }
742 let client = clients
743 .get(&identity)
744 .expect("An unexpected error occurred.")
745 .read()
746 .expect("An unexpected error occurred.");
747 let client_group =
748 client.groups.read().expect("An unexpected error occurred.");
749 let client_group = client_group
750 .get(&group.group_id)
751 .expect("An unexpected error occurred.");
752 target_member_leaf_indices.push(client_group.own_leaf_index());
753 target_member_identities.push(identity);
754 }
755 self.remove_clients(
756 action_type,
757 group,
758 &member_id.1,
759 &target_member_leaf_indices,
760 authentication_service,
761 )
762 .await?
763 };
764 }
765 2 => {
766 let clients_left = self
768 .clients
769 .read()
770 .expect("An unexpected error occurred.")
771 .len()
772 - group.members.len();
773 if clients_left > 0 {
774 let number_of_adds = (((rng.next_u32() as usize) % clients_left) % 5) + 1;
775 let new_member_ids = self
776 .random_new_members_for_group(group, number_of_adds)
777 .expect("An unexpected error occurred.");
778 println!("{action_type:?}: Adding new clients: {new_member_ids:?}");
779 self.add_clients(
781 action_type,
782 group,
783 &member_id.1,
784 new_member_ids,
785 authentication_service,
786 )
787 .await?;
788 }
789 }
790 _ => return Err(SetupError::Unknown),
791 };
792 Ok(())
793 }
794}
795
796pub fn noop_authentication_service(_cred: &Credential) -> bool {
800 true
801}