ping-openmls-sdk-core 0.6.7

//! Conversation state — wraps an OpenMLS `MlsGroup`.
//!
//! Each external conversation maps 1:1 to an MLS group whose leaves are devices. The DeviceGroup
//! (one per user, devices only) is just a special-cased conversation with the same wrapper.
//!
//! Persistence: we snapshot the `MlsGroup` after every state-changing operation under
//! `groups/{conversation_id}` and cache the result in-memory.

use openmls::{
    framing::{MlsMessageOut, ProcessedMessageContent},
    group::{MlsGroup, MlsGroupCreateConfig, MlsGroupJoinConfig},
    prelude::{
        tls_codec::{Deserialize as TlsDeserialize, Serialize as TlsSerialize},
        BasicCredential, Ciphersuite, CredentialWithKey, Extension, Extensions, MlsMessageBodyIn,
        MlsMessageIn, ProcessedMessage, ProtocolMessage, ProtocolVersion, UnknownExtension,
    },
};
use openmls_basic_credential::SignatureKeyPair;
use openmls_traits::OpenMlsProvider;
use ping_mls_store::PersistentMlsProvider;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::sync::Arc;
use ulid::Ulid;
use zeroize::Zeroizing;

use crate::{
    clock::Hlc,
    codec,
    device::{DeviceId, GroupSnapshotEntry, GroupStateSnapshot, GROUP_SNAPSHOT_VERSION},
    error::{Error, Result},
    identity::UserId,
    message::{IncomingMessage, MessageEnvelope, MessageKind},
    storage::Storage,
    sync::SyncCursor,
};

const DEFAULT_CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;

/// MLS GroupContext extension type carrying the human conversation `name`.
///
/// Stored in the group context so the name is part of shared MLS group state and
/// therefore present on EVERY device that holds the group — including a device
/// that joins via a Welcome. Without this the name was creator-local and a
/// joiner/linked device saw `None`. A private-use `Unknown` extension type
/// (not a GREASE `0x?A?A` value, not a registered type 0x0001–0x0005).
///
/// It is added WITHOUT a `RequiredCapabilities` extension, so openmls imposes no
/// per-member capability check — existing KeyPackages keep working and no
/// re-link is required. The value is UTF-8 bytes of the name.
const GROUP_NAME_EXTENSION_TYPE: u16 = 0xFF00;

/// Read the conversation `name` from a group's GroupContext extensions, if set.
fn group_name_from_extensions(extensions: &Extensions) -> Option<String> {
    extensions.iter().find_map(|ext| match ext {
        Extension::Unknown(ext_type, data) if *ext_type == GROUP_NAME_EXTENSION_TYPE => {
            String::from_utf8(data.0.clone())
                .ok()
                .filter(|s| !s.is_empty())
        }
        _ => None,
    })
}

/// Build the GroupContext extensions carrying `name` (empty when `name` is
/// `None`/blank), for `MlsGroupCreateConfig::with_group_context_extensions`.
fn group_context_extensions_for_name(name: Option<&str>) -> Extensions {
    match name {
        Some(n) if !n.is_empty() => Extensions::single(Extension::Unknown(
            GROUP_NAME_EXTENSION_TYPE,
            UnknownExtension(n.as_bytes().to_vec()),
        )),
        _ => Extensions::empty(),
    }
}

/// 16-byte conversation identifier (ULID encoded). Stable across epochs.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ConversationId(#[serde(with = "serde_bytes_array16")] pub [u8; 16]);

impl ConversationId {
    pub fn new() -> Self {
        Self(Ulid::new().to_bytes())
    }
    pub fn as_hex(&self) -> String {
        hex::encode(self.0)
    }
}

impl Default for ConversationId {
    fn default() -> Self {
        Self::new()
    }
}

mod serde_bytes_array16 {
    use serde::{Deserializer, Serializer};
    pub fn serialize<S: Serializer>(b: &[u8; 16], s: S) -> Result<S::Ok, S::Error> {
        serde_bytes::serialize(b.as_slice(), s)
    }
    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<[u8; 16], D::Error> {
        let v: Vec<u8> = serde_bytes::deserialize(d)?;
        v.try_into()
            .map_err(|_| serde::de::Error::custom("expected 16 bytes"))
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConversationMeta {
    pub id: ConversationId,
    pub name: Option<String>,
    pub epoch: u64,
    pub member_count: u32,
    pub is_device_group: bool,
    pub created_at_ms: u64,
}

/// One member leaf of a conversation's MLS group: the member's [`UserId`]
/// (recovered from the leaf's `BasicCredential`) and its ratchet-tree leaf
/// index. A user with multiple devices appears once **per device leaf** —
/// callers that want a per-user roster should dedup by `user_id`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemberInfo {
    pub user_id: UserId,
    pub leaf_index: u32,
}

/// In-memory conversation handle. Holds the OpenMLS group plus our wire-level cursor.
pub struct Conversation {
    pub(crate) id: ConversationId,
    pub(crate) meta: ConversationMeta,
    pub(crate) group: MlsGroup,
    pub(crate) crypto: Arc<PersistentMlsProvider>,
    pub(crate) signing: Arc<SignatureKeyPair>,
    pub(crate) own_device: DeviceId,
    pub(crate) seq: u64,
    pub(crate) hlc: Hlc,
    pub(crate) cursor: SyncCursor,
    pub(crate) storage: Arc<dyn Storage>,
    /// Local device→leaf-index map for [CR-2] revocation.
    ///
    /// Populated when this device either (a) admits a peer via [`Self::add_members`] —
    /// every entry in the `Vec<(DeviceId, KeyPackage)>` is recorded after the commit
    /// merges — or (b) joins as the receiving device via [`Self::join`], at which point
    /// we record our own leaf. Pruned when [`Self::remove_members`] is called.
    ///
    /// Not authoritative for *peers' devices we didn't admit*: those are visible in
    /// `group.members()` but their device_ids are opaque to this client. `revoke_device`
    /// is therefore best-effort across conversations we ourselves invited the device
    /// into; see [`MessagingClient::revoke_device`] for the documented scope.
    pub(crate) device_leaves: BTreeMap<DeviceId, u32>,
}

impl std::fmt::Debug for Conversation {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Conversation")
            .field("id", &self.id.as_hex())
            .field("meta", &self.meta)
            .finish()
    }
}

impl Conversation {
    pub fn id(&self) -> ConversationId {
        self.id
    }
    pub fn meta(&self) -> &ConversationMeta {
        &self.meta
    }

    /// Current member roster, recovered locally from the MLS group's leaf
    /// credentials — no network and no out-of-band `ping.profile` message.
    /// Each `BasicCredential` was built from the member's `UserId`
    /// (`BasicCredential::new(own_user.0.clone())` in [`Self::create`] /
    /// [`Self::join`]), so we round-trip it back here. Each entry is one
    /// leaf; a multi-device user appears once per device leaf.
    pub fn members(&self) -> Vec<MemberInfo> {
        self.group
            .members()
            .filter_map(|m| {
                let basic = BasicCredential::try_from(m.credential).ok()?;
                Some(MemberInfo {
                    user_id: UserId(basic.identity().to_vec()),
                    leaf_index: m.index.u32(),
                })
            })
            .collect()
    }

    pub fn epoch(&self) -> u64 {
        self.group.epoch().as_u64()
    }
    pub fn cursor(&self) -> &SyncCursor {
        &self.cursor
    }

    /// Create a new conversation, with `self` as the only initial member.
    // 8 args is a lot, but they're all needed for an internal constructor and a builder
    // would be over-engineered for v0.1.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn create(
        id: ConversationId,
        name: Option<String>,
        own_device: DeviceId,
        own_user: &UserId,
        crypto: Arc<PersistentMlsProvider>,
        signing: Arc<SignatureKeyPair>,
        storage: Arc<dyn Storage>,
        now_ms: u64,
    ) -> Result<Self> {
        let credential = BasicCredential::new(own_user.0.clone());
        let credential_with_key = CredentialWithKey {
            credential: credential.into(),
            signature_key: signing.public().into(),
        };
        // Carry the conversation name in the GroupContext so it travels in the
        // Welcome to every joiner (see `GROUP_NAME_EXTENSION_TYPE`). No
        // RequiredCapabilities is added, so this imposes no member capability
        // check and stays compatible with existing KeyPackages.
        let cfg = MlsGroupCreateConfig::builder()
            .ciphersuite(DEFAULT_CIPHERSUITE)
            .use_ratchet_tree_extension(true)
            .with_group_context_extensions(group_context_extensions_for_name(name.as_deref()))
            .map_err(Error::mls)?
            .build();
        let group = MlsGroup::new_with_group_id(
            crypto.as_ref(),
            signing.as_ref(),
            &cfg,
            openmls::group::GroupId::from_slice(&id.0),
            credential_with_key,
        )
        .map_err(Error::mls)?;

        let meta = ConversationMeta {
            id,
            name,
            epoch: 0,
            member_count: 1,
            is_device_group: false,
            created_at_ms: now_ms,
        };
        // [CR-2] Group creator is always leaf 0; record so revoke_device can target it.
        let mut device_leaves = BTreeMap::new();
        device_leaves.insert(own_device.clone(), group.own_leaf_index().u32());
        Ok(Self {
            id,
            meta,
            group,
            crypto,
            signing,
            own_device,
            seq: 0,
            hlc: Hlc::ZERO.tick(now_ms),
            cursor: SyncCursor::default(),
            storage,
            device_leaves,
        })
    }

    /// Join an existing conversation from a Welcome message.
    pub(crate) fn join(
        welcome_bytes: &[u8],
        own_device: DeviceId,
        crypto: Arc<PersistentMlsProvider>,
        signing: Arc<SignatureKeyPair>,
        storage: Arc<dyn Storage>,
        now_ms: u64,
    ) -> Result<Self> {
        let mls_in = MlsMessageIn::tls_deserialize_exact(welcome_bytes).map_err(Error::mls)?;
        let welcome = match mls_in.extract() {
            MlsMessageBodyIn::Welcome(w) => w,
            _ => return Err(Error::Invalid("expected Welcome".into())),
        };
        let cfg = MlsGroupJoinConfig::builder()
            .use_ratchet_tree_extension(true)
            .build();
        let staged =
            openmls::group::StagedWelcome::new_from_welcome(crypto.as_ref(), &cfg, welcome, None)
                .map_err(Error::mls)?;
        let group = staged.into_group(crypto.as_ref()).map_err(Error::mls)?;

        let id_bytes: [u8; 16] = group
            .group_id()
            .as_slice()
            .try_into()
            .map_err(|_| Error::Invalid("group id must be 16 bytes".into()))?;
        let id = ConversationId(id_bytes);
        // The creator stamped the name into the GroupContext, which is part of
        // the shared group state the Welcome reconstructs — so a joining device
        // recovers the real name here instead of starting with `None`.
        let name = group_name_from_extensions(group.extensions());
        let meta = ConversationMeta {
            id,
            name,
            epoch: group.epoch().as_u64(),
            member_count: group.members().count() as u32,
            is_device_group: false,
            created_at_ms: now_ms,
        };

        // Seed the cursor at the join epoch so subsequent fetches skip pre-join Commits
        // (notably the Add commit that produced this Welcome — it lives in the conversation
        // log at `epoch - 1`, which the joiner must not try to apply on top of its
        // already-advanced group state).
        let join_epoch = group.epoch().as_u64();
        // [CR-2] Record our own (device_id → leaf_index) so the host can later revoke us
        // via the standard `revoke_device` flow. `own_leaf_index()` is stable for the
        // lifetime of this group membership.
        let own_leaf = group.own_leaf_index().u32();
        let mut device_leaves = BTreeMap::new();
        device_leaves.insert(own_device.clone(), own_leaf);
        Ok(Self {
            id,
            meta,
            group,
            crypto,
            signing,
            own_device,
            seq: 0,
            hlc: Hlc::ZERO.tick(now_ms),
            cursor: SyncCursor {
                epoch: join_epoch,
                ..Default::default()
            },
            storage,
            device_leaves,
        })
    }

    /// The conversation name recovered from MLS GroupContext state (set at
    /// creation, carried in the Welcome to every joiner). `None` if unnamed.
    pub(crate) fn name_from_group_state(&self) -> Option<String> {
        group_name_from_extensions(self.group.extensions())
    }

    /// [CR-4] Rehydrate a previously-persisted conversation on cold restart.
    ///
    /// Loads the OpenMLS group state via `MlsGroup::load` (which reads from the
    /// provider's storage — populated by the SQLite-backed checkpoint on the
    /// previous run). Pairs the loaded MLS state with the meta + cursor + device→leaf
    /// map the host-side `Storage` trait kept for us. Returns `Ok(None)` if OpenMLS
    /// finds no state for `id` — the host's `groups` namespace had a stale entry.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn load(
        id: ConversationId,
        meta: ConversationMeta,
        cursor: SyncCursor,
        device_leaves: BTreeMap<DeviceId, u32>,
        own_device: DeviceId,
        crypto: Arc<PersistentMlsProvider>,
        signing: Arc<SignatureKeyPair>,
        storage: Arc<dyn Storage>,
        now_ms: u64,
    ) -> Result<Option<Self>> {
        use openmls::group::GroupId;
        let group_id = GroupId::from_slice(&id.0);
        let group = match MlsGroup::load(crypto.storage(), &group_id).map_err(Error::mls)? {
            Some(g) => g,
            None => return Ok(None),
        };
        // Restore the local outgoing-send counter from the persisted cursor. The cursor
        // tracks the highest applied (epoch, sender, seq) for every device — including
        // our own — so we can recover `self.seq` from `cursor.last_seq_per_device[own]`.
        // Without this, the next `send_application()` re-uses an already-consumed seq
        // and receivers silently dedupe (cursor.is_new returns false on their side).
        let seq = cursor
            .last_seq_per_device
            .get(&own_device)
            .copied()
            .unwrap_or(0);
        Ok(Some(Self {
            id,
            meta,
            group,
            crypto,
            signing,
            own_device,
            seq,
            hlc: Hlc::ZERO.tick(now_ms),
            cursor,
            storage,
            device_leaves,
        }))
    }

    /// Encrypt an application message and produce a wire envelope ready for transport.
    ///
    /// Uses the [CR-6] plaintext content_hash path: the envelope's `content_hash` is
    /// `SHA-256(plaintext)`, not the MLS ciphertext. This is what makes rebase clean
    /// and gives cross-binding hash parity.
    pub fn send_application(&mut self, plaintext: &[u8], now_ms: u64) -> Result<MessageEnvelope> {
        let out = self
            .group
            .create_message(self.crypto.as_ref(), self.signing.as_ref(), plaintext)
            .map_err(Error::mls)?;

        self.seq += 1;
        self.hlc = self.hlc.tick(now_ms);
        let bytes = out.tls_serialize_detached().map_err(Error::mls)?;
        let env = MessageEnvelope::new_application(
            self.id,
            self.epoch(),
            self.own_device.clone(),
            self.seq,
            self.hlc,
            bytes,
            plaintext,
        );
        // Advance the local cursor past our own send so a subsequent catch-up sync doesn't
        // pull this envelope back to us (we've already applied it locally — re-processing
        // would either fail or duplicate-deliver).
        self.cursor.advance(
            env.epoch,
            self.own_device.clone(),
            self.seq,
            self.hlc,
            now_ms,
        );
        Ok(env)
    }

    /// Add members by KeyPackage. Produces the Commit envelope to broadcast plus the Welcome
    /// envelope(s) to deliver out-of-band to the newly-added devices.
    ///
    /// [CR-2] takes a `Vec<(DeviceId, KeyPackage)>` instead of a bare `Vec<KeyPackage>`. The
    /// `DeviceId` for each entry is the *caller's* assertion of which device owns that
    /// KeyPackage — hosts typically get it from the directory service alongside the
    /// KeyPackage itself. The mapping is persisted per-conversation so [`MessagingClient::revoke_device`]
    /// can later locate the leaf to remove without a fresh directory lookup. The SDK does
    /// not cryptographically verify the device claim; that's a host policy concern
    /// (typically: the directory authenticates the key_package_id → device_id mapping).
    pub fn add_members(
        &mut self,
        entries: Vec<(DeviceId, Vec<u8>)>,
        now_ms: u64,
    ) -> Result<AddOutcome> {
        // All-in-one (stage + immediate merge) for callers that commit and
        // persist synchronously with NO networked rollback window — e.g. the
        // device-group / device-linking paths. The networked group path in
        // `client.rs` instead uses `stage_add_members` + `confirm_staged` /
        // `abort_staged`, so a Commit the server REJECTS can be rolled back
        // rather than leaving the local epoch ahead of the server (the desync
        // that permanently bricks a group: every later Commit then 409s and
        // peers can't decrypt our epoch).
        let staged = self.stage_add_members(entries, now_ms)?;
        self.confirm_staged(&staged, now_ms)?;
        let StagedCommit {
            commit, welcome, ..
        } = staged;
        let welcome =
            welcome.ok_or_else(|| Error::Invalid("add_members produced no Welcome".into()))?;
        Ok(AddOutcome { commit, welcome })
    }

    /// Stage an add-members Commit WITHOUT merging it — the group keeps a
    /// *pending* commit and the local epoch is UNCHANGED. Returns the Commit +
    /// Welcome envelopes to send. The caller MUST follow with exactly one of
    /// [`Self::confirm_staged`] (server accepted the Commit → merge locally) or
    /// [`Self::abort_staged`] (server rejected it → discard, epoch never moves).
    ///
    /// This is what makes the local epoch advance *only after* the server
    /// accepts the Commit (send-then-merge), so a rejected/conflicting Commit
    /// can't desync the group. Safe because the JS worker serializes top-level
    /// requests — nothing else touches this group during the send round-trip
    /// between stage and confirm/abort (worker.ts "SERIALIZE … dispatch").
    pub(crate) fn stage_add_members(
        &mut self,
        entries: Vec<(DeviceId, Vec<u8>)>,
        now_ms: u64,
    ) -> Result<StagedCommit> {
        let mut kps = Vec::with_capacity(entries.len());
        // Track signature_key → device_id so we can resolve leaf indices post-commit.
        let mut sig_to_device: Vec<(Vec<u8>, DeviceId)> = Vec::with_capacity(entries.len());
        for (device_id, raw) in &entries {
            let mls_in = MlsMessageIn::tls_deserialize_exact(raw).map_err(Error::mls)?;
            let kp_in = match mls_in.extract() {
                MlsMessageBodyIn::KeyPackage(kp) => kp,
                _ => return Err(Error::Invalid("expected KeyPackage".into())),
            };
            // KeyPackages on the wire are unvalidated (`KeyPackageIn`); validate against the
            // crypto provider before handing them to OpenMLS.
            let kp = kp_in
                .validate(self.crypto.crypto(), ProtocolVersion::default())
                .map_err(Error::mls)?;
            let sig_key = kp.leaf_node().signature_key().as_slice().to_vec();
            sig_to_device.push((sig_key, device_id.clone()));
            kps.push(kp);
        }

        // The Commit's wire `epoch` is the *source* epoch (where it was crafted). The
        // Welcome's `epoch` is the *post-commit* epoch. An MLS Commit advances the epoch by
        // exactly 1, so we can name the post-commit epoch as `pre + 1` WITHOUT merging.
        let pre_commit_epoch = self.epoch();
        let post_commit_epoch = pre_commit_epoch + 1;

        let (commit_out, welcome_out, _gi) = self
            .group
            .add_members(self.crypto.as_ref(), self.signing.as_ref(), &kps)
            .map_err(Error::mls)?;
        // NB: NO merge here — the pending commit is merged by `confirm_staged`
        // only once the server has accepted the Commit send.

        let next_seq = self.seq + 1;
        let next_hlc = self.hlc.tick(now_ms);

        let commit_bytes = mls_message_out_bytes(commit_out)?;
        let commit_env = MessageEnvelope::new(
            self.id,
            pre_commit_epoch,
            MessageKind::Commit,
            self.own_device.clone(),
            next_seq,
            next_hlc,
            commit_bytes,
        );

        let welcome_bytes = mls_message_out_bytes(welcome_out)?;
        let welcome_env = MessageEnvelope::new(
            self.id,
            post_commit_epoch,
            MessageKind::Welcome,
            self.own_device.clone(),
            next_seq,
            next_hlc,
            welcome_bytes,
        );

        Ok(StagedCommit {
            commit: commit_env,
            welcome: Some(welcome_env),
            next_seq,
            next_hlc,
            leaf_update: StagedLeafUpdate::Add(sig_to_device),
        })
    }

    pub fn remove_members(
        &mut self,
        leaf_indexes: Vec<u32>,
        now_ms: u64,
    ) -> Result<MessageEnvelope> {
        // All-in-one (stage + immediate merge). See `add_members` for why the
        // networked path uses stage/confirm/abort instead.
        let staged = self.stage_remove_members(leaf_indexes, now_ms)?;
        self.confirm_staged(&staged, now_ms)?;
        let StagedCommit { commit, .. } = staged;
        Ok(commit)
    }

    /// Stage a remove-members Commit WITHOUT merging it — see
    /// [`Self::stage_add_members`]. No Welcome (removals don't admit anyone).
    pub(crate) fn stage_remove_members(
        &mut self,
        leaf_indexes: Vec<u32>,
        now_ms: u64,
    ) -> Result<StagedCommit> {
        use openmls::prelude::LeafNodeIndex;
        let leaves: Vec<LeafNodeIndex> = leaf_indexes
            .iter()
            .copied()
            .map(LeafNodeIndex::new)
            .collect();

        let pre_commit_epoch = self.epoch();

        let (commit_out, _welcome_opt, _gi) = self
            .group
            .remove_members(self.crypto.as_ref(), self.signing.as_ref(), &leaves)
            .map_err(Error::mls)?;
        // NB: NO merge here — see `stage_add_members`.

        let next_seq = self.seq + 1;
        let next_hlc = self.hlc.tick(now_ms);
        let bytes = mls_message_out_bytes(commit_out)?;
        let commit_env = MessageEnvelope::new(
            self.id,
            pre_commit_epoch,
            MessageKind::Commit,
            self.own_device.clone(),
            next_seq,
            next_hlc,
            bytes,
        );

        let removed: std::collections::HashSet<u32> = leaf_indexes.iter().copied().collect();
        Ok(StagedCommit {
            commit: commit_env,
            welcome: None,
            next_seq,
            next_hlc,
            leaf_update: StagedLeafUpdate::Remove(removed),
        })
    }

    /// Merge a previously [staged](Self::stage_add_members) Commit into the local
    /// group — call ONLY after the server has accepted the Commit send. Advances
    /// the epoch, updates the roster + device→leaf map, bumps seq/hlc, and moves
    /// the sync cursor past our own Commit so catch-up doesn't re-apply it.
    pub(crate) fn confirm_staged(&mut self, staged: &StagedCommit, now_ms: u64) -> Result<()> {
        self.group
            .merge_pending_commit(self.crypto.as_ref())
            .map_err(Error::mls)?;
        self.meta.epoch = self.epoch();
        self.meta.member_count = self.group.members().count() as u32;

        match &staged.leaf_update {
            StagedLeafUpdate::Add(sig_to_device) => {
                // [CR-2] Resolve leaf indexes for the devices we just added (match by the
                // per-device MLS signature_key, unique per device).
                for member in self.group.members() {
                    if let Some((_, device_id)) = sig_to_device
                        .iter()
                        .find(|(sig, _)| sig.as_slice() == member.signature_key.as_slice())
                    {
                        self.device_leaves
                            .insert(device_id.clone(), member.index.u32());
                    }
                }
            }
            StagedLeafUpdate::Remove(removed) => {
                // [CR-2] Prune the device→leaf map for removed leaves. Other entries' leaf
                // indexes are stable (OpenMLS reuses blank slots, doesn't reshuffle).
                self.device_leaves.retain(|_, idx| !removed.contains(idx));
            }
        }

        self.seq = staged.next_seq;
        self.hlc = staged.next_hlc;
        self.cursor.advance(
            self.meta.epoch,
            self.own_device.clone(),
            self.seq,
            self.hlc,
            now_ms,
        );
        Ok(())
    }

    /// Discard a previously [staged](Self::stage_add_members) Commit — call when
    /// the server REJECTED the Commit send. Clears the pending commit so the
    /// local epoch stays exactly where it was (no desync) and the conversation
    /// is operational again. Idempotent / safe if there is no pending commit.
    pub(crate) fn abort_staged(&mut self) -> Result<()> {
        self.group
            .clear_pending_commit(self.crypto.storage())
            .map_err(Error::mls)?;
        Ok(())
    }

    /// Process an inbound envelope. Returns Some(IncomingMessage) for application traffic.
    pub fn process(
        &mut self,
        env: &MessageEnvelope,
        now_ms: u64,
    ) -> Result<Option<IncomingMessage>> {
        if !self.cursor.is_new(env.epoch, &env.sender_device, env.seq) {
            return Ok(None); // dedupe: already applied
        }
        let mls_in = MlsMessageIn::tls_deserialize_exact(&env.payload).map_err(Error::mls)?;

        // OpenMLS' `process_message` expects an `impl Into<ProtocolMessage>`. `MlsMessageIn`
        // itself doesn't implement that; we have to extract the body and convert the inner
        // private/public message. Welcomes are handled at the client level, not here.
        let protocol_msg: ProtocolMessage = match mls_in.extract() {
            MlsMessageBodyIn::PrivateMessage(m) => m.into(),
            MlsMessageBodyIn::PublicMessage(m) => m.into(),
            MlsMessageBodyIn::Welcome(_) => {
                return Err(Error::Invalid(
                    "Welcome must be handled at client level, not in-group".into(),
                ));
            }
            _ => return Err(Error::Invalid("unsupported MLS message body".into())),
        };

        let processed: ProcessedMessage = self
            .group
            .process_message(self.crypto.as_ref(), protocol_msg)
            .map_err(Error::mls)?;

        // Recover the sender's account-level `UserId` from their authenticated
        // leaf credential BEFORE `into_content()` consumes `processed`. Same
        // round-trip as `members()`: the leaf was built as
        // `BasicCredential::new(user.0)`, so `identity()` is the `UserId` bytes.
        // This lets the host attribute messages to the right account across
        // every linked device without any device→account side channel.
        let sender_user_id = BasicCredential::try_from(processed.credential().clone())
            .map(|c| UserId(c.identity().to_vec()))
            .unwrap_or_else(|_| UserId(Vec::new()));

        let out = match processed.into_content() {
            ProcessedMessageContent::ApplicationMessage(app) => {
                let pt = app.into_bytes();
                // CR-6: for v=2 application envelopes the wire-contract validator can't
                // check `content_hash` (the hash is over plaintext, which it didn't have).
                // We can now: verify SHA-256(pt) == env.content_hash and reject mismatches.
                // For v=1 envelopes the wire-contract validator already checked the
                // ciphertext-based hash, so no extra work here.
                if env.v >= 2 {
                    let computed = crate::message::hash_application_plaintext(&pt);
                    if computed != env.content_hash {
                        return Err(Error::Invalid(
                            "v=2 application content_hash mismatch".into(),
                        ));
                    }
                }
                Some(IncomingMessage {
                    conversation_id: self.id,
                    sender_device: env.sender_device.clone(),
                    sender_user_id,
                    epoch: env.epoch,
                    hlc: env.hlc,
                    plaintext: pt,
                    content_hash: env.content_hash,
                })
            }
            ProcessedMessageContent::StagedCommitMessage(staged) => {
                self.group
                    .merge_staged_commit(self.crypto.as_ref(), *staged)
                    .map_err(Error::mls)?;
                self.meta.epoch = self.epoch();
                self.meta.member_count = self.group.members().count() as u32;
                None
            }
            ProcessedMessageContent::ProposalMessage(_)
            | ProcessedMessageContent::ExternalJoinProposalMessage(_) => {
                // Proposals are buffered by OpenMLS until the next Commit; nothing to surface
                // to the application.
                None
            }
        };

        self.cursor.advance(
            env.epoch,
            env.sender_device.clone(),
            env.seq,
            env.hlc,
            now_ms,
        );
        Ok(out)
    }

    /// Export a derived secret keyed to this group's current epoch ([CR-8]).
    ///
    /// Wraps `MlsGroup::export_secret` (the MLS exporter, RFC 9420 §8.5) and surfaces the
    /// bytes in a `Zeroizing<Vec<u8>>` so the local copy is wiped on drop. Used by the host
    /// to seed:
    ///   * the ephemeral channel (`ping/ephemeral`, §5.4 of the architecture)
    ///   * call media keys (`ping/calls/media/{call_id}`, §7.2)
    ///   * call-ephemeral framer keys (`ping/calls/ephemeral/{call_id}`, §7.5)
    ///
    /// `label` should use the documented `ping/*` namespacing convention. There is no
    /// runtime enforcement — cross-binding parity is enforced by conformance fixtures
    /// pinning specific label strings.
    ///
    /// Output is the secret. Callers MUST treat the buffer as a secret: never log, never
    /// persist unencrypted. The wrapper zeroes our local copy on drop; the caller is
    /// responsible for zeroing any copy they make.
    pub fn export_secret(
        &self,
        label: &str,
        context: &[u8],
        length: usize,
    ) -> Result<Zeroizing<Vec<u8>>> {
        if length == 0 {
            return Err(Error::Invalid("export_secret length must be > 0".into()));
        }
        // Soft cap to prevent runaway allocations from a malformed caller. Real labels never
        // need more than ~64 bytes (AES-256 key + 96-bit nonce + slack); 1 KiB is generous.
        if length > 1024 {
            return Err(Error::Invalid(
                "export_secret length exceeds 1024-byte cap".into(),
            ));
        }
        let bytes = self
            .group
            .export_secret(self.crypto.as_ref(), label, context, length)
            .map_err(Error::mls)?;
        Ok(Zeroizing::new(bytes))
    }

    /// [CR-7] Export a portable snapshot of this group's MLS state.
    ///
    /// Walks the provider's working set, picks every entry whose key references this
    /// group's id, and bundles them with format metadata. Returns CBOR-encoded bytes
    /// suitable for inclusion in:
    ///   * `LinkingTicket.catchup_snapshot.conversation_metas[i].group_state_bytes`
    ///     (via [CR-13] — host calls this and passes the bytes through);
    ///   * `IdentityBackup.device_group_snapshot` (the Permissive-recovery path per
    ///     `docs/architecture/recovery.md`).
    ///
    /// Returns `Err` if the encoded snapshot exceeds [`GROUP_SNAPSHOT_HARD_CAP`].
    /// Output is wrapped in `Zeroizing` because the bytes contain past epoch secrets;
    /// the caller's copy on the FFI side is the host's responsibility to wipe.
    pub fn export_state_snapshot(&self, now_ms: u64) -> Result<Zeroizing<Vec<u8>>> {
        let entries = self.crypto.group_scoped_entries(&self.id.0);
        let snap = GroupStateSnapshot {
            v: GROUP_SNAPSHOT_VERSION,
            group_id: self.id,
            openmls_storage_version: openmls_traits::storage::CURRENT_VERSION,
            snapshot_created_at_ms: now_ms,
            entries: entries
                .into_iter()
                .map(|(key, value)| GroupSnapshotEntry { key, value })
                .collect(),
        };
        Ok(Zeroizing::new(snap.encode()?))
    }

    /// Look up the leaf index this device controls, if known ([CR-2]).
    ///
    /// Returns the locally-tracked leaf for `device_id`. Only populated for devices we
    /// added via [`Self::add_members`] or for our own leaf via [`Self::create`] /
    /// [`Self::join`]. Devices a peer admitted on our behalf are not in this map.
    pub fn leaf_index_of(&self, device_id: &DeviceId) -> Option<u32> {
        self.device_leaves.get(device_id).copied()
    }

    /// Synchronously capture everything [`ConversationSnapshot::flush`]
    /// needs to persist this conversation, so a caller can DROP the
    /// `conversations` lock BEFORE awaiting the async writes.
    ///
    /// Holding a `parking_lot` guard across `.await` is a latent bug: on
    /// the single-threaded wasm worker, a second client call that lands
    /// while the first is suspended (a waiting writer + a new reader)
    /// makes `parking_lot` try to PARK, and its wasm stub `panic!`s with
    /// "Parking not supported on this platform" — poisoning the module.
    /// Splitting the synchronous capture (under the lock) from the async
    /// flush (lock released) removes that hazard everywhere the snapshot
    /// runs for a conversation that lives inside the shared map. The
    /// capture is a consistent point-in-time view (cursor + meta + leaves
    /// + the Arc'd provider/storage handles).
    pub(crate) fn snapshot_inputs(&self) -> Result<ConversationSnapshot> {
        // [CR-2] Stable BTreeMap-of-pairs encoding → canonical CBOR so
        // every platform decodes identical bytes.
        let leaves_vec: Vec<(DeviceId, u32)> = self
            .device_leaves
            .iter()
            .map(|(d, i)| (d.clone(), *i))
            .collect();
        Ok(ConversationSnapshot {
            id: self.id,
            crypto: self.crypto.clone(),
            storage: self.storage.clone(),
            cursor: self.cursor.encode()?,
            meta: codec::encode(&self.meta)?,
            device_leaves: codec::encode(&leaves_vec)?,
        })
    }

    /// Persist this conversation's state. Convenience wrapper used by
    /// call sites that hold an OWNED `Conversation` (not borrowed from the
    /// shared map) — e.g. just-created/just-joined conversations before
    /// they're inserted, where no lock is held across the await. Map-
    /// resident callers MUST instead use `snapshot_inputs()` + drop the
    /// guard + `flush().await` (see client.rs) to avoid the wasm parking
    /// panic described on `snapshot_inputs`.
    pub(crate) async fn snapshot_to_storage(&self) -> Result<()> {
        self.snapshot_inputs()?.flush().await
    }
}

/// Point-in-time, lock-free snapshot of a [`Conversation`]'s persistable
/// state. Produced synchronously by [`Conversation::snapshot_inputs`] (so
/// the `conversations` lock can be dropped) and flushed asynchronously by
/// [`Self::flush`].
pub(crate) struct ConversationSnapshot {
    id: ConversationId,
    crypto: Arc<PersistentMlsProvider>,
    storage: Arc<dyn Storage>,
    cursor: Vec<u8>,
    meta: Vec<u8>,
    device_leaves: Vec<u8>,
}

impl ConversationSnapshot {
    /// Flush the captured state to storage. Safe to `.await` with NO
    /// `conversations` lock held — it only touches the Arc'd provider +
    /// storage handles, never the shared map.
    pub(crate) async fn flush(self) -> Result<()> {
        // [CR-4] Flush the MLS working set to the configured backend (no-op
        // for `StorageBackend::Memory`). MUST happen on every state-changing
        // op so a cold restart — iOS NSE, web SW — finds the latest epoch on
        // disk. `checkpoint_async` is required for the WASM `IndexedDb`
        // backend (IDB is async-only); native Memory / Sqlite paths await
        // trivially since their I/O is sync internally.
        self.crypto
            .checkpoint_async()
            .await
            .map_err(|e| Error::Storage(format!("checkpoint: {e}")))?;

        let hex = self.id.as_hex();
        self.storage.put("cursors", &hex, self.cursor).await?;
        self.storage
            .put("groups", &format!("{hex}/meta"), self.meta)
            .await?;
        // [CR-2] device→leaf map, persisted alongside meta + cursor so
        // revoke_device works after a cold restart.
        self.storage
            .put("device_leaves", &hex, self.device_leaves)
            .await?;
        Ok(())
    }
}

/// Both halves of an Add commit. The Commit goes on the conversation channel; the Welcome is
/// delivered to the new members via whatever out-of-band path the host uses (often the same
/// transport, addressed to the new device's mailbox).
#[derive(Debug, Clone)]
pub struct AddOutcome {
    pub commit: MessageEnvelope,
    pub welcome: MessageEnvelope,
}

/// The local-state mutation a staged Commit will apply on
/// [`Conversation::confirm_staged`]. Captured at stage time so confirm can run
/// after the (async) Commit send without re-deriving anything.
pub(crate) enum StagedLeafUpdate {
    /// Add: signature_key → device_id for each added device, resolved to a leaf
    /// index against the merged tree in `confirm_staged`.
    Add(Vec<(Vec<u8>, DeviceId)>),
    /// Remove: the leaf indexes being dropped from the device→leaf map.
    Remove(std::collections::HashSet<u32>),
}

/// A Commit produced but NOT yet merged (see [`Conversation::stage_add_members`]).
/// Held by the client across the Commit send; merged via
/// [`Conversation::confirm_staged`] on success or discarded via
/// [`Conversation::abort_staged`] on a server rejection. This is the unit of the
/// send-then-merge protocol that keeps the local epoch from ever running ahead of
/// the server.
pub(crate) struct StagedCommit {
    pub commit: MessageEnvelope,
    pub welcome: Option<MessageEnvelope>,
    next_seq: u64,
    next_hlc: Hlc,
    leaf_update: StagedLeafUpdate,
}

fn mls_message_out_bytes(m: MlsMessageOut) -> Result<Vec<u8>> {
    m.tls_serialize_detached().map_err(Error::mls)
}