xmtp_mls/groups/mls_sync/

post_commit.rs

//! Post-commit work: installations, welcomes, and HMAC keys.

use super::publish::OutgoingPreparationError;
use super::*;
use xmtp_common::snippet::Snippet;
use xmtp_mls_group_keys::group_keys::{GroupKeys, SealError, sign_group_message};

fn log_missing_installations(
    inbox_id: &str,
    installation_id: &[u8],
    intent_data: &UpdateGroupMembershipIntentData,
) {
    debug!(
        inbox_id,
        installation_id = installation_id.snippet(),
        membership_update_count = intent_data.membership_updates.len(),
        removed_member_count = intent_data.removed_members.len(),
        failed_installation_count = intent_data.failed_installations.len(),
        "Adding missing installations"
    );
}

pub(crate) async fn membership_update_intent(
    context: &impl XmtpSharedContext,
    group_id: &GroupId,
    existing_group_membership: &GroupMembership,
    retry_leaves: Option<&[Vec<u8>]>,
    inbox_ids_to_add: &[InboxIdRef<'_>],
    inbox_ids_to_remove: &[InboxIdRef<'_>],
) -> Result<UpdateGroupMembershipIntentData, GroupError> {
    // TODO:nm prevent querying for updates on members who are being removed
    let mut inbox_ids = existing_group_membership.inbox_ids();
    inbox_ids.extend_from_slice(inbox_ids_to_add);
    let conn = context.db();
    // Load any missing updates from the network
    load_identity_updates_for_client(context, &conn, &inbox_ids).await?;

    let latest_sequence_id_map = conn.get_latest_sequence_id(&inbox_ids as &[&str])?;

    // Get a list of all inbox IDs that have increased sequence_id for the group
    let mut changed_inbox_ids =
        inbox_ids
            .iter()
            .try_fold(HashMap::new(), |mut updates, inbox_id| {
                match (
                    latest_sequence_id_map.get(inbox_id as &str),
                    existing_group_membership.get(inbox_id),
                ) {
                    // This is an update. We have a new sequence ID and an existing one
                    (Some(latest_sequence_id), Some(current_sequence_id)) => {
                        let latest_sequence_id_u64 = *latest_sequence_id as u64;
                        if latest_sequence_id_u64.gt(current_sequence_id) {
                            updates.insert(inbox_id.to_string(), latest_sequence_id_u64);
                        }
                    }
                    // This is for new additions to the group
                    (Some(latest_sequence_id), None) => {
                        // This is the case for net new members to the group
                        updates.insert(inbox_id.to_string(), *latest_sequence_id as u64);
                    }
                    (_, _) => {
                        tracing::warn!(
                            "Could not find existing sequence ID for inbox {}",
                            inbox_id
                        );
                        return Err(GroupError::MissingSequenceId);
                    }
                }

                Ok(updates)
            })?;
    let old_group_membership = existing_group_membership.clone();
    let mut new_membership = old_group_membership.clone();
    for (inbox_id, sequence_id) in changed_inbox_ids.iter() {
        new_membership.add(inbox_id.clone(), *sequence_id);
    }
    for inbox_id in inbox_ids_to_remove {
        new_membership.remove(inbox_id);
    }

    let changes_with_kps = calculate_membership_changes_with_keypackages(
        context,
        group_id,
        &new_membership,
        &old_group_membership,
        retry_leaves,
    )
    .await?;

    // A retry alone changes no entry: refresh the own entry at its current
    // sequence so the intent is not empty. Apply keeps the higher sequence.
    if changed_inbox_ids.is_empty()
        && !changes_with_kps.new_installations.is_empty()
        && let Some(&sequence_id) = existing_group_membership.get(context.inbox_id())
    {
        changed_inbox_ids.insert(context.inbox_id().to_string(), sequence_id);
    }

    // If we fail to fetch or verify all the added members' KeyPackage, return an error.
    // skip if the inbox ids is 0 from the beginning
    if !inbox_ids_to_add.is_empty()
        && !changes_with_kps.failed_installations.is_empty()
        && changes_with_kps.new_installations.is_empty()
    {
        return Err(GroupError::FailedToVerifyInstallations(
            FailedInstallationIds(changes_with_kps.failed_installations.clone()),
        ));
    }

    // Only an inbox the app asked to add may be inserted. Every other
    // entry refreshes an existing member.
    let refreshed_members = changed_inbox_ids
        .keys()
        .filter(|inbox_id| {
            existing_group_membership.get(inbox_id.as_str()).is_some()
                && !inbox_ids_to_add.contains(&inbox_id.as_str())
        })
        .cloned()
        .collect();
    Ok(UpdateGroupMembershipIntentData::new(
        changed_inbox_ids,
        inbox_ids_to_remove
            .iter()
            .map(|s| s.to_string())
            .collect::<Vec<String>>(),
        changes_with_kps.failed_installations,
    )
    .with_refreshed_members(refreshed_members))
}

impl<Context> MlsGroup<Context>
where
    Context: XmtpSharedContext,
{
    #[tracing::instrument(skip_all)]
    pub(crate) async fn post_commit(&self) -> Result<(), GroupError> {
        self.publish_required_welcomes().await
    }

    pub async fn maybe_update_installations(
        &self,
        update_interval_ns: Option<i64>,
    ) -> Result<(), GroupError> {
        let db = self.context.db();
        let Some(stored_group) = db.find_group(&self.group_id)? else {
            return Err(GroupError::NotFound(NotFound::GroupById(self.group_id)));
        };
        if stored_group.conversation_type.is_virtual() {
            return Ok(());
        }

        // determine how long of an interval in time to use before updating list
        let interval_ns = update_interval_ns.unwrap_or(SYNC_UPDATE_INSTALLATIONS_INTERVAL_NS);

        let now_ns = xmtp_common::time::now_ns();
        let last_ns = db.get_installations_time_checked(&self.group_id)?;
        let elapsed_ns = now_ns - last_ns;
        if elapsed_ns > interval_ns && self.is_active()? {
            self.add_missing_installations().await?;
            db.update_installations_time_checked(&self.group_id)?;
        }

        Ok(())
    }

    /**
     * Checks each member of the group for `IdentityUpdates` after their current sequence_id. If updates
     * are found the method will construct an [`UpdateGroupMembershipIntentData`] and create a change
     * to the [`GroupMembership`] that will add any missing installations.
     *
     * This is designed to handle cases where existing members have added a new installation to their inbox or revoked an installation
     * and the group has not been updated to include it.
     */
    #[cfg_attr(any(test, feature = "test-utils"), tracing::instrument(level = "info", fields(inbox_id = %self.context.inbox_id()), skip_all))]
    #[cfg_attr(
        not(any(test, feature = "test-utils")),
        tracing::instrument(level = "trace", skip_all)
    )]
    pub(crate) async fn add_missing_installations(&self) -> Result<(), GroupError> {
        let intent_data = self.get_membership_update_intent(&[], &[]).await?;

        // If there is nothing to do, stop here
        if intent_data.is_empty() {
            return Ok(());
        }

        log_missing_installations(
            self.context.inbox_id(),
            self.context.installation_id().as_slice(),
            &intent_data,
        );

        let intent = QueueIntent::update_group_membership()
            .data(intent_data)
            .queue(self)?;

        let _ = self.sync_until_intent_resolved(intent.id).await?;
        Ok(())
    }

    /// Refresh membership, removing revoked installations but adding none.
    // implements: SYNC-034
    pub(crate) async fn remove_revoked_installations(&self) -> Result<(), GroupError> {
        let mut intent_data = self.get_membership_update_intent(&[], &[]).await?;
        if intent_data.is_empty() {
            return Ok(());
        }
        intent_data.removal_only = true;
        let intent = QueueIntent::update_group_membership()
            .data(intent_data)
            .queue(self)?;
        self.sync_until_intent_resolved(intent.id).await?;
        Ok(())
    }

    #[tracing::instrument(level = "trace", skip_all)]
    /**
     * get_membership_update_intent will query the network for any new [`IdentityUpdate`]s for any of the existing
     * group members
     *
     * Callers may also include a list of added or removed inboxes
     */
    pub(crate) async fn get_membership_update_intent(
        &self,
        inbox_ids_to_add: &[InboxIdRef<'_>],
        inbox_ids_to_remove: &[InboxIdRef<'_>],
    ) -> Result<UpdateGroupMembershipIntentData, GroupError> {
        let (existing_group_membership, retry_leaves) = self.with_group_snapshot(|group| {
            Ok::<_, GroupError>((
                extract_group_membership(group.extensions())?,
                retry_leaves(group.extensions(), group.members()),
            ))
        })?;
        membership_update_intent(
            &self.context,
            &self.group_id,
            &existing_group_membership,
            retry_leaves.as_deref(),
            inbox_ids_to_add,
            inbox_ids_to_remove,
        )
        .await
    }

    pub(crate) async fn initial_add(
        context: &Context,
        draft: &(
            super::super::lifecycle::NewGroupId,
            openmls::group::MlsGroupCreateConfig,
        ),
        ids: &[InboxIdRef<'_>],
    ) -> Result<Option<UpdateGroupMembershipIntentData>, GroupError> {
        use super::super::lifecycle::NewGroupId;
        let id = match draft.0 {
            NewGroupId::Keyed { group_id, .. } | NewGroupId::Placeholder(group_id) => group_id,
        };
        // The prepared dictionary supplies the stored group's membership and
        // floor; its only leaf is this installation.
        let extensions = draft.1.group_context_extensions();
        let existing = extract_group_membership(extensions)?;
        let leaves = retry_leaves(extensions, std::iter::empty())
            .map(|_| vec![context.installation_id().to_vec()]);
        let add =
            membership_update_intent(context, &id, &existing, leaves.as_deref(), ids, &[]).await?;
        if add.is_empty() {
            return Ok(None);
        }
        let max_members = context
            .server_configuration()
            .configuration()
            .mls
            .max_group_members;
        let resulting: HashSet<&str> = existing
            .inbox_ids()
            .into_iter()
            .chain(add.membership_updates.keys().map(String::as_str))
            .collect();
        if resulting.len() > max_members {
            return Err(GroupError::UserLimitExceeded);
        }
        Ok(Some(add))
    }

    pub(crate) fn intent_pending(&self, id: i32) -> Result<bool, GroupError> {
        if self.membership_child_follow_up_failed(id)? {
            return Ok(false);
        }
        let intent: Option<StoredGroupIntent> = self.context.db().fetch(&id)?;
        match intent.map(|intent| intent.state) {
            Some(IntentState::ToPublish | IntentState::Published | IntentState::Committed) => {
                Ok(true)
            }
            Some(IntentState::Processed) => Ok(!self.pending_membership_children(id)?.is_empty()),
            _ => Ok(false),
        }
    }

    #[cfg(test)]
    pub(in crate::groups) async fn send_welcomes(
        &self,
        action: SendWelcomesAction,
        message_cursor: Option<i64>,
    ) -> Result<(), GroupError> {
        let message_cursor = u64::try_from(message_cursor.unwrap_or(0))
            .map_err(|_| xmtp_proto::ConversionError::Unspecified("negative Welcome cursor"))?;
        let units = crate::state_tx::state_write(self.context.mls_storage(), |_tx| {
            self.prepare_welcome_envelopes(action, message_cursor)?
                .into_iter()
                .map(PublishUnit::single)
                .collect::<Result<Vec<_>, _>>()
                .map(Continue)
                .map_err(GroupError::from)
        })?
        .into_continued();
        self.context.api().publish_units(units).await?;
        Ok(())
    }

    /// Provides hmac keys for a range of epochs around current epoch
    /// `group.hmac_keys(-1..=1)`` will provide 3 keys consisting of last epoch, current epoch, and next epoch
    /// `group.hmac_keys(0..=0) will provide 1 key, consisting of only the current epoch
    #[tracing::instrument(level = "trace", skip_all)]
    pub fn hmac_keys(
        &self,
        epoch_delta_range: RangeInclusive<i64>,
    ) -> Result<Zeroizing<Vec<HmacKey>>, StorageError> {
        crate::state_tx::state_write_with_events(
            self.context.mls_storage(),
            self.context.events(),
            |tx, events| {
                self.hmac_keys_in(tx.storage().db(), events, epoch_delta_range)
                    .map(Continue)
            },
        )
        .map(TransactionOutcome::into_continued)
    }

    /// Derive the group's HMAC key for each epoch in `epoch_delta_range`,
    /// relative to the current HMAC epoch. Generates and stores a root key
    /// when none is stored, and reports it to `events`, which the caller
    /// emits once its transaction commits.
    ///
    /// The vector is wiped on drop, including slots that `pop()` vacated.
    // implements: PUSH-258
    pub(crate) fn hmac_keys_in(
        &self,
        conn: impl xmtp_db::DbQuery,
        events: &impl xmtp_events::EventWriter<crate::subscriptions::internal::InternalEvent>,
        epoch_delta_range: RangeInclusive<i64>,
    ) -> Result<Zeroizing<Vec<HmacKey>>, StorageError> {
        // implements: SYNC-026
        let (root_key, made) = ensure_root_key(&conn)?;
        if made {
            crate::subscriptions::internal::emit_root_key_changed(events, &conn, true)?;
        }
        let root_key = root_key.as_slice();
        let group_id = self.group_id.as_ref();
        // Sized up front: growing would free an unwiped partial copy.
        let mut ikm = Zeroizing::new(Vec::with_capacity(root_key.len() + group_id.len()));
        ikm.extend_from_slice(root_key);
        ikm.extend_from_slice(group_id);
        let hkdf = Hkdf::<Sha256>::new(Some(HMAC_SALT), &ikm);

        let current_epoch = hmac_epoch();
        let mut result = Zeroizing::new(Vec::with_capacity(epoch_delta_range.size_hint().0));
        for delta in epoch_delta_range {
            let epoch = current_epoch + delta;

            let mut info = group_id.to_vec();
            info.extend(&epoch.to_le_bytes());

            let mut key = HmacKey {
                key: [0; 42],
                epoch,
            };
            hkdf.expand(&info, &mut key.key).expect("Length is correct");

            result.push(key);
        }

        Ok(result)
    }

    #[cfg(test)]
    #[tracing::instrument(level = "trace", skip_all)]
    pub(in crate::groups) fn prepare_group_messages(
        &self,
        payloads: Vec<(&[u8], bool)>,
    ) -> Result<Vec<PublishUnit>, GroupError> {
        crate::state_tx::state_write_with_events(
            self.context.mls_storage(),
            self.context.events(),
            |tx, events| {
                tx.with_group(self.group_id, |group, storage| {
                    let envelopes =
                        self.prepare_group_envelopes_in(storage.db(), events, group, payloads)?;
                    Ok::<_, GroupError>(Continue(vec![PublishUnit::new(envelopes)?]))
                })
            },
        )
        .map(TransactionOutcome::into_continued)
    }

    /// One signed group message envelope per payload, for `group` as the
    /// caller's writer loaded it. The keys come from the group's own
    /// `GROUP_ROOT_KEY`, and the root must derive the group's id, so a group
    /// without a root (a Restored placeholder or a legacy group) fails here,
    /// before anything is saved or sent.
    pub(super) fn prepare_group_envelopes_in(
        &self,
        conn: impl xmtp_db::DbQuery,
        events: &impl xmtp_events::EventWriter<crate::subscriptions::internal::InternalEvent>,
        group: &OpenMlsGroup,
        payloads: Vec<(&[u8], bool)>,
    ) -> Result<Vec<ClientEnvelope>, GroupError> {
        let keys = xmtp_mls_group_state::app_data::group_root::check_group_root(
            group.extensions(),
            group.group_id().as_slice(),
        )
        .map_err(OutgoingPreparationError::GroupKeys)?;
        let hmac_keys = self.hmac_keys_in(conn, events, 0..=0)?;
        let hmac_key = hmac_keys.last().expect("Range of count 1 was provided.");
        let signed = group_message_envelopes(hmac_key, &keys, payloads);
        #[cfg(test)]
        let signed = sign_fault::apply(group.group_id().as_slice(), signed);
        signed.map_err(|error| OutgoingPreparationError::SignGroupMessage(error).into())
    }
}

/// Test-only faults at the result of signing a group's messages, by group id.
/// They let tests drive a failed self-check, or an envelope the local
/// admission parse refuses, through preparation.
#[cfg(test)]
pub(crate) mod sign_fault {
    use super::*;
    use std::collections::HashMap;
    use std::sync::{Arc, LazyLock, Mutex};

    pub(crate) type SignFault = Arc<
        dyn Fn(Result<Vec<ClientEnvelope>, SealError>) -> Result<Vec<ClientEnvelope>, SealError>
            + Send
            + Sync,
    >;

    static FAULTS: LazyLock<Mutex<HashMap<Vec<u8>, SignFault>>> = LazyLock::new(Default::default);

    /// Pass every later signing result of `group_id` through `fault`.
    pub(crate) fn set(group_id: &[u8], fault: SignFault) {
        FAULTS
            .lock()
            .expect("sign faults lock")
            .insert(group_id.to_vec(), fault);
    }

    pub(super) fn apply(
        group_id: &[u8],
        signed: Result<Vec<ClientEnvelope>, SealError>,
    ) -> Result<Vec<ClientEnvelope>, SealError> {
        let fault = FAULTS
            .lock()
            .expect("sign faults lock")
            .get(group_id)
            .cloned();
        match fault {
            Some(fault) => fault(signed),
            None => signed,
        }
    }
}

/// One group message envelope per payload, whose `sender_hmac` is the
/// HMAC-SHA256 of the payload under `hmac_key` and whose `signature` is the
/// group signing key's signature over every other field. Signing runs after
/// `sender_hmac` is set and passes the self-check before an envelope is
/// returned, so the bytes a caller saves are the bytes the backend verifies.
// implements: PUSH-258, GMOD-041
pub(super) fn group_message_envelopes(
    hmac_key: &HmacKey,
    keys: &GroupKeys,
    payloads: Vec<(&[u8], bool)>,
) -> Result<Vec<ClientEnvelope>, SealError> {
    let sender_hmac =
        Hmac::<Sha256>::new_from_slice(&hmac_key.key).expect("HMAC can take key of any size");

    let mut result = vec![];
    for (payload, should_push) in payloads {
        let mut sender_hmac = sender_hmac.clone();
        sender_hmac.update(payload);
        let sender_hmac = sender_hmac.finalize();

        let mut message = BackendGroupMessage {
            data: payload.to_vec(),
            sender_hmac: sender_hmac.into_bytes().to_vec(),
            should_push,
            signature: Vec::new(),
        };
        sign_group_message(keys, &mut message)?;
        result.push(ClientEnvelope {
            payload: Some(Payload::GroupMessage(message)),
        });
    }
    Ok(result)
}

#[cfg(all(test, not(target_arch = "wasm32")))]
mod tests {
    use super::*;
    use xmtp_logging::{Level, test_logging::LogCapture};

    #[xmtp_common::test(unwrap_try = true)]
    fn missing_installation_log_omits_full_ids() {
        let failed_id = (32_u8..64).collect::<Vec<_>>();
        let local_id = (64_u8..96).collect::<Vec<_>>();
        let intent = UpdateGroupMembershipIntentData::new(
            HashMap::from([("updated-inbox".to_string(), 7)]),
            vec!["removed-inbox".to_string()],
            vec![failed_id.clone()],
        );
        let encoded: Vec<u8> = intent.into();
        let restored = UpdateGroupMembershipIntentData::try_from(encoded)?;
        assert_eq!(restored.failed_installations, vec![failed_id.clone()]);
        let capture = LogCapture::new(Level::Debug);
        tracing::dispatcher::with_default(&capture.dispatch(), || {
            log_missing_installations("local-inbox", &local_id, &restored);
        });
        let output = capture.output();
        assert!(!output.contains(&format!("{failed_id:?}")));
        assert!(!output.contains(&hex::encode(&failed_id)));
        assert!(!output.contains(&hex::encode(&local_id)));
        let record: serde_json::Value = serde_json::from_str(output.trim())?;
        assert_eq!(record["message"], "Adding missing installations");
        assert_eq!(record["membership_update_count"], 1);
        assert_eq!(record["removed_member_count"], 1);
        assert_eq!(record["failed_installation_count"], 1);
    }
}