xmtp_mls/groups/mls_sync/
post_commit.rs

1//! Post-commit work: installations, welcomes, and HMAC keys.
2
3use super::*;
4use xmtp_common::snippet::Snippet;
5
6fn log_missing_installations(
7    inbox_id: &str,
8    installation_id: &[u8],
9    intent_data: &UpdateGroupMembershipIntentData,
10) {
11    debug!(
12        inbox_id,
13        installation_id = installation_id.snippet(),
14        membership_update_count = intent_data.membership_updates.len(),
15        removed_member_count = intent_data.removed_members.len(),
16        failed_installation_count = intent_data.failed_installations.len(),
17        "Adding missing installations"
18    );
19}
20
21impl<Context> MlsGroup<Context>
22where
23    Context: XmtpSharedContext,
24{
25    #[tracing::instrument(skip_all)]
26    pub(crate) async fn post_commit(&self) -> Result<(), GroupError> {
27        self.publish_required_welcomes().await
28    }
29
30    pub async fn maybe_update_installations(
31        &self,
32        update_interval_ns: Option<i64>,
33    ) -> Result<(), GroupError> {
34        let db = self.context.db();
35        let Some(stored_group) = db.find_group(&self.group_id)? else {
36            return Err(GroupError::NotFound(NotFound::GroupById(self.group_id)));
37        };
38        if stored_group.conversation_type.is_virtual() {
39            return Ok(());
40        }
41
42        // determine how long of an interval in time to use before updating list
43        let interval_ns = update_interval_ns.unwrap_or(SYNC_UPDATE_INSTALLATIONS_INTERVAL_NS);
44
45        let now_ns = xmtp_common::time::now_ns();
46        let last_ns = db.get_installations_time_checked(&self.group_id)?;
47        let elapsed_ns = now_ns - last_ns;
48        if elapsed_ns > interval_ns && self.is_active()? {
49            self.add_missing_installations().await?;
50            db.update_installations_time_checked(&self.group_id)?;
51        }
52
53        Ok(())
54    }
55
56    /**
57     * Checks each member of the group for `IdentityUpdates` after their current sequence_id. If updates
58     * are found the method will construct an [`UpdateGroupMembershipIntentData`] and create a change
59     * to the [`GroupMembership`] that will add any missing installations.
60     *
61     * This is designed to handle cases where existing members have added a new installation to their inbox or revoked an installation
62     * and the group has not been updated to include it.
63     */
64    #[cfg_attr(any(test, feature = "test-utils"), tracing::instrument(level = "info", fields(inbox_id = %self.context.inbox_id()), skip_all))]
65    #[cfg_attr(
66        not(any(test, feature = "test-utils")),
67        tracing::instrument(level = "trace", skip_all)
68    )]
69    pub(crate) async fn add_missing_installations(&self) -> Result<(), GroupError> {
70        let intent_data = self.get_membership_update_intent(&[], &[]).await?;
71
72        // If there is nothing to do, stop here
73        if intent_data.is_empty() {
74            return Ok(());
75        }
76
77        log_missing_installations(
78            self.context.inbox_id(),
79            self.context.installation_id().as_slice(),
80            &intent_data,
81        );
82
83        let intent = QueueIntent::update_group_membership()
84            .data(intent_data)
85            .queue(self)?;
86
87        let _ = self.sync_until_intent_resolved(intent.id).await?;
88        Ok(())
89    }
90
91    #[tracing::instrument(level = "trace", skip_all)]
92    /**
93     * get_membership_update_intent will query the network for any new [`IdentityUpdate`]s for any of the existing
94     * group members
95     *
96     * Callers may also include a list of added or removed inboxes
97     */
98    pub(crate) async fn get_membership_update_intent(
99        &self,
100        inbox_ids_to_add: &[InboxIdRef<'_>],
101        inbox_ids_to_remove: &[InboxIdRef<'_>],
102    ) -> Result<UpdateGroupMembershipIntentData, GroupError> {
103        let existing_group_membership = self.with_group_snapshot(|group| {
104            extract_group_membership(group.extensions()).map_err(Into::into)
105        })?;
106        {
107            // TODO:nm prevent querying for updates on members who are being removed
108            let mut inbox_ids = existing_group_membership.inbox_ids();
109            inbox_ids.extend_from_slice(inbox_ids_to_add);
110            let conn = self.context.db();
111            // Load any missing updates from the network
112            load_identity_updates_for_client(&self.context, &conn, &inbox_ids).await?;
113
114            let latest_sequence_id_map = conn.get_latest_sequence_id(&inbox_ids as &[&str])?;
115
116            // Get a list of all inbox IDs that have increased sequence_id for the group
117            let changed_inbox_ids =
118                inbox_ids
119                    .iter()
120                    .try_fold(HashMap::new(), |mut updates, inbox_id| {
121                        match (
122                            latest_sequence_id_map.get(inbox_id as &str),
123                            existing_group_membership.get(inbox_id),
124                        ) {
125                            // This is an update. We have a new sequence ID and an existing one
126                            (Some(latest_sequence_id), Some(current_sequence_id)) => {
127                                let latest_sequence_id_u64 = *latest_sequence_id as u64;
128                                if latest_sequence_id_u64.gt(current_sequence_id) {
129                                    updates.insert(inbox_id.to_string(), latest_sequence_id_u64);
130                                }
131                            }
132                            // This is for new additions to the group
133                            (Some(latest_sequence_id), None) => {
134                                // This is the case for net new members to the group
135                                updates.insert(inbox_id.to_string(), *latest_sequence_id as u64);
136                            }
137                            (_, _) => {
138                                tracing::warn!(
139                                    "Could not find existing sequence ID for inbox {}",
140                                    inbox_id
141                                );
142                                return Err(GroupError::MissingSequenceId);
143                            }
144                        }
145
146                        Ok(updates)
147                    })?;
148            let old_group_membership = existing_group_membership.clone();
149            let mut new_membership = old_group_membership.clone();
150            for (inbox_id, sequence_id) in changed_inbox_ids.iter() {
151                new_membership.add(inbox_id.clone(), *sequence_id);
152            }
153            for inbox_id in inbox_ids_to_remove {
154                new_membership.remove(inbox_id);
155            }
156
157            let changes_with_kps = calculate_membership_changes_with_keypackages(
158                &self.context,
159                &self.group_id,
160                &new_membership,
161                &old_group_membership,
162            )
163            .await?;
164
165            // If we fail to fetch or verify all the added members' KeyPackage, return an error.
166            // skip if the inbox ids is 0 from the beginning
167            if !inbox_ids_to_add.is_empty()
168                && !changes_with_kps.failed_installations.is_empty()
169                && changes_with_kps.new_installations.is_empty()
170            {
171                return Err(GroupError::FailedToVerifyInstallations(
172                    FailedInstallationIds(changes_with_kps.failed_installations.clone()),
173                ));
174            }
175
176            Ok(UpdateGroupMembershipIntentData::new(
177                changed_inbox_ids,
178                inbox_ids_to_remove
179                    .iter()
180                    .map(|s| s.to_string())
181                    .collect::<Vec<String>>(),
182                changes_with_kps.failed_installations,
183            ))
184        }
185    }
186
187    #[cfg(test)]
188    pub(in crate::groups) async fn send_welcomes(
189        &self,
190        action: SendWelcomesAction,
191        message_cursor: Option<i64>,
192    ) -> Result<(), GroupError> {
193        let message_cursor = u64::try_from(message_cursor.unwrap_or(0))
194            .map_err(|_| xmtp_proto::ConversionError::Unspecified("negative Welcome cursor"))?;
195        let units = crate::state_tx::state_write(self.context.mls_storage(), |_tx| {
196            self.prepare_welcome_envelopes(action, message_cursor)?
197                .into_iter()
198                .map(PublishUnit::single)
199                .collect::<Result<Vec<_>, _>>()
200                .map(Continue)
201                .map_err(GroupError::from)
202        })?
203        .into_continued();
204        self.context.api().publish_units(units).await?;
205        Ok(())
206    }
207
208    /// Provides hmac keys for a range of epochs around current epoch
209    /// `group.hmac_keys(-1..=1)`` will provide 3 keys consisting of last epoch, current epoch, and next epoch
210    /// `group.hmac_keys(0..=0) will provide 1 key, consisting of only the current epoch
211    #[tracing::instrument(level = "trace", skip_all)]
212    pub fn hmac_keys(
213        &self,
214        epoch_delta_range: RangeInclusive<i64>,
215    ) -> Result<Vec<HmacKey>, StorageError> {
216        crate::state_tx::state_write(self.context.mls_storage(), |tx| {
217            self.hmac_keys_in(tx.storage().db(), epoch_delta_range)
218                .map(Continue)
219        })
220        .map(TransactionOutcome::into_continued)
221    }
222
223    pub(crate) fn hmac_keys_in(
224        &self,
225        conn: impl xmtp_db::DbQuery,
226        epoch_delta_range: RangeInclusive<i64>,
227    ) -> Result<Vec<HmacKey>, StorageError> {
228        let preferences = StoredUserPreferences::load(&conn)?;
229        let mut ikm = match preferences.hmac_key {
230            Some(ikm) => ikm,
231            None => {
232                let key = HmacKey::random_key();
233                StoredUserPreferences::store_hmac_key(&conn, &key, None)?;
234                key
235            }
236        };
237        ikm.extend_from_slice(self.group_id.as_ref());
238        let hkdf = Hkdf::<Sha256>::new(Some(HMAC_SALT), &ikm);
239
240        let mut result = vec![];
241        let current_epoch = hmac_epoch();
242        for delta in epoch_delta_range {
243            let epoch = current_epoch + delta;
244
245            let mut info = self.group_id.to_vec();
246            info.extend(&epoch.to_le_bytes());
247
248            let mut key = [0; 42];
249            hkdf.expand(&info, &mut key).expect("Length is correct");
250
251            result.push(HmacKey { key, epoch });
252        }
253
254        Ok(result)
255    }
256
257    #[cfg(test)]
258    #[tracing::instrument(level = "trace", skip_all)]
259    pub(in crate::groups) fn prepare_group_messages(
260        &self,
261        payloads: Vec<(&[u8], bool)>,
262    ) -> Result<Vec<PublishUnit>, GroupError> {
263        crate::state_tx::state_write(self.context.mls_storage(), |tx| {
264            let envelopes = self.prepare_group_envelopes_in(tx.storage().db(), payloads)?;
265            Ok::<_, GroupError>(Continue(vec![PublishUnit::new(envelopes)?]))
266        })
267        .map(TransactionOutcome::into_continued)
268    }
269
270    pub(super) fn prepare_group_envelopes_in(
271        &self,
272        conn: impl xmtp_db::DbQuery,
273        payloads: Vec<(&[u8], bool)>,
274    ) -> Result<Vec<ClientEnvelope>, GroupError> {
275        let hmac_key = self
276            .hmac_keys_in(conn, 0..=0)?
277            .pop()
278            .expect("Range of count 1 was provided.");
279        let sender_hmac =
280            Hmac::<Sha256>::new_from_slice(&hmac_key.key).expect("HMAC can take key of any size");
281
282        let mut result = vec![];
283        for (payload, should_push) in payloads {
284            let mut sender_hmac = sender_hmac.clone();
285            sender_hmac.update(payload);
286            let sender_hmac = sender_hmac.finalize();
287
288            result.push(ClientEnvelope {
289                payload: Some(Payload::GroupMessage(BackendGroupMessage {
290                    data: payload.to_vec(),
291                    sender_hmac: sender_hmac.into_bytes().to_vec(),
292                    should_push,
293                })),
294            });
295        }
296
297        Ok(result)
298    }
299}
300
301#[cfg(all(test, not(target_arch = "wasm32")))]
302mod tests {
303    use super::*;
304    use xmtp_logging::{Level, test_logging::LogCapture};
305
306    #[xmtp_common::test(unwrap_try = true)]
307    fn missing_installation_log_omits_full_ids() {
308        let failed_id = (32_u8..64).collect::<Vec<_>>();
309        let local_id = (64_u8..96).collect::<Vec<_>>();
310        let intent = UpdateGroupMembershipIntentData::new(
311            HashMap::from([("updated-inbox".to_string(), 7)]),
312            vec!["removed-inbox".to_string()],
313            vec![failed_id.clone()],
314        );
315        let encoded: Vec<u8> = intent.into();
316        let restored = UpdateGroupMembershipIntentData::try_from(encoded)?;
317        assert_eq!(restored.failed_installations, vec![failed_id.clone()]);
318        let capture = LogCapture::new(Level::Debug);
319        tracing::dispatcher::with_default(&capture.dispatch(), || {
320            log_missing_installations("local-inbox", &local_id, &restored);
321        });
322        let output = capture.output();
323        assert!(!output.contains(&format!("{failed_id:?}")));
324        assert!(!output.contains(&hex::encode(&failed_id)));
325        assert!(!output.contains(&hex::encode(&local_id)));
326        let record: serde_json::Value = serde_json::from_str(output.trim())?;
327        assert_eq!(record["message"], "Adding missing installations");
328        assert_eq!(record["membership_update_count"], 1);
329        assert_eq!(record["removed_member_count"], 1);
330        assert_eq!(record["failed_installation_count"], 1);
331    }
332}
xmtp_mls/groups/mls_sync/post_commit.rs

xmtp_mls/groups/mls_sync/
post_commit.rs
