xmtp_mls/groups/welcomes/
validated_membership.rs

1use crate::context::XmtpSharedContext;
2use crate::groups::{GroupError, GroupMembership};
3use crate::identity::parse_credential;
4use crate::identity_updates::{
5    IdentityDependencyError, IdentityRequirement, InstallationDiffError, require_association_state,
6    resolve_identity_requirements,
7};
8use openmls::prelude::{BasicCredential, StagedWelcome};
9use std::collections::{HashMap, HashSet};
10use xmtp_common::snippet::Snippet;
11use xmtp_db::DbQuery;
12use xmtp_mls_validation::commit::extract_group_membership;
13
14/// Validate public trial membership, then recheck exact proofs under the writer.
15#[allow(async_fn_in_trait)]
16pub trait ValidateGroupMembership {
17    /// Resolve identity proofs and check that the tree matches the membership extension.
18    async fn check_initial_membership(&self, welcome: &WelcomeMembership)
19    -> Result<(), GroupError>;
20
21    /// Recheck exact proofs against the database read under the state writer.
22    fn check_verified_membership(
23        &self,
24        _welcome: &WelcomeMembership,
25        _db: &impl DbQuery,
26    ) -> Result<(), GroupError> {
27        Ok(())
28    }
29}
30
31/// Public membership data from a trial decode. It contains no mutable MLS state.
32#[derive(Debug, PartialEq)]
33pub struct WelcomeMembership {
34    /// Authenticated inbox membership and exact identity sequence references.
35    membership: GroupMembership,
36    /// Tree members as inbox IDs and installation signature keys.
37    members: Vec<(String, Vec<u8>)>,
38}
39
40impl WelcomeMembership {
41    /// Copy public data from a trial without retaining staged MLS state.
42    pub(crate) fn from_staged(welcome: &StagedWelcome) -> Result<Self, GroupError> {
43        let extensions = welcome.public_group().group_context().extensions();
44        let membership =
45            extract_group_membership(extensions).map_err(|_| GroupError::InvalidWelcomeMetadata)?;
46        let members = welcome
47            .public_group()
48            .members()
49            .map(|member| {
50                let credential = BasicCredential::try_from(member.credential.clone())?;
51                Ok((
52                    parse_credential(credential.identity())?,
53                    member.signature_key,
54                ))
55            })
56            .collect::<Result<_, GroupError>>()?;
57        Ok(Self {
58            membership,
59            members,
60        })
61    }
62
63    fn requirements(&self) -> impl Iterator<Item = IdentityRequirement> + '_ {
64        self.membership
65            .members
66            .iter()
67            .map(|(inbox_id, sequence_id)| IdentityRequirement {
68                inbox_id: inbox_id.clone(),
69                sequence_id: *sequence_id,
70            })
71    }
72
73    /// Reject zero identity references and references at or after this Welcome.
74    // implements: JOIN-050
75    pub(crate) fn validate_sequences(&self, welcome_sequence: u64) -> Result<(), GroupError> {
76        for requirement in self.requirements() {
77            if requirement.sequence_id == 0 || requirement.sequence_id >= welcome_sequence {
78                return Err(
79                    InstallationDiffError::from(IdentityDependencyError::InvalidSequence(
80                        requirement.sequence_id,
81                    ))
82                    .into(),
83                );
84            }
85        }
86        Ok(())
87    }
88}
89
90/// Check every installation against the exact identity state named by the Welcome.
91pub struct InitialMembershipValidator<C> {
92    context: C,
93}
94
95impl<C> InitialMembershipValidator<C> {
96    pub fn new(context: C) -> InitialMembershipValidator<C> {
97        Self { context }
98    }
99}
100
101impl<C> ValidateGroupMembership for InitialMembershipValidator<C>
102where
103    C: XmtpSharedContext,
104{
105    // implements: JOIN-051
106    async fn check_initial_membership(
107        &self,
108        welcome: &WelcomeMembership,
109    ) -> Result<(), GroupError> {
110        for (_, result) in
111            resolve_identity_requirements(&self.context, welcome.requirements()).await
112        {
113            result.map_err(InstallationDiffError::from)?;
114        }
115        self.check_verified_membership(welcome, &self.context.db())
116    }
117
118    // implements: JOIN-052, JOIN-053
119    fn check_verified_membership(
120        &self,
121        welcome: &WelcomeMembership,
122        db: &impl DbQuery,
123    ) -> Result<(), GroupError> {
124        let membership = &welcome.membership;
125        let mut expected_members = HashMap::<String, HashSet<Vec<u8>>>::new();
126        for requirement in welcome.requirements() {
127            let association_state =
128                require_association_state(db, &requirement).map_err(InstallationDiffError::from)?;
129            expected_members.insert(
130                association_state.inbox_id().to_string(),
131                HashSet::from_iter(association_state.installation_ids()),
132            );
133        }
134
135        for (claimed_inbox_id, signature_key) in &welcome.members {
136            let Some(installation_ids) = expected_members.get_mut(claimed_inbox_id) else {
137                tracing::error!(
138                    claimed_inbox_id = claimed_inbox_id,
139                    "Inbox ID not found in expected members",
140                );
141                return Err(GroupError::InvalidGroupMembership);
142            };
143            if !installation_ids.contains(signature_key) {
144                tracing::error!(
145                    claimed_inbox_id = claimed_inbox_id,
146                    "Installation ID not found in expected members for inbox ID",
147                );
148                return Err(GroupError::InvalidGroupMembership);
149            }
150            installation_ids.remove(signature_key);
151        }
152        for installation_set in expected_members.values() {
153            for remaining_installation_id in installation_set {
154                if !membership
155                    .failed_installations
156                    .contains(remaining_installation_id)
157                {
158                    tracing::error!(
159                        installation_id = %remaining_installation_id.as_slice().snippet(),
160                        "Installation ID in expected members not found in ratchet tree",
161                    );
162                    return Err(GroupError::InvalidGroupMembership);
163                }
164            }
165        }
166        // TODO: Is it an error if there are 'failed installations' that are not in the expected members list?
167
168        tracing::info!("Group membership validated");
169
170        Ok(())
171    }
172}
173
174#[cfg(any(test, feature = "test-utils"))]
175pub mod test {
176    use super::*;
177
178    #[derive(Default)]
179    pub struct NoopValidator;
180
181    impl ValidateGroupMembership for NoopValidator {
182        async fn check_initial_membership(
183            &self,
184            _welcome: &WelcomeMembership,
185        ) -> Result<(), GroupError> {
186            Ok(())
187        }
188    }
189}
xmtp_mls/groups/welcomes/validated_membership.rs

xmtp_mls/groups/welcomes/
validated_membership.rs
