openmls/test_utils/single_group_test_framework/
mod.rs

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
16use crate::test_utils::storage_state::GroupStorageState;
17
18mod assertions;
19
20mod errors;
21pub use errors::GroupError;
22use errors::*;
23
24use std::collections::HashMap;
25
26// type alias for &'static str
27type Name = &'static str;
28
29// TODO: only define this once
30/// Helper function for generating a credential.
31pub fn generate_credential(
32    identity: Vec<u8>,
33    signature_algorithm: SignatureScheme,
34    provider: &impl crate::storage::OpenMlsProvider,
35) -> (CredentialWithKey, SignatureKeyPair) {
36    let credential = BasicCredential::new(identity);
37    let signature_keys = SignatureKeyPair::new(signature_algorithm).unwrap();
38    signature_keys.store(provider.storage()).unwrap();
39
40    (
41        CredentialWithKey {
42            credential: credential.into(),
43            signature_key: signature_keys.to_public_vec().into(),
44        },
45        signature_keys,
46    )
47}
48
49// TODO: only define this once
50pub(crate) fn generate_key_package(
51    ciphersuite: Ciphersuite,
52    credential_with_key: CredentialWithKey,
53    extensions: Extensions<KeyPackage>,
54    provider: &impl crate::storage::OpenMlsProvider,
55    lifetime: impl Into<Option<Lifetime>>,
56    signer: &impl Signer,
57) -> KeyPackageBundle {
58    let mut builder = KeyPackage::builder().key_package_extensions(extensions);
59
60    if let Some(lifetime) = lifetime.into() {
61        builder = builder.key_package_lifetime(lifetime);
62    }
63
64    builder
65        .build(ciphersuite, provider, signer, credential_with_key)
66        .unwrap()
67}
68
69/// Struct representing a party's global state
70pub struct CorePartyState<Provider> {
71    pub name: Name,
72    pub provider: Provider,
73}
74
75impl<Provider: Default> CorePartyState<Provider> {
76    pub fn new(name: Name) -> Self {
77        Self {
78            name,
79            provider: Provider::default(),
80        }
81    }
82}
83
84/// Struct representing a party's state before joining a group
85pub struct PreGroupPartyState<'a, Provider> {
86    pub credential_with_key: CredentialWithKey,
87    // TODO: regenerate?
88    pub key_package_bundle: KeyPackageBundle,
89    pub signer: SignatureKeyPair,
90    pub core_state: &'a CorePartyState<Provider>,
91}
92
93pub struct PreGroupPartyStateBuilder<'a, Provider: OpenMlsProvider> {
94    ciphersuite: Ciphersuite,
95    lifetime: Option<Lifetime>,
96    key_package_extensions: Option<Extensions<KeyPackage>>,
97    leaf_node_extensions: Option<Extensions<LeafNode>>,
98    leaf_node_capabilities: Option<Capabilities>,
99    core_state: &'a CorePartyState<Provider>,
100}
101
102impl<'a, Provider: OpenMlsProvider> PreGroupPartyStateBuilder<'a, Provider> {
103    pub fn with_lifetime(mut self, lifetime: impl Into<Option<Lifetime>>) -> Self {
104        self.lifetime = lifetime.into();
105
106        self
107    }
108    pub fn with_key_package_extensions(
109        mut self,
110        extensions: impl Into<Option<Extensions<KeyPackage>>>,
111    ) -> Self {
112        self.key_package_extensions = extensions.into();
113
114        self
115    }
116    pub fn with_leaf_node_extensions(
117        mut self,
118        extensions: impl Into<Option<Extensions<LeafNode>>>,
119    ) -> Self {
120        self.leaf_node_extensions = extensions.into();
121
122        self
123    }
124    pub fn with_leaf_node_capabilities(
125        mut self,
126        capabilities: impl Into<Option<Capabilities>>,
127    ) -> Self {
128        self.leaf_node_capabilities = capabilities.into();
129
130        self
131    }
132
133    pub fn build(self) -> PreGroupPartyState<'a, Provider> {
134        let (credential_with_key, signer) = generate_credential(
135            self.core_state.name.into(),
136            self.ciphersuite.signature_algorithm(),
137            &self.core_state.provider,
138        );
139        let mut builder = KeyPackage::builder()
140            .leaf_node_extensions(self.leaf_node_extensions.unwrap_or_default())
141            .key_package_extensions(self.key_package_extensions.unwrap_or_default())
142            .leaf_node_capabilities(self.leaf_node_capabilities.unwrap_or_default());
143
144        if let Some(lifetime) = self.lifetime {
145            builder = builder.key_package_lifetime(lifetime);
146        }
147
148        let key_package_bundle = builder
149            .build(
150                self.ciphersuite,
151                &self.core_state.provider,
152                &signer,
153                credential_with_key.clone(),
154            )
155            .unwrap();
156
157        PreGroupPartyState {
158            credential_with_key,
159            key_package_bundle,
160            signer,
161            core_state: self.core_state,
162        }
163    }
164}
165
166impl<Provider: OpenMlsProvider> CorePartyState<Provider> {
167    /// Returns a builder for the [`CorePartyState`].
168    pub fn pre_group_builder<'a>(
169        &'a self,
170        ciphersuite: Ciphersuite,
171    ) -> PreGroupPartyStateBuilder<'a, Provider> {
172        PreGroupPartyStateBuilder {
173            ciphersuite,
174            lifetime: None,
175            key_package_extensions: None,
176            leaf_node_extensions: None,
177            leaf_node_capabilities: None,
178            core_state: self,
179        }
180    }
181
182    /// Generates a simple pre-group state for a `CorePartyState`
183    pub fn generate_pre_group(&self, ciphersuite: Ciphersuite) -> PreGroupPartyState<'_, Provider> {
184        self.pre_group_builder(ciphersuite).build()
185    }
186}
187
188/// Represents a group member's `MlsGroup` instance and pre-group state
189pub struct MemberState<'a, Provider> {
190    pub party: PreGroupPartyState<'a, Provider>,
191    pub group: MlsGroup,
192}
193
194impl<Provider: OpenMlsProvider> MemberState<'_, Provider> {
195    /// Get member's `SignatureKeyPair` if available
196    pub fn get_storage_signature_key_pair(&self) -> Option<SignatureKeyPair> {
197        let ciphersuite = self
198            .party
199            .key_package_bundle
200            .key_package()
201            .ciphersuite()
202            .into();
203
204        SignatureKeyPair::read(
205            self.party.core_state.provider.storage(),
206            self.party.signer.public(),
207            ciphersuite,
208        )
209    }
210    /// Get the `GroupStorageState` for this group
211    pub fn group_storage_state(&self) -> GroupStorageState {
212        let storage_provider = self.party.core_state.provider.storage();
213        let group_id = self.group.group_id();
214
215        GroupStorageState::from_storage(storage_provider, group_id)
216    }
217    /// Deliver_and_apply a message to this member's `MlsGroup`
218    pub fn deliver_and_apply(&mut self, message: MlsMessageIn) -> Result<(), GroupError<Provider>> {
219        let message = message.try_into_protocol_message()?;
220
221        // process message
222        let processed_message = self
223            .group
224            .process_message(&self.party.core_state.provider, message)?;
225
226        match processed_message.into_content() {
227            ProcessedMessageContent::ApplicationMessage(_) => todo!(),
228            ProcessedMessageContent::ProposalMessage(_) => todo!(),
229            ProcessedMessageContent::ExternalJoinProposalMessage(_) => todo!(),
230            ProcessedMessageContent::StagedCommitMessage(m) => self
231                .group
232                .merge_staged_commit(&self.party.core_state.provider, *m)?,
233            ProcessedMessageContent::OwnPendingCommit => self
234                .group
235                .merge_pending_commit(&self.party.core_state.provider)?,
236            // Own PrivateMessages echoed by the DS cannot be decrypted, so skip
237            // them.
238            ProcessedMessageContent::OwnPrivateMessage => {}
239            #[cfg(feature = "extensions-draft")]
240            ProcessedMessageContent::UnresolvedAppDataCommit(_) => todo!(),
241        };
242
243        Ok(())
244    }
245}
246
247impl<'commit_builder, 'b: 'commit_builder, 'a: 'b, Provider> MemberState<'a, Provider>
248where
249    Provider: openmls_traits::OpenMlsProvider,
250{
251    /// Build and stage a commit, using the provided closure to add proposals
252    pub fn build_commit_and_stage(
253        &'b mut self,
254        f: impl FnOnce(
255            CommitBuilder<'commit_builder, Initial>,
256        ) -> CommitBuilder<'commit_builder, Initial>,
257    ) -> Result<CommitMessageBundle, GroupError<Provider>> {
258        let commit_builder = f(self.group.commit_builder());
259
260        let provider = &self.party.core_state.provider;
261
262        // TODO: most of the steps here cannot be done via the closure (yet)
263        let bundle = commit_builder
264            .load_psks(provider.storage())?
265            .build(
266                provider.rand(),
267                provider.crypto(),
268                &self.party.signer,
269                |_| true,
270            )?
271            .stage_commit(provider)?;
272
273        Ok(bundle)
274    }
275}
276
277impl<'a, Provider: OpenMlsProvider> MemberState<'a, Provider> {
278    /// Create a `MemberState` from a `PreGroupPartyState`. This creates a new `MlsGroup` with one
279    /// member
280    pub fn create_from_pre_group(
281        party: PreGroupPartyState<'a, Provider>,
282        mls_group_create_config: MlsGroupCreateConfig,
283        group_id: GroupId,
284    ) -> Result<Self, GroupError<Provider>> {
285        // initialize MlsGroup
286        let group = MlsGroup::new_with_group_id(
287            &party.core_state.provider,
288            &party.signer,
289            &mls_group_create_config,
290            group_id,
291            party.credential_with_key.clone(),
292        )?;
293
294        Ok(Self { party, group })
295    }
296    /// Create a `MemberState` from a `Welcome`, which creates a new `MlsGroup` using a `Welcome`
297    /// invitation from an existing group
298    pub fn join_from_pre_group(
299        party: PreGroupPartyState<'a, Provider>,
300        mls_group_join_config: MlsGroupJoinConfig,
301        welcome: Welcome,
302        tree: Option<RatchetTreeIn>,
303    ) -> Result<Self, GroupError<Provider>> {
304        let staged_join = StagedWelcome::new_from_welcome(
305            &party.core_state.provider,
306            &mls_group_join_config,
307            welcome,
308            tree,
309        )?;
310
311        let group = staged_join.into_group(&party.core_state.provider)?;
312
313        Ok(Self { party, group })
314    }
315}
316
317/// All of the state for a group and its members
318pub struct GroupState<'a, Provider> {
319    group_id: GroupId,
320    members: HashMap<Name, MemberState<'a, Provider>>,
321}
322
323impl<'a, Provider: OpenMlsProvider> GroupState<'a, Provider> {
324    /// Create a new `GroupState` from a single party
325    pub fn new_from_party(
326        group_id: GroupId,
327        pre_group_state: PreGroupPartyState<'a, Provider>,
328        mls_group_create_config: MlsGroupCreateConfig,
329    ) -> Result<Self, GroupError<Provider>> {
330        let mut members = HashMap::new();
331
332        let name = pre_group_state.core_state.name;
333        let member_state = MemberState::create_from_pre_group(
334            pre_group_state,
335            mls_group_create_config,
336            group_id.clone(),
337        )?;
338
339        members.insert(name, member_state);
340
341        Ok(Self { group_id, members })
342    }
343
344    /// Get mutable references to specified `MemberState`s as a fixed-size array,
345    /// in the order of the names provided in `names`.
346    /// At least one member must be requested.
347    pub fn members_mut<const N: usize>(
348        &mut self,
349        names: &[Name; N],
350    ) -> [&mut MemberState<'a, Provider>; N] {
351        assert!(N > 0, "must request at least one member");
352        assert!(
353            N <= self.members.len(),
354            "cannot request more members than available"
355        );
356
357        // map each member in `self.members` to its name's index in `names`
358        let mut members: [(_, _); N] = self
359            .members
360            .iter_mut()
361            .filter_map(|(member_name, member)| {
362                // Find the index of the member's name in `names`
363                // NOTE: the list of names provided to this method will generally be short,
364                // so not many comparisons are made here.
365                let index = names.iter().position(|name| name == member_name)?;
366
367                Some((index, member))
368            })
369            // collect into Vec, then into fixed-size array
370            .collect::<Vec<_>>()
371            .try_into()
372            .ok()
373            .expect("At least one requested member not found");
374
375        // sort by index
376        members.sort_by_key(|(pos, _member)| *pos);
377
378        members.map(|(_pos, member)| member)
379    }
380
381    /// Deliver_and_apply a message to all parties
382    pub fn deliver_and_apply(&mut self, message: MlsMessageIn) -> Result<(), GroupError<Provider>> {
383        self.deliver_and_apply_if(message, |_| true)
384    }
385    /// Deliver_and_apply a message to all parties if a provided condition is met
386    pub fn deliver_and_apply_if(
387        &mut self,
388        message: MlsMessageIn,
389        condition: impl Fn(&MemberState<'a, Provider>) -> bool,
390    ) -> Result<(), GroupError<Provider>> {
391        self.members
392            .values_mut()
393            .filter(|member| condition(member))
394            .try_for_each(|member| member.deliver_and_apply(message.clone()))?;
395
396        Ok(())
397    }
398
399    /// Deliver_and_apply a welcome to a single party, and initialize a group for that party
400    pub fn deliver_and_apply_welcome(
401        &mut self,
402        recipient: PreGroupPartyState<'a, Provider>,
403        mls_group_join_config: MlsGroupJoinConfig,
404        welcome: Welcome,
405        tree: Option<RatchetTreeIn>,
406    ) -> Result<(), GroupError<Provider>> {
407        // create new group
408        let name = recipient.core_state.name;
409
410        let member_state =
411            MemberState::join_from_pre_group(recipient, mls_group_join_config, welcome, tree)?;
412
413        // insert after success
414        self.members.insert(name, member_state);
415
416        Ok(())
417    }
418
419    /// Drop a member from the internal hashmap. This does not delete the member from any
420    /// `MlsGroup`
421    pub fn untrack_member(&mut self, name: Name) {
422        let _ = self.members.remove(&name);
423    }
424
425    pub fn add_member(
426        &mut self,
427        add_config: AddMemberConfig<'a, Provider>,
428    ) -> Result<(), GroupError<Provider>> {
429        let adder = self
430            .members
431            .get_mut(add_config.adder)
432            .ok_or(TestError::NoSuchMember)?;
433
434        let key_packages: Vec<_> = add_config
435            .addees
436            .iter()
437            .map(|addee| addee.key_package_bundle.key_package.clone())
438            .collect();
439
440        let (commit, welcome, _) = adder.group.add_members(
441            &adder.party.core_state.provider,
442            &adder.party.signer,
443            &key_packages,
444        )?;
445
446        // Deliver_and_apply to all members but adder
447        self.deliver_and_apply_if(commit.into(), |member| {
448            member.party.core_state.name != add_config.adder
449        })?;
450
451        // Deliver_and_apply welcome to addee
452        let welcome = match welcome.body() {
453            MlsMessageBodyOut::Welcome(welcome) => welcome.clone(),
454            _ => panic!("No welcome returned"),
455        };
456
457        for addee in add_config.addees.into_iter() {
458            self.deliver_and_apply_welcome(
459                addee,
460                add_config.join_config.clone(),
461                welcome.clone(),
462                None,
463            )?;
464        }
465
466        let adder = self
467            .members
468            .get_mut(add_config.adder)
469            .ok_or(TestError::NoSuchMember)?;
470
471        let staged_commit = adder.group.pending_commit().unwrap().clone();
472
473        adder
474            .group
475            .merge_staged_commit(&adder.party.core_state.provider, staged_commit)?;
476
477        Ok(())
478    }
479
480    /// Returns a copy of the GroupId
481    pub fn group_id(&self) -> GroupId {
482        self.group_id.clone()
483    }
484}
485
486impl MlsGroupCreateConfig {
487    /// Default config for test framework
488    pub fn test_default_from_ciphersuite(ciphersuite: Ciphersuite) -> Self {
489        MlsGroupCreateConfig::builder()
490            .ciphersuite(ciphersuite)
491            .use_ratchet_tree_extension(true)
492            .wire_format_policy(PURE_PLAINTEXT_WIRE_FORMAT_POLICY) // Important because the secret tree might diverge otherwise
493            .build()
494    }
495}
496
497pub struct AddMemberConfig<'a, Provider> {
498    pub adder: Name,
499    pub addees: Vec<PreGroupPartyState<'a, Provider>>,
500    pub join_config: MlsGroupJoinConfig,
501    pub tree: Option<RatchetTreeIn>,
502}
503
504#[cfg(test)]
505mod test {
506
507    use super::*;
508    use openmls_test::openmls_test;
509
510    #[openmls_test]
511    fn test_members_mut() {
512        let alice_party = CorePartyState::<Provider>::new("alice");
513        let bob_party = CorePartyState::<Provider>::new("bob");
514        let charlie_party = CorePartyState::<Provider>::new("charlie");
515        let dave_party = CorePartyState::<Provider>::new("dave");
516
517        let alice_pre_group = alice_party.generate_pre_group(ciphersuite);
518        let bob_pre_group = bob_party.generate_pre_group(ciphersuite);
519        let charlie_pre_group = charlie_party.generate_pre_group(ciphersuite);
520        let dave_pre_group = dave_party.generate_pre_group(ciphersuite);
521
522        // Create config
523        let mls_group_create_config = MlsGroupCreateConfig::builder()
524            .ciphersuite(ciphersuite)
525            .use_ratchet_tree_extension(true)
526            .build();
527
528        // Join config
529        let mls_group_join_config = mls_group_create_config.join_config().clone();
530
531        // Initialize the group state
532        let group_id = GroupId::from_slice(b"test");
533        let mut group_state =
534            GroupState::new_from_party(group_id, alice_pre_group, mls_group_create_config).unwrap();
535
536        group_state
537            .add_member(AddMemberConfig {
538                adder: "alice",
539                addees: vec![bob_pre_group, charlie_pre_group, dave_pre_group],
540                join_config: mls_group_join_config.clone(),
541                tree: None,
542            })
543            .expect("Could not add member");
544
545        // test different orderings
546        let [alice, bob, charlie, dave] =
547            group_state.members_mut(&["alice", "bob", "charlie", "dave"]);
548        assert_eq!(alice.party.core_state.name, "alice");
549        assert_eq!(bob.party.core_state.name, "bob");
550        assert_eq!(charlie.party.core_state.name, "charlie");
551        assert_eq!(dave.party.core_state.name, "dave");
552
553        let [dave, charlie, bob, alice] =
554            group_state.members_mut(&["dave", "charlie", "bob", "alice"]);
555        assert_eq!(alice.party.core_state.name, "alice");
556        assert_eq!(bob.party.core_state.name, "bob");
557        assert_eq!(charlie.party.core_state.name, "charlie");
558        assert_eq!(dave.party.core_state.name, "dave");
559
560        let [dave, bob, charlie, alice] =
561            group_state.members_mut(&["dave", "bob", "charlie", "alice"]);
562        assert_eq!(alice.party.core_state.name, "alice");
563        assert_eq!(bob.party.core_state.name, "bob");
564        assert_eq!(charlie.party.core_state.name, "charlie");
565        assert_eq!(dave.party.core_state.name, "dave");
566
567        let [dave, bob] = group_state.members_mut(&["dave", "bob"]);
568        assert_eq!(bob.party.core_state.name, "bob");
569        assert_eq!(dave.party.core_state.name, "dave");
570
571        let [alice, charlie] = group_state.members_mut(&["alice", "charlie"]);
572        assert_eq!(alice.party.core_state.name, "alice");
573        assert_eq!(charlie.party.core_state.name, "charlie");
574    }
575    #[openmls_test]
576    pub fn simpler_example() {
577        let alice_party = CorePartyState::<Provider>::new("alice");
578        let bob_party = CorePartyState::<Provider>::new("bob");
579        let charlie_party = CorePartyState::<Provider>::new("charlie");
580        let dave_party = CorePartyState::<Provider>::new("dave");
581
582        let alice_pre_group = alice_party.generate_pre_group(ciphersuite);
583        let bob_pre_group = bob_party.generate_pre_group(ciphersuite);
584        let charlie_pre_group = charlie_party.generate_pre_group(ciphersuite);
585        let dave_pre_group = dave_party.generate_pre_group(ciphersuite);
586
587        // Create config
588        let mls_group_create_config = MlsGroupCreateConfig::builder()
589            .ciphersuite(ciphersuite)
590            .use_ratchet_tree_extension(true)
591            .build();
592
593        // Join config
594        let mls_group_join_config = mls_group_create_config.join_config().clone();
595
596        // Initialize the group state
597        let group_id = GroupId::from_slice(b"test");
598        let mut group_state =
599            GroupState::new_from_party(group_id, alice_pre_group, mls_group_create_config).unwrap();
600
601        group_state
602            .add_member(AddMemberConfig {
603                adder: "alice",
604                addees: vec![bob_pre_group, charlie_pre_group],
605                join_config: mls_group_join_config.clone(),
606                tree: None,
607            })
608            .expect("Could not add member");
609
610        group_state.assert_membership();
611
612        group_state
613            .add_member(AddMemberConfig {
614                adder: "bob",
615                addees: vec![dave_pre_group],
616                join_config: mls_group_join_config,
617                tree: None,
618            })
619            .expect("Could not add member");
620
621        group_state.assert_membership();
622    }
623
624    #[openmls_test]
625    pub fn simple_example() {
626        let alice_party = CorePartyState::<Provider>::new("alice");
627        let bob_party = CorePartyState::<Provider>::new("bob");
628
629        let alice_pre_group = alice_party.generate_pre_group(ciphersuite);
630        let bob_pre_group = bob_party.generate_pre_group(ciphersuite);
631
632        // Get the key package for Bob
633        // TODO: should key package be regenerated each time?
634        let bob_key_package = bob_pre_group.key_package_bundle.key_package.clone();
635
636        // Create config
637        let mls_group_create_config = MlsGroupCreateConfig::builder()
638            .ciphersuite(ciphersuite)
639            .use_ratchet_tree_extension(true)
640            .build();
641
642        // Join config
643        let mls_group_join_config = mls_group_create_config.join_config().clone();
644
645        // Initialize the group state
646        let group_id = GroupId::from_slice(b"test");
647        let mut group_state =
648            GroupState::new_from_party(group_id, alice_pre_group, mls_group_create_config).unwrap();
649
650        // Get a mutable reference to Alice's group representation
651        let [alice] = group_state.members_mut(&["alice"]);
652
653        // Build a commit with a single add proposal
654        let bundle = alice
655            .build_commit_and_stage(move |builder| {
656                let add_proposal = Proposal::add(AddProposal {
657                    key_package: bob_key_package,
658                });
659
660                // ...add more proposals here...
661
662                builder
663                    .consume_proposal_store(false)
664                    .add_proposal(add_proposal)
665            })
666            .expect("Could not stage commit");
667
668        // Deliver and apply welcome to Bob
669        let welcome = bundle.welcome().unwrap().clone();
670        group_state
671            .deliver_and_apply_welcome(bob_pre_group, mls_group_join_config, welcome, None)
672            .expect("Error delivering and applying welcome");
673
674        let [alice] = group_state.members_mut(&["alice"]);
675
676        let staged_commit = alice.group.pending_commit().unwrap().clone();
677
678        alice
679            .group
680            .merge_staged_commit(&alice.party.core_state.provider, staged_commit)
681            .expect("Error merging staged commit");
682
683        group_state.assert_membership();
684    }
685}
openmls/test_utils/single_group_test_framework/mod.rs

openmls/test_utils/single_group_test_framework/
mod.rs
