feat(xmtp_mls): same-commit registry lookahead — register-then-write in one commit

crates/xmtp_mls/src/groups/app_data/mod.rs

@@ -41,6 +41,109 @@ use self::component_source::{
     ComponentSourceError, apply_app_data_update_payload, read_from_app_data_dict,
 };
 
+/// Same-commit registry lookahead: tracks the `COMPONENT_REGISTRY`
+/// snapshot as a commit's `AppDataUpdate` proposals are walked in queue
+/// order, so a component *registered* by an earlier proposal in the
+/// same commit dispatches and policy-checks exactly like one registered
+/// in an earlier commit. Without this, register-then-write batched into
+/// one commit is rejected by every receiver (the write's type dispatch
+/// and deny-by-default policy lookup would only see the pre-commit
+/// snapshot) — a permanent trap enforced by nothing on the send side.
+///
+/// Both receive stages fold through this one type so their semantics
+/// cannot drift:
+/// - the apply path ([`accumulate_app_data_updates`]) folds
+///   optimistically — its results are discarded unless commit
+///   validation later passes;
+/// - the validate path
+///   (`validate_app_data_update_proposals_in_commit`) folds only after
+///   the registry proposal itself has passed its hardcoded
+///   super-admin-only policy check, so the overlaid entries carry the
+///   same trust as ones committed earlier.
+///
+/// Queue order is fixed by the commit, so every honest receiver folds
+/// identically; write-before-register in queue order deterministically
+/// still fails.
+pub(crate) struct RegistryOverlay {
+    registry: ComponentRegistry,
+    snapshot_bytes: Option<Vec<u8>>,
+}
+
+impl RegistryOverlay {
+    /// Start from the pre-commit dict state of `mls_group`.
+    pub(crate) fn from_group(mls_group: &OpenMlsGroup) -> Result<Self, ComponentSourceError> {
+        Ok(Self {
+            registry: load_component_registry(mls_group)?,
+            snapshot_bytes: read_from_app_data_dict(ComponentId::COMPONENT_REGISTRY, mls_group),
+        })
+    }
+
+    /// Start from an already-loaded registry (the validator pre-loads
+    /// one per commit) plus the matching pre-commit dict bytes.
+    pub(crate) fn new(registry: ComponentRegistry, snapshot_bytes: Option<Vec<u8>>) -> Self {
+        Self {
+            registry,
+            snapshot_bytes,
+        }
+    }
+
+    /// The registry as of the last folded proposal (initially the
+    /// pre-commit snapshot). Hand this to type dispatch and policy
+    /// lookups for each subsequent proposal.
+    pub(crate) fn registry(&self) -> &ComponentRegistry {
+        &self.registry
+    }
+
+    /// Fold one `COMPONENT_REGISTRY` `AppDataUpdate` operation. Callers
+    /// invoke this only for proposals targeting the registry component,
+    /// in queue order.
+    pub(crate) fn fold(
+        &mut self,
+        operation: &AppDataUpdateOperation,
+    ) -> Result<(), ComponentSourceError> {
+        use xmtp_mls_common::app_data::components::tls_map_components::ComponentRegistryComponent;
+        use xmtp_mls_common::app_data::typed::Component;
+        match operation {
+            AppDataUpdateOperation::Update(payload) => {
+                let new_bytes = <ComponentRegistryComponent as Component>::apply_update_payload(
+                    payload.as_slice(),
+                    self.snapshot_bytes.as_deref(),
+                )
+                .map_err(ComponentSourceError::from)?;
+                self.set_snapshot(new_bytes)
+            }
+            AppDataUpdateOperation::Remove => {
+                // Removing the whole registry component is never legal
+                // and is rejected at validation; the optimistic
+                // apply-path view of "whatever follows" is deny-all.
+                self.registry = ComponentRegistry::new();
+                self.snapshot_bytes = None;
+                Ok(())
+            }
+        }
+    }
+
+    /// Adopt already-computed post-proposal snapshot bytes (the apply
+    /// path computes them anyway via its in-batch chaining; re-applying
+    /// the delta here would duplicate work).
+    pub(crate) fn set_snapshot(&mut self, bytes: Vec<u8>) -> Result<(), ComponentSourceError> {
+        self.registry = ComponentRegistry::from_bytes(&bytes).map_err(|e| {
+            ComponentSourceError::MalformedComponentValue {
+                component_id: ComponentId::COMPONENT_REGISTRY,
+                reason: format!("same-commit registry overlay: {e}"),
+            }
+        })?;
+        self.snapshot_bytes = Some(bytes);
+        Ok(())
+    }
+
+    /// Deny-all view — the registry component was removed in-batch.
+    pub(crate) fn clear(&mut self) {
+        self.registry = ComponentRegistry::new();
+        self.snapshot_bytes = None;
+    }
+}
+
 #[cfg(any(test, feature = "test-utils"))]
 tokio::task_local! {
     /// Test-only override returned by [`load_component_registry`].
@@ -101,15 +204,12 @@ where
 {
     let mut in_batch: BTreeMap<openmls::component::ComponentId, Option<Vec<u8>>> = BTreeMap::new();
 
-    // Load the pre-commit registry once. It supplies the
-    // `ComponentType` tag the type-aware dispatcher in
-    // `apply_app_data_update_payload` uses when an unknown component id
-    // arrives. Registry updates that land in the same commit don't
-    // retroactively change this snapshot — the typed path would need
-    // an in-batch registry overlay to handle the corner case where the
-    // very same commit both registers a new component and writes to
-    // it.
-    let registry = load_component_registry(mls_group)?;
+    // Registry state for type dispatch, starting from the pre-commit
+    // snapshot. `COMPONENT_REGISTRY` proposals in this batch fold into
+    // the overlay in queue order, so a component registered earlier in
+    // this same commit dispatches like one registered in an earlier
+    // commit (see [`RegistryOverlay`]).
+    let mut overlay = RegistryOverlay::from_group(mls_group)?;
 
     for (openmls_id, operation) in proposals {
         let xmtp_id = ComponentId::from(openmls_id);
@@ -127,18 +227,21 @@ where
                         xmtp_id,
                         payload.as_slice(),
                         Some(bytes.as_slice()),
-                        &registry,
+                        overlay.registry(),
+                    ),
+                    Some(None) => apply_app_data_update_payload(
+                        xmtp_id,
+                        payload.as_slice(),
+                        None,
+                        overlay.registry(),
                     ),
-                    Some(None) => {
-                        apply_app_data_update_payload(xmtp_id, payload.as_slice(), None, &registry)
-                    }
                     None => {
                         let from_dict = read_from_app_data_dict(xmtp_id, mls_group);
                         apply_app_data_update_payload(
                             xmtp_id,
                             payload.as_slice(),
                             from_dict.as_deref(),
-                            &registry,
+                            overlay.registry(),
                         )
                     }
                 }
@@ -149,9 +252,18 @@ where
                         "Failed to apply AppDataUpdate payload"
                     );
                 })?;
+                if xmtp_id == ComponentId::COMPONENT_REGISTRY {
+                    // The in-batch chaining above already computed the
+                    // post-proposal snapshot; adopt it instead of
+                    // re-applying the delta.
+                    overlay.set_snapshot(new_value.clone())?;
+                }
                 in_batch.insert(openmls_id, Some(new_value));
             }
             AppDataUpdateOperation::Remove => {
+                if xmtp_id == ComponentId::COMPONENT_REGISTRY {
+                    overlay.clear();
+                }
                 in_batch.insert(openmls_id, None);
             }
         }

crates/xmtp_mls/src/groups/tests/test_proposals.rs

@@ -3571,6 +3571,206 @@ async fn test_accumulate_app_data_updates_chains_intra_batch() {
     );
 }
 
+/// Build the two `AppDataUpdateOperation`s for the same-commit
+/// registry-lookahead tests: a `COMPONENT_REGISTRY` delta registering
+/// `custom_id` as an allow-all String component, and a write of
+/// `value` to that component.
+fn lookahead_register_and_write_ops(
+    custom_id: xmtp_mls_common::app_data::component_id::ComponentId,
+    value: &[u8],
+) -> (
+    openmls::messages::proposals::AppDataUpdateOperation,
+    openmls::messages::proposals::AppDataUpdateOperation,
+) {
+    use openmls::messages::proposals::AppDataUpdateOperation;
+    use prost::Message;
+    use tls_codec::{Serialize, VLBytes};
+    use xmtp_mls_common::{
+        app_data::{
+            component_id::ComponentId, component_permissions::component_permissions,
+            component_registry::new_component_metadata,
+        },
+        tls_map::TlsMapDelta,
+    };
+    use xmtp_proto::xmtp::mls::message_contents::{
+        ComponentType, MetadataPolicy as MetadataPolicyProto,
+        metadata_policy::{Kind as MetadataPolicyKind, MetadataBasePolicy},
+    };
+
+    let allow = || MetadataPolicyProto {
+        kind: Some(MetadataPolicyKind::Base(MetadataBasePolicy::Allow as i32)),
+    };
+    let meta = new_component_metadata(
+        component_permissions()
+            .insert(allow())
+            .update(allow())
+            .delete(allow())
+            .call(),
+        ComponentType::String,
+    );
+    let registry_delta = TlsMapDelta::<ComponentId, VLBytes>::new()
+        .insert(custom_id, VLBytes::new(meta.encode_to_vec()));
+    let register = AppDataUpdateOperation::Update(
+        registry_delta
+            .tls_serialize_detached()
+            .expect("delta serializes")
+            .into(),
+    );
+    let write = AppDataUpdateOperation::Update(value.to_vec().into());
+    (register, write)
+}
+
+/// Same-commit registry lookahead, apply path: registering a custom
+/// component and writing to it inside ONE batch (queue order:
+/// register first) must dispatch the write through the just-registered
+/// type — exactly as if the registration had landed in an earlier
+/// commit. Pre-lookahead, the write failed `UnknownComponent` because
+/// dispatch only consulted the pre-commit registry snapshot.
+#[xmtp_common::test(unwrap_try = true)]
+async fn test_accumulate_registry_lookahead_register_then_write() {
+    use crate::groups::app_data::accumulate_app_data_updates;
+    use xmtp_mls_common::app_data::component_id::ComponentId;
+
+    tester!(alix);
+    let alix_group = alix.create_group(None, None)?;
+
+    let custom_id = ComponentId::new(0xC300);
+    let (op_register, op_write) = lookahead_register_and_write_ops(custom_id, b"hello");
+
+    let updates: openmls::group::AppDataUpdates = alix_group
+        .load_mls_group_with_lock_async(async |g| {
+            let out = accumulate_app_data_updates(
+                &g,
+                [
+                    (ComponentId::COMPONENT_REGISTRY.as_u16(), &op_register),
+                    (custom_id.as_u16(), &op_write),
+                ],
+            )
+            .map_err(crate::groups::GroupError::from)?;
+            Ok::<openmls::group::AppDataUpdates, crate::groups::GroupError>(
+                out.expect("two updates should produce dict changes"),
+            )
+        })
+        .await?;
+
+    let mut wrote: Option<Vec<u8>> = None;
+    let mut registered = false;
+    for (id, value) in updates {
+        if id == custom_id.as_u16() {
+            wrote = value;
+        } else if id == ComponentId::COMPONENT_REGISTRY.as_u16() {
+            registered = value.is_some();
+        }
+    }
+    assert!(registered, "registry snapshot missing from dict updates");
+    assert_eq!(
+        wrote.as_deref(),
+        Some(&b"hello"[..]),
+        "write to the just-registered component must dispatch via the overlay"
+    );
+}
+
+/// Queue-order discipline: the write coming BEFORE its registration in
+/// the same batch must still fail `UnknownComponent` — the overlay only
+/// looks back, never ahead, so every honest receiver rejects
+/// identically regardless of version.
+#[xmtp_common::test(unwrap_try = true)]
+async fn test_accumulate_registry_lookahead_write_before_register_fails() {
+    use crate::groups::app_data::{accumulate_app_data_updates, component_source};
+    use xmtp_mls_common::app_data::component_id::ComponentId;
+
+    tester!(alix);
+    let alix_group = alix.create_group(None, None)?;
+
+    let custom_id = ComponentId::new(0xC301);
+    let (op_register, op_write) = lookahead_register_and_write_ops(custom_id, b"hello");
+
+    let err = alix_group
+        .load_mls_group_with_lock_async(async |g| {
+            Ok::<_, crate::groups::GroupError>(accumulate_app_data_updates(
+                &g,
+                [
+                    (custom_id.as_u16(), &op_write),
+                    (ComponentId::COMPONENT_REGISTRY.as_u16(), &op_register),
+                ],
+            ))
+        })
+        .await?
+        .unwrap_err();
+    assert!(
+        matches!(
+            err,
+            component_source::ComponentSourceError::UnknownComponent(id) if id == custom_id
+        ),
+        "write-before-register must fail UnknownComponent, got {err:?}"
+    );
+}
+
+/// Same-commit registry lookahead, validate path: the per-proposal
+/// policy check for a write to a just-registered component succeeds
+/// against the folded overlay and fails (deny-by-default
+/// `NoRegistryEntry` → `InsufficientPermissions`) without it. Drives
+/// `validate_one_app_data_update` + `RegistryOverlay::fold` directly —
+/// the same pair `validate_app_data_update_proposals_in_commit` walks
+/// in queue order.
+#[xmtp_common::test(unwrap_try = true)]
+async fn test_validate_registry_lookahead_policy_sees_folded_entry() {
+    use crate::groups::app_data::RegistryOverlay;
+    use crate::groups::validated_commit::validate_one_app_data_update;
+    use xmtp_mls_common::app_data::{
+        component_id::ComponentId, component_registry::ComponentRegistry,
+        validation::ActorAuthority,
+    };
+
+    tester!(alix);
+    let alix_group = alix.create_group(None, None)?;
+
+    let custom_id = ComponentId::new(0xC302);
+    let (op_register, op_write) = lookahead_register_and_write_ops(custom_id, b"hello");
+    let member = ActorAuthority {
+        is_admin: false,
+        is_super_admin: false,
+    };
+
+    alix_group
+        .load_mls_group_with_lock_async(async |g| {
+            let mut overlay = RegistryOverlay::new(ComponentRegistry::new(), None);
+
+            // Without the fold: deny-by-default, no registry entry.
+            assert!(
+                validate_one_app_data_update(
+                    custom_id,
+                    &op_write,
+                    member,
+                    "test-inbox",
+                    overlay.registry(),
+                    &g,
+                )
+                .is_err(),
+                "write to an unregistered component must be rejected"
+            );
+
+            // Fold the registration (in the real loop this happens only
+            // after the registry proposal passes its super-admin check).
+            overlay
+                .fold(&op_register)
+                .expect("valid registry delta folds");
+
+            // With the fold: the allow-all entry admits the write.
+            validate_one_app_data_update(
+                custom_id,
+                &op_write,
+                member,
+                "test-inbox",
+                overlay.registry(),
+                &g,
+            )
+            .expect("write after same-batch registration must validate");
+            Ok::<(), crate::groups::GroupError>(())
+        })
+        .await?;
+}
+
 // =============================================================================
 // Sender-path tests for `IntentKind::UpdateAdminList` and
 // `IntentKind::UpdatePermission`. The AppDataUpdate path activates on