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