1use openmls_basic_credential::SignatureKeyPair;
2use openmls_traits::{signatures::Signer, types::SignatureScheme};
3pub use openmls_traits::{
4 storage::StorageProvider as StorageProviderTrait,
5 types::{Ciphersuite, HpkeKeyPair},
6 OpenMlsProvider,
7};
8
9pub use crate::utils::*;
10use crate::{
11 credentials::CredentialWithKey,
12 key_packages::KeyPackageBuilder,
13 prelude::{commit_builder::*, *},
14};
15
16pub use crate::test_utils::minimal_capabilities_for;
17use crate::test_utils::storage_state::GroupStorageState;
18
19mod assertions;
20
21mod errors;
22pub use errors::GroupError;
23use errors::*;
24
25use std::collections::HashMap;
26
27type Name = &'static str;
29
30#[openmls_traits::maybe_async]
33pub async fn generate_credential(
34 identity: Vec<u8>,
35 signature_algorithm: SignatureScheme,
36 provider: &impl crate::storage::OpenMlsProvider,
37) -> (CredentialWithKey, SignatureKeyPair) {
38 let credential = BasicCredential::new(identity);
39 let signature_keys = SignatureKeyPair::new(signature_algorithm).unwrap();
40 signature_keys.store(provider.storage()).await.unwrap();
41
42 (
43 CredentialWithKey {
44 credential: credential.into(),
45 signature_key: signature_keys.to_public_vec().into(),
46 },
47 signature_keys,
48 )
49}
50
51#[openmls_traits::maybe_async]
53pub(crate) async fn generate_key_package(
54 ciphersuite: Ciphersuite,
55 credential_with_key: CredentialWithKey,
56 extensions: Extensions<KeyPackage>,
57 provider: &impl crate::storage::OpenMlsProvider,
58 lifetime: impl Into<Option<Lifetime>>,
59 signer: &impl Signer,
60) -> KeyPackageBundle {
61 let mut builder = KeyPackage::builder().key_package_extensions(extensions);
62
63 if let Some(lifetime) = lifetime.into() {
64 builder = builder.key_package_lifetime(lifetime);
65 }
66
67 builder
68 .build(ciphersuite, provider, signer, credential_with_key)
69 .await
70 .unwrap()
71}
72
73pub struct CorePartyState<Provider> {
75 pub name: Name,
76 pub provider: Provider,
77}
78
79impl<Provider: Default> CorePartyState<Provider> {
80 pub fn new(name: Name) -> Self {
81 Self {
82 name,
83 provider: Provider::default(),
84 }
85 }
86}
87
88pub struct PreGroupPartyState<'a, Provider> {
90 pub credential_with_key: CredentialWithKey,
91 pub key_package_bundle: KeyPackageBundle,
93 pub signer: SignatureKeyPair,
94 pub core_state: &'a CorePartyState<Provider>,
95}
96
97pub struct PreGroupPartyStateBuilder<'a, Provider: OpenMlsProvider> {
98 ciphersuite: Ciphersuite,
99 lifetime: Option<Lifetime>,
100 key_package_extensions: Option<Extensions<KeyPackage>>,
101 leaf_node_extensions: Option<Extensions<LeafNode>>,
102 leaf_node_capabilities: Option<Capabilities>,
103 core_state: &'a CorePartyState<Provider>,
104}
105
106impl<'a, Provider: OpenMlsProvider> PreGroupPartyStateBuilder<'a, Provider> {
107 pub fn with_lifetime(mut self, lifetime: impl Into<Option<Lifetime>>) -> Self {
108 self.lifetime = lifetime.into();
109
110 self
111 }
112 pub fn with_key_package_extensions(
113 mut self,
114 extensions: impl Into<Option<Extensions<KeyPackage>>>,
115 ) -> Self {
116 self.key_package_extensions = extensions.into();
117
118 self
119 }
120 pub fn with_leaf_node_extensions(
121 mut self,
122 extensions: impl Into<Option<Extensions<LeafNode>>>,
123 ) -> Self {
124 self.leaf_node_extensions = extensions.into();
125
126 self
127 }
128 pub fn with_leaf_node_capabilities(
129 mut self,
130 capabilities: impl Into<Option<Capabilities>>,
131 ) -> Self {
132 self.leaf_node_capabilities = capabilities.into();
133
134 self
135 }
136
137 #[openmls_traits::maybe_async]
138 pub async fn build(self) -> PreGroupPartyState<'a, Provider> {
139 let (credential_with_key, signer) = generate_credential(
140 self.core_state.name.into(),
141 self.ciphersuite.signature_algorithm(),
142 &self.core_state.provider,
143 )
144 .await;
145 let mut builder = KeyPackage::builder()
146 .leaf_node_extensions(self.leaf_node_extensions.unwrap_or_default())
147 .key_package_extensions(self.key_package_extensions.unwrap_or_default())
148 .leaf_node_capabilities(
149 self.leaf_node_capabilities
150 .unwrap_or_else(|| minimal_capabilities_for(self.ciphersuite).build()),
151 );
152
153 if let Some(lifetime) = self.lifetime {
154 builder = builder.key_package_lifetime(lifetime);
155 }
156
157 let key_package_bundle = builder
158 .build(
159 self.ciphersuite,
160 &self.core_state.provider,
161 &signer,
162 credential_with_key.clone(),
163 )
164 .await
165 .unwrap();
166
167 PreGroupPartyState {
168 credential_with_key,
169 key_package_bundle,
170 signer,
171 core_state: self.core_state,
172 }
173 }
174}
175
176impl<Provider: OpenMlsProvider> CorePartyState<Provider> {
177 pub fn pre_group_builder<'a>(
179 &'a self,
180 ciphersuite: Ciphersuite,
181 ) -> PreGroupPartyStateBuilder<'a, Provider> {
182 PreGroupPartyStateBuilder {
183 ciphersuite,
184 lifetime: None,
185 key_package_extensions: None,
186 leaf_node_extensions: None,
187 leaf_node_capabilities: None,
188 core_state: self,
189 }
190 }
191
192 #[openmls_traits::maybe_async]
194 pub async fn generate_pre_group(
195 &self,
196 ciphersuite: Ciphersuite,
197 ) -> PreGroupPartyState<'_, Provider> {
198 self.pre_group_builder(ciphersuite).build().await
199 }
200}
201
202pub struct MemberState<'a, Provider> {
204 pub party: PreGroupPartyState<'a, Provider>,
205 pub group: MlsGroup,
206}
207
208impl<Provider: OpenMlsProvider> MemberState<'_, Provider> {
209 #[openmls_traits::maybe_async]
211 pub async fn get_storage_signature_key_pair(&self) -> Option<SignatureKeyPair> {
212 let ciphersuite = self
213 .party
214 .key_package_bundle
215 .key_package()
216 .ciphersuite()
217 .into();
218
219 SignatureKeyPair::read(
220 self.party.core_state.provider.storage(),
221 self.party.signer.public(),
222 ciphersuite,
223 )
224 .await
225 }
226 #[openmls_traits::maybe_async]
228 pub async fn group_storage_state(&self) -> GroupStorageState {
229 let storage_provider = self.party.core_state.provider.storage();
230 let group_id = self.group.group_id();
231
232 GroupStorageState::from_storage(storage_provider, group_id).await
233 }
234 #[openmls_traits::maybe_async]
236 pub async fn deliver_and_apply(
237 &mut self,
238 message: MlsMessageIn,
239 ) -> Result<(), GroupError<Provider>> {
240 let message = message.try_into_protocol_message()?;
241
242 let processed_message = self
244 .group
245 .process_message(&self.party.core_state.provider, message)
246 .await?;
247
248 match processed_message.into_content() {
249 ProcessedMessageContent::ApplicationMessage(_) => todo!(),
250 ProcessedMessageContent::ProposalMessage(_) => todo!(),
251 ProcessedMessageContent::ExternalJoinProposalMessage(_) => todo!(),
252 ProcessedMessageContent::StagedCommitMessage(m) => {
253 self.group
254 .merge_staged_commit(&self.party.core_state.provider, *m)
255 .await?
256 }
257 ProcessedMessageContent::OwnPendingCommit => {
258 self.group
259 .merge_pending_commit(&self.party.core_state.provider)
260 .await?
261 }
262 ProcessedMessageContent::OwnPrivateMessage => {}
265 #[cfg(feature = "extensions-draft")]
266 ProcessedMessageContent::UnresolvedAppDataCommit(_) => todo!(),
267 };
268
269 Ok(())
270 }
271}
272
273impl<'commit_builder, 'b: 'commit_builder, 'a: 'b, Provider> MemberState<'a, Provider>
274where
275 Provider: openmls_traits::OpenMlsProvider,
276{
277 #[openmls_traits::maybe_async]
279 pub async fn build_commit_and_stage(
280 &'b mut self,
281 f: impl FnOnce(
282 CommitBuilder<'commit_builder, Initial>,
283 ) -> CommitBuilder<'commit_builder, Initial>,
284 ) -> Result<CommitMessageBundle, GroupError<Provider>> {
285 let commit_builder = f(self.group.commit_builder());
286
287 let provider = &self.party.core_state.provider;
288
289 let bundle = commit_builder
291 .load_psks(provider.storage())
292 .await?
293 .build(
294 provider.rand(),
295 provider.crypto(),
296 &self.party.signer,
297 |_| true,
298 )?
299 .stage_commit(provider)
300 .await?;
301
302 Ok(bundle)
303 }
304}
305
306impl<'a, Provider: OpenMlsProvider> MemberState<'a, Provider> {
307 #[openmls_traits::maybe_async]
310 pub async fn create_from_pre_group(
311 party: PreGroupPartyState<'a, Provider>,
312 mls_group_create_config: MlsGroupCreateConfig,
313 group_id: GroupId,
314 ) -> Result<Self, GroupError<Provider>> {
315 let group = MlsGroup::new_with_group_id(
317 &party.core_state.provider,
318 &party.signer,
319 &mls_group_create_config,
320 group_id,
321 party.credential_with_key.clone(),
322 )
323 .await?;
324
325 Ok(Self { party, group })
326 }
327 #[openmls_traits::maybe_async]
330 pub async fn join_from_pre_group(
331 party: PreGroupPartyState<'a, Provider>,
332 mls_group_join_config: MlsGroupJoinConfig,
333 welcome: Welcome,
334 tree: Option<RatchetTreeIn>,
335 ) -> Result<Self, GroupError<Provider>> {
336 let staged_join = StagedWelcome::new_from_welcome(
337 &party.core_state.provider,
338 &mls_group_join_config,
339 welcome,
340 tree,
341 )
342 .await?;
343
344 let group = staged_join.into_group(&party.core_state.provider).await?;
345
346 Ok(Self { party, group })
347 }
348}
349
350pub struct GroupState<'a, Provider> {
352 group_id: GroupId,
353 members: HashMap<Name, MemberState<'a, Provider>>,
354}
355
356impl<'a, Provider: OpenMlsProvider> GroupState<'a, Provider> {
357 #[openmls_traits::maybe_async]
359 pub async fn new_from_party(
360 group_id: GroupId,
361 pre_group_state: PreGroupPartyState<'a, Provider>,
362 mls_group_create_config: MlsGroupCreateConfig,
363 ) -> Result<Self, GroupError<Provider>> {
364 let mut members = HashMap::new();
365
366 let name = pre_group_state.core_state.name;
367 let member_state = MemberState::create_from_pre_group(
368 pre_group_state,
369 mls_group_create_config,
370 group_id.clone(),
371 )
372 .await?;
373
374 members.insert(name, member_state);
375
376 Ok(Self { group_id, members })
377 }
378
379 pub fn members_mut<const N: usize>(
383 &mut self,
384 names: &[Name; N],
385 ) -> [&mut MemberState<'a, Provider>; N] {
386 assert!(N > 0, "must request at least one member");
387 assert!(
388 N <= self.members.len(),
389 "cannot request more members than available"
390 );
391
392 let mut members: [(_, _); N] = self
394 .members
395 .iter_mut()
396 .filter_map(|(member_name, member)| {
397 let index = names.iter().position(|name| name == member_name)?;
401
402 Some((index, member))
403 })
404 .collect::<Vec<_>>()
406 .try_into()
407 .ok()
408 .expect("At least one requested member not found");
409
410 members.sort_by_key(|(pos, _member)| *pos);
412
413 members.map(|(_pos, member)| member)
414 }
415
416 #[openmls_traits::maybe_async]
418 pub async fn deliver_and_apply(
419 &mut self,
420 message: MlsMessageIn,
421 ) -> Result<(), GroupError<Provider>> {
422 self.deliver_and_apply_if(message, |_| true).await
423 }
424 #[openmls_traits::maybe_async]
426 pub async fn deliver_and_apply_if(
427 &mut self,
428 message: MlsMessageIn,
429 condition: impl Fn(&MemberState<'a, Provider>) -> bool,
430 ) -> Result<(), GroupError<Provider>> {
431 for member in self.members.values_mut() {
432 if condition(member) {
433 member.deliver_and_apply(message.clone()).await?;
434 }
435 }
436
437 Ok(())
438 }
439
440 #[openmls_traits::maybe_async]
442 pub async fn deliver_and_apply_welcome(
443 &mut self,
444 recipient: PreGroupPartyState<'a, Provider>,
445 mls_group_join_config: MlsGroupJoinConfig,
446 welcome: Welcome,
447 tree: Option<RatchetTreeIn>,
448 ) -> Result<(), GroupError<Provider>> {
449 let name = recipient.core_state.name;
451
452 let member_state =
453 MemberState::join_from_pre_group(recipient, mls_group_join_config, welcome, tree)
454 .await?;
455
456 self.members.insert(name, member_state);
458
459 Ok(())
460 }
461
462 pub fn untrack_member(&mut self, name: Name) {
465 let _ = self.members.remove(&name);
466 }
467
468 #[openmls_traits::maybe_async]
469 pub async fn add_member(
470 &mut self,
471 add_config: AddMemberConfig<'a, Provider>,
472 ) -> Result<(), GroupError<Provider>> {
473 let adder = self
474 .members
475 .get_mut(add_config.adder)
476 .ok_or(TestError::NoSuchMember)?;
477
478 let key_packages: Vec<_> = add_config
479 .addees
480 .iter()
481 .map(|addee| addee.key_package_bundle.key_package.clone())
482 .collect();
483
484 let (commit, welcome, _) = adder
485 .group
486 .add_members(
487 &adder.party.core_state.provider,
488 &adder.party.signer,
489 &key_packages,
490 )
491 .await?;
492
493 self.deliver_and_apply_if(commit.into(), |member| {
495 member.party.core_state.name != add_config.adder
496 })
497 .await?;
498
499 let welcome = match welcome.body() {
501 MlsMessageBodyOut::Welcome(welcome) => welcome.clone(),
502 _ => panic!("No welcome returned"),
503 };
504
505 for addee in add_config.addees.into_iter() {
506 self.deliver_and_apply_welcome(
507 addee,
508 add_config.join_config.clone(),
509 welcome.clone(),
510 None,
511 )
512 .await?;
513 }
514
515 let adder = self
516 .members
517 .get_mut(add_config.adder)
518 .ok_or(TestError::NoSuchMember)?;
519
520 let staged_commit = adder.group.pending_commit().unwrap().clone();
521
522 adder
523 .group
524 .merge_staged_commit(&adder.party.core_state.provider, staged_commit)
525 .await?;
526
527 Ok(())
528 }
529
530 pub fn group_id(&self) -> GroupId {
532 self.group_id.clone()
533 }
534}
535
536impl MlsGroupCreateConfig {
537 pub fn test_default_from_ciphersuite(ciphersuite: Ciphersuite) -> Self {
539 MlsGroupCreateConfig::builder()
540 .ciphersuite(ciphersuite)
541 .use_ratchet_tree_extension(true)
542 .wire_format_policy(PURE_PLAINTEXT_WIRE_FORMAT_POLICY) .build()
544 }
545}
546
547pub struct AddMemberConfig<'a, Provider> {
548 pub adder: Name,
549 pub addees: Vec<PreGroupPartyState<'a, Provider>>,
550 pub join_config: MlsGroupJoinConfig,
551 pub tree: Option<RatchetTreeIn>,
552}
553
554#[cfg(test)]
555mod test {
556
557 use super::*;
558 use openmls_test::openmls_test;
559
560 #[openmls_test]
561 fn test_members_mut() {
562 let alice_party = CorePartyState::<Provider>::new("alice");
563 let bob_party = CorePartyState::<Provider>::new("bob");
564 let charlie_party = CorePartyState::<Provider>::new("charlie");
565 let dave_party = CorePartyState::<Provider>::new("dave");
566
567 let alice_pre_group = alice_party.generate_pre_group(ciphersuite);
568 let bob_pre_group = bob_party.generate_pre_group(ciphersuite);
569 let charlie_pre_group = charlie_party.generate_pre_group(ciphersuite);
570 let dave_pre_group = dave_party.generate_pre_group(ciphersuite);
571
572 let mls_group_create_config = MlsGroupCreateConfig::builder()
574 .ciphersuite(ciphersuite)
575 .use_ratchet_tree_extension(true)
576 .build();
577
578 let mls_group_join_config = mls_group_create_config.join_config().clone();
580
581 let group_id = GroupId::from_slice(b"test");
583 let mut group_state =
584 GroupState::new_from_party(group_id, alice_pre_group, mls_group_create_config).unwrap();
585
586 group_state
587 .add_member(AddMemberConfig {
588 adder: "alice",
589 addees: vec![bob_pre_group, charlie_pre_group, dave_pre_group],
590 join_config: mls_group_join_config.clone(),
591 tree: None,
592 })
593 .expect("Could not add member");
594
595 let [alice, bob, charlie, dave] =
597 group_state.members_mut(&["alice", "bob", "charlie", "dave"]);
598 assert_eq!(alice.party.core_state.name, "alice");
599 assert_eq!(bob.party.core_state.name, "bob");
600 assert_eq!(charlie.party.core_state.name, "charlie");
601 assert_eq!(dave.party.core_state.name, "dave");
602
603 let [dave, charlie, bob, alice] =
604 group_state.members_mut(&["dave", "charlie", "bob", "alice"]);
605 assert_eq!(alice.party.core_state.name, "alice");
606 assert_eq!(bob.party.core_state.name, "bob");
607 assert_eq!(charlie.party.core_state.name, "charlie");
608 assert_eq!(dave.party.core_state.name, "dave");
609
610 let [dave, bob, charlie, alice] =
611 group_state.members_mut(&["dave", "bob", "charlie", "alice"]);
612 assert_eq!(alice.party.core_state.name, "alice");
613 assert_eq!(bob.party.core_state.name, "bob");
614 assert_eq!(charlie.party.core_state.name, "charlie");
615 assert_eq!(dave.party.core_state.name, "dave");
616
617 let [dave, bob] = group_state.members_mut(&["dave", "bob"]);
618 assert_eq!(bob.party.core_state.name, "bob");
619 assert_eq!(dave.party.core_state.name, "dave");
620
621 let [alice, charlie] = group_state.members_mut(&["alice", "charlie"]);
622 assert_eq!(alice.party.core_state.name, "alice");
623 assert_eq!(charlie.party.core_state.name, "charlie");
624 }
625 #[openmls_test]
626 pub fn simpler_example() {
627 let alice_party = CorePartyState::<Provider>::new("alice");
628 let bob_party = CorePartyState::<Provider>::new("bob");
629 let charlie_party = CorePartyState::<Provider>::new("charlie");
630 let dave_party = CorePartyState::<Provider>::new("dave");
631
632 let alice_pre_group = alice_party.generate_pre_group(ciphersuite);
633 let bob_pre_group = bob_party.generate_pre_group(ciphersuite);
634 let charlie_pre_group = charlie_party.generate_pre_group(ciphersuite);
635 let dave_pre_group = dave_party.generate_pre_group(ciphersuite);
636
637 let mls_group_create_config = MlsGroupCreateConfig::builder()
639 .ciphersuite(ciphersuite)
640 .use_ratchet_tree_extension(true)
641 .build();
642
643 let mls_group_join_config = mls_group_create_config.join_config().clone();
645
646 let group_id = GroupId::from_slice(b"test");
648 let mut group_state =
649 GroupState::new_from_party(group_id, alice_pre_group, mls_group_create_config).unwrap();
650
651 group_state
652 .add_member(AddMemberConfig {
653 adder: "alice",
654 addees: vec![bob_pre_group, charlie_pre_group],
655 join_config: mls_group_join_config.clone(),
656 tree: None,
657 })
658 .expect("Could not add member");
659
660 group_state.assert_membership();
661
662 group_state
663 .add_member(AddMemberConfig {
664 adder: "bob",
665 addees: vec![dave_pre_group],
666 join_config: mls_group_join_config,
667 tree: None,
668 })
669 .expect("Could not add member");
670
671 group_state.assert_membership();
672 }
673
674 #[openmls_test]
675 pub fn simple_example() {
676 let alice_party = CorePartyState::<Provider>::new("alice");
677 let bob_party = CorePartyState::<Provider>::new("bob");
678
679 let alice_pre_group = alice_party.generate_pre_group(ciphersuite);
680 let bob_pre_group = bob_party.generate_pre_group(ciphersuite);
681
682 let bob_key_package = bob_pre_group.key_package_bundle.key_package.clone();
685
686 let mls_group_create_config = MlsGroupCreateConfig::builder()
688 .ciphersuite(ciphersuite)
689 .use_ratchet_tree_extension(true)
690 .build();
691
692 let mls_group_join_config = mls_group_create_config.join_config().clone();
694
695 let group_id = GroupId::from_slice(b"test");
697 let mut group_state =
698 GroupState::new_from_party(group_id, alice_pre_group, mls_group_create_config).unwrap();
699
700 let [alice] = group_state.members_mut(&["alice"]);
702
703 let bundle = alice
705 .build_commit_and_stage(move |builder| {
706 let add_proposal = Proposal::add(AddProposal {
707 key_package: bob_key_package,
708 });
709
710 builder
713 .consume_proposal_store(false)
714 .add_proposal(add_proposal)
715 })
716 .expect("Could not stage commit");
717
718 let welcome = bundle.welcome().unwrap().clone();
720 group_state
721 .deliver_and_apply_welcome(bob_pre_group, mls_group_join_config, welcome, None)
722 .expect("Error delivering and applying welcome");
723
724 let [alice] = group_state.members_mut(&["alice"]);
725
726 let staged_commit = alice.group.pending_commit().unwrap().clone();
727
728 alice
729 .group
730 .merge_staged_commit(&alice.party.core_state.provider, staged_commit)
731 .expect("Error merging staged commit");
732
733 group_state.assert_membership();
734 }
735}