xmtp_mls/groups/welcomes/
validated_membership.rs

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

xmtp_mls/groups/welcomes/
validated_membership.rs
