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, MissingProofs,
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::{
13    commit::extract_group_membership_with_failures, group_membership::FailedInstallations,
14};
15
16/// Validate public trial membership, then recheck exact proofs under the writer.
17#[allow(async_fn_in_trait)]
18pub trait ValidateGroupMembership {
19    /// Resolve identity proofs and check that the tree matches the membership extension.
20    async fn check_initial_membership(&self, welcome: &WelcomeMembership)
21    -> Result<(), GroupError>;
22
23    /// Recheck exact proofs against the database read under the state writer.
24    fn check_verified_membership(
25        &self,
26        _welcome: &WelcomeMembership,
27        _db: &impl DbQuery,
28    ) -> Result<(), GroupError> {
29        Ok(())
30    }
31}
32
33/// Public membership data from a trial decode. It contains no mutable MLS state.
34#[derive(Debug, PartialEq)]
35pub struct WelcomeMembership {
36    /// Authenticated inbox membership and exact identity sequence references.
37    membership: GroupMembership,
38    failed: FailedInstallations,
39    /// Tree members as inbox IDs and installation signature keys.
40    members: Vec<(String, Vec<u8>)>,
41}
42
43impl WelcomeMembership {
44    /// Copy public data from a trial without retaining staged MLS state.
45    pub(crate) fn from_staged(welcome: &StagedWelcome) -> Result<Self, GroupError> {
46        let extensions = welcome.group_context().extensions();
47        let (membership, failed) = extract_group_membership_with_failures(extensions)
48            .map_err(|_| GroupError::InvalidWelcomeMetadata)?;
49        let members = welcome
50            .members()
51            .map(|member| {
52                let credential = BasicCredential::try_from(member.credential.clone())?;
53                Ok((
54                    parse_credential(credential.identity())?,
55                    member.signature_key,
56                ))
57            })
58            .collect::<Result<_, GroupError>>()?;
59        Ok(Self {
60            membership,
61            failed,
62            members,
63        })
64    }
65
66    fn requirements(&self) -> impl Iterator<Item = IdentityRequirement> + '_ {
67        self.membership
68            .members
69            .iter()
70            .map(|(inbox_id, sequence_id)| IdentityRequirement {
71                inbox_id: inbox_id.clone(),
72                sequence_id: *sequence_id,
73            })
74    }
75
76    /// Reject zero identity references and references at or after this Welcome.
77    // implements: JOIN-050
78    pub(crate) fn validate_sequences(&self, welcome_sequence: u64) -> Result<(), GroupError> {
79        for requirement in self.requirements() {
80            if requirement.sequence_id == 0 || requirement.sequence_id >= welcome_sequence {
81                return Err(
82                    InstallationDiffError::from(IdentityDependencyError::InvalidSequence(
83                        requirement.sequence_id,
84                    ))
85                    .into(),
86                );
87            }
88        }
89        Ok(())
90    }
91}
92
93/// Check every installation against the exact identity state named by the Welcome.
94pub struct InitialMembershipValidator<C> {
95    context: C,
96}
97
98impl<C> InitialMembershipValidator<C> {
99    pub fn new(context: C) -> InitialMembershipValidator<C> {
100        Self { context }
101    }
102}
103
104impl<C> ValidateGroupMembership for InitialMembershipValidator<C>
105where
106    C: XmtpSharedContext,
107{
108    // implements: JOIN-051
109    async fn check_initial_membership(
110        &self,
111        welcome: &WelcomeMembership,
112    ) -> Result<(), GroupError> {
113        resolve_identity_requirements(&self.context, &welcome.requirements().collect())
114            .await
115            .map_err(InstallationDiffError::from)?;
116        self.check_verified_membership(welcome, &self.context.db())
117    }
118
119    /// Name every missing proof in one `Need`, so one resolution fetches them all.
120    // implements: JOIN-052, JOIN-053
121    fn check_verified_membership(
122        &self,
123        welcome: &WelcomeMembership,
124        db: &impl DbQuery,
125    ) -> Result<(), GroupError> {
126        let mut expected_members = HashMap::<String, HashSet<Vec<u8>>>::new();
127        let mut missing = MissingProofs::default();
128        for requirement in welcome.requirements() {
129            let Some(association_state) = missing
130                .read(db, &requirement)
131                .map_err(InstallationDiffError::from)?
132            else {
133                continue;
134            };
135            expected_members.insert(
136                association_state.inbox_id().to_string(),
137                HashSet::from_iter(association_state.installation_ids()),
138            );
139        }
140        missing.check().map_err(InstallationDiffError::from)?;
141
142        for (claimed_inbox_id, signature_key) in &welcome.members {
143            let Some(installation_ids) = expected_members.get_mut(claimed_inbox_id) else {
144                tracing::error!(
145                    claimed_inbox_id = claimed_inbox_id,
146                    "Inbox ID not found in expected members",
147                );
148                return Err(GroupError::InvalidGroupMembership);
149            };
150            if !installation_ids.contains(signature_key) {
151                tracing::error!(
152                    claimed_inbox_id = claimed_inbox_id,
153                    "Installation ID not found in expected members for inbox ID",
154                );
155                return Err(GroupError::InvalidGroupMembership);
156            }
157            installation_ids.remove(signature_key);
158        }
159        for (inbox, installation_set) in &expected_members {
160            for remaining_installation_id in installation_set {
161                if !welcome
162                    .failed
163                    .of(inbox)
164                    .is_some_and(|keys| keys.contains(remaining_installation_id))
165                {
166                    tracing::error!(
167                        installation_id = %remaining_installation_id.as_slice().snippet(),
168                        "Installation ID in expected members not found in ratchet tree",
169                    );
170                    return Err(GroupError::InvalidGroupMembership);
171                }
172            }
173        }
174
175        tracing::info!("Group membership validated");
176
177        Ok(())
178    }
179}
180
181#[cfg(any(test, feature = "test-utils"))]
182pub mod test {
183    use super::*;
184
185    #[derive(Default)]
186    pub struct NoopValidator;
187
188    impl ValidateGroupMembership for NoopValidator {
189        async fn check_initial_membership(
190            &self,
191            _welcome: &WelcomeMembership,
192        ) -> Result<(), GroupError> {
193            Ok(())
194        }
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    use crate::tester;
202    use std::collections::BTreeMap;
203    use xmtp_db::identity_update::QueryIdentityUpdates;
204    use xmtp_mls_group_state::{
205        app_data::migration::encode_group_membership_dict, inbox_id::InboxId,
206    };
207    use xmtp_mls_validation::commit::membership_entries;
208    use xmtp_proto::xmtp::mls::message_contents::{
209        GroupMembershipEntry,
210        group_membership_entry::{V1, Version},
211    };
212
213    // verifies: JOIN-053
214    #[xmtp_common::test(unwrap_try = true)]
215    async fn misattributed_failed_installation_is_rejected() {
216        tester!(alix, disable_workers);
217        tester!(bo, disable_workers);
218        tester!(caro, disable_workers);
219        tester!(bo2, from: bo, disable_workers);
220        for client in [&*alix, &*bo, &*caro] {
221            client.inbox_state(true).await?;
222        }
223        let sequence = |client: &crate::utils::FullXmtpClient| -> u64 {
224            client
225                .context
226                .db()
227                .get_latest_sequence_id_for_inbox(client.inbox_id())
228                .unwrap() as u64
229        };
230        let absent_id = bo2.context.installation_id().to_vec();
231        let members = [&*alix, &*bo, &*caro]
232            .into_iter()
233            .map(|client| {
234                (
235                    client.inbox_id().to_string(),
236                    client.context.installation_id().to_vec(),
237                )
238            })
239            .collect::<Vec<_>>();
240        let validator = InitialMembershipValidator::new(caro.context.clone());
241        for owner in [alix.inbox_id(), bo.inbox_id()] {
242            let entries = [&*alix, &*bo, &*caro]
243                .into_iter()
244                .map(|client| {
245                    (
246                        InboxId::from_hex(client.inbox_id()).unwrap(),
247                        GroupMembershipEntry {
248                            version: Some(Version::V1(V1 {
249                                sequence_id: sequence(client),
250                                failed_installations: if client.inbox_id() == owner {
251                                    vec![absent_id.clone()]
252                                } else {
253                                    vec![]
254                                },
255                            })),
256                        },
257                    )
258                })
259                .collect::<BTreeMap<_, _>>();
260            let (membership, failed) =
261                membership_entries(&encode_group_membership_dict(&entries)?)?;
262            let welcome = WelcomeMembership {
263                membership,
264                failed,
265                members: members.clone(),
266            };
267            let result = validator.check_initial_membership(&welcome).await;
268            if owner == bo.inbox_id() {
269                result?;
270            } else {
271                assert!(matches!(result, Err(GroupError::InvalidGroupMembership)));
272            }
273        }
274    }
275}
xmtp_mls/groups/welcomes/validated_membership.rs

xmtp_mls/groups/welcomes/
validated_membership.rs
