ping-openmls-sdk-core 0.5.0

//! `MessagingClient` — top-level handle. Owns the OpenMLS provider, identity, local device,
//! and the set of open conversations.
//!
//! All operations are `async`. The intent is that the FFI generators emit Swift `async`,
//! Kotlin `suspend`, and the WASM glue exposes Promises.

use openmls::framing::MlsMessageOut;
use openmls::prelude::{
    tls_codec::Serialize as TlsSerialize, BasicCredential, Ciphersuite, CredentialWithKey,
    KeyPackageBuilder,
};
use openmls_basic_credential::SignatureKeyPair;
use openmls_traits::OpenMlsProvider;
use parking_lot::RwLock;
use ping_mls_store::{PersistentMlsProvider, StorageBackend};
use std::collections::HashMap;
use std::sync::Arc;
use zeroize::Zeroizing;

use crate::{
    codec,
    conversation::{Conversation, ConversationId, ConversationMeta},
    device::{
        CatchupAppEventEntry, CatchupConversationEntry, CatchupSnapshot, DeviceId, DeviceInfo,
        LinkingTicket, LocalDevice, CATCHUP_SNAPSHOT_VERSION,
    },
    error::{Error, Result},
    identity::{Identity, UserId},
    message::{IncomingMessage, MessageEnvelope, MessageKind},
    storage::Storage,
    sync::SyncCursor,
    transport::Transport,
};

const DEFAULT_CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;

#[derive(Debug)]
pub struct ClientConfig {
    pub identity: Identity,
    pub device_label: String,
    pub storage: Arc<dyn Storage>,
    pub transport: Arc<dyn Transport>,
    /// Wall clock in ms. Pulled from the host so we can use a synthetic clock in tests.
    pub now_ms: u64,
    /// [CR-4] OpenMLS-provider backend. Defaults to in-memory; iOS NSE and web SW
    /// cold-start paths MUST pass `StorageBackend::Sqlite { path, encryption_key }`
    /// (native) or `StorageBackend::IndexedDb { db_name }` (WASM, when that lands).
    /// See `docs/design/CR4_CR7_PERSISTENCE.md`.
    pub storage_backend: StorageBackend,
    /// Optional 32-byte Ed25519 secret key the SDK should use as the
    /// device signing key. When set AND no `LocalDevice` is yet
    /// persisted in `storage`, the SDK constructs its first
    /// `LocalDevice` from this key instead of generating a fresh
    /// random one — so `device_id = SHA-256(public_key_of(secret))`
    /// is fully determined by what the host provided.
    ///
    /// Use case: align the SDK's `device_id` (which it stamps into
    /// every envelope's `sender_device` field) with an externally-
    /// computed device id — typically `SHA-256(device_signing_pubkey)`
    /// in the host's auth layer, where the JWT carries that same
    /// value as its `device_id` claim. Without this alignment, a
    /// server that validates `envelope.sender_device ==
    /// jwt.device_id` would reject every send.
    ///
    /// Ignored on re-init (when storage already has a persisted
    /// `LocalDevice`) so the device identity remains stable across
    /// restarts.
    pub device_signing_secret_key: Option<[u8; 32]>,
}

impl ClientConfig {
    /// Construct a config with `StorageBackend::Memory` — convenient for tests and
    /// the existing v0.1 in-memory flow.
    pub fn new_in_memory(
        identity: Identity,
        device_label: String,
        storage: Arc<dyn Storage>,
        transport: Arc<dyn Transport>,
        now_ms: u64,
    ) -> Self {
        Self {
            identity,
            device_label,
            storage,
            transport,
            now_ms,
            storage_backend: StorageBackend::Memory,
            device_signing_secret_key: None,
        }
    }
}

pub struct MessagingClient {
    pub(crate) identity: Identity,
    pub(crate) local_device: LocalDevice,
    pub(crate) crypto: Arc<PersistentMlsProvider>,
    pub(crate) signing: Arc<SignatureKeyPair>,
    pub(crate) storage: Arc<dyn Storage>,
    pub(crate) transport: Arc<dyn Transport>,
    conversations: RwLock<HashMap<ConversationId, Conversation>>,
}

impl std::fmt::Debug for MessagingClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MessagingClient")
            .field("user_id", &self.identity.user_id().as_hex())
            .field("device_id", &self.local_device.device_id.as_hex())
            .field("conversation_count", &self.conversations.read().len())
            .finish()
    }
}

impl MessagingClient {
    /// Initialise. Creates a new local device if none is recorded in storage; otherwise rehydrates.
    pub async fn init(cfg: ClientConfig) -> Result<Arc<Self>> {
        // [CR-4] OpenMLS provider is now pluggable. For `StorageBackend::Memory` this
        // behaves like the old `OpenMlsRustCrypto::default()`. For `Sqlite`, the
        // working set is hydrated from the on-disk blob; subsequent `checkpoint` calls
        // flush it back. iOS NSE / web SW cold-start lives here.
        let crypto = PersistentMlsProvider::open(cfg.storage_backend.clone())
            .map_err(|e| Error::Storage(format!("provider open: {e}")))?;
        let local_device = match cfg.storage.get("device", "local").await? {
            Some(bytes) => decode_local_device(&bytes, cfg.identity.user_id().clone())?,
            None => {
                // First-init path. If the host supplied a signing secret
                // (typically to align the device_id with their auth
                // layer), use it; otherwise mint a fresh random key.
                // Either way, the constructed `LocalDevice` is
                // immediately persisted so future inits load from
                // storage without consulting the override again.
                let dev = match cfg.device_signing_secret_key.as_ref() {
                    Some(secret) => LocalDevice::from_signing_secret(
                        cfg.identity.user_id().clone(),
                        cfg.device_label,
                        cfg.now_ms,
                        secret,
                    ),
                    None => LocalDevice::generate(
                        cfg.identity.user_id().clone(),
                        cfg.device_label,
                        cfg.now_ms,
                    ),
                };
                let bytes = encode_local_device(&dev)?;
                cfg.storage.put("device", "local", bytes).await?;
                dev
            }
        };

        // [CR-4] MLS signing keypair MUST be stable across cold restarts — otherwise the
        // leaf-key stored on disk no longer matches the per-client key on re-init, and any
        // send-after-restart silently misroutes. We derive deterministically from the
        // already-persistent `LocalDevice::signing` (Ed25519, 32 raw bytes), and the
        // ciphersuite's signature scheme is Ed25519 too — so the device signing key and the
        // MLS leaf signing key are the same bytes. The MLS storage provider also receives
        // a copy via `store()` so OpenMLS-internal lookups (process_message, etc.) succeed.
        let signing = {
            let sk_bytes = local_device.signing.to_bytes().to_vec();
            let pk_bytes = local_device.signing.verifying_key().to_bytes().to_vec();
            let kp = SignatureKeyPair::from_raw(
                DEFAULT_CIPHERSUITE.signature_algorithm(),
                sk_bytes,
                pk_bytes,
            );
            kp.store(crypto.storage()).map_err(Error::mls)?;
            Arc::new(kp)
        };

        let client = Arc::new(Self {
            identity: cfg.identity,
            local_device,
            crypto,
            signing,
            storage: cfg.storage,
            transport: cfg.transport,
            conversations: RwLock::new(HashMap::new()),
        });

        client.rehydrate_conversations(cfg.now_ms).await?;

        // [CR-10] Ensure the DeviceGroup exists at init, not lazily inside
        // build_linking_ticket. Single-device users need somewhere to write
        // personal events (drafts, read pointers, notes, vault wrapper)
        // even before they pair a second device. Lazy creation in
        // build_linking_ticket left them with no DG → no place for
        // personal state to land.
        //
        // Idempotent — re-init after a cold restart finds the DG via
        // rehydrate_conversations and this becomes a no-op.
        client.ensure_device_group(cfg.now_ms).await?;

        Ok(client)
    }

    /// [CR-10] Idempotently ensures this user's DeviceGroup exists in
    /// `self.conversations`. Called from `init` (so single-device users
    /// have a DG immediately) and from `build_linking_ticket` (the legacy
    /// lazy path; still safe to call when the DG already exists, since
    /// rehydrate_conversations would have re-attached it before init
    /// returned).
    ///
    /// The DeviceGroup is a one-leaf MLS group at creation time —
    /// `add_members` (called by `build_linking_ticket` when a second
    /// device pairs in) is what grows it. We persist the snapshot so a
    /// cold restart picks it up before this function runs again.
    pub(crate) async fn ensure_device_group(self: &Arc<Self>, now_ms: u64) -> Result<()> {
        let dg_id = device_group_id_for(self.identity.user_id());
        if self.conversations.read().contains_key(&dg_id) {
            return Ok(());
        }
        let mut new_dg = Conversation::create(
            dg_id,
            Some("device-group".into()),
            self.local_device.device_id.clone(),
            self.identity.user_id(),
            self.crypto.clone(),
            self.signing.clone(),
            self.storage.clone(),
            now_ms,
        )?;
        new_dg.meta.is_device_group = true;
        new_dg.snapshot_to_storage().await?;
        self.conversations.write().insert(dg_id, new_dg);
        Ok(())
    }

    pub fn user_id(&self) -> UserId {
        self.identity.user_id().clone()
    }
    pub fn device_id(&self) -> DeviceId {
        self.local_device.device_id.clone()
    }
    pub fn device_info(&self, now_ms: u64) -> DeviceInfo {
        self.local_device.info(now_ms)
    }

    /// Generate a fresh KeyPackage to publish to the directory. Hosts call this when registering
    /// a device or topping up the directory.
    pub fn fresh_key_package(&self) -> Result<Vec<u8>> {
        let credential_with_key = CredentialWithKey {
            credential: BasicCredential::new(self.identity.user_id().0.clone()).into(),
            signature_key: self.signing.public().to_vec().into(),
        };
        let bundle = KeyPackageBuilder::new()
            .build(
                DEFAULT_CIPHERSUITE,
                self.crypto.as_ref(),
                self.signing.as_ref(),
                credential_with_key,
            )
            .map_err(Error::mls)?;
        // KeyPackages are serialized as MlsMessage(KeyPackage) per the MLS framing spec.
        let msg: MlsMessageOut = bundle.key_package().clone().into();
        msg.tls_serialize_detached().map_err(Error::mls)
    }

    /// Create a new conversation owned by this client (and seeded with a single member: this device).
    pub async fn create_conversation(
        self: &Arc<Self>,
        name: Option<String>,
        now_ms: u64,
    ) -> Result<ConversationId> {
        let id = ConversationId::new();
        let convo = Conversation::create(
            id,
            name,
            self.local_device.device_id.clone(),
            self.identity.user_id(),
            self.crypto.clone(),
            self.signing.clone(),
            self.storage.clone(),
            now_ms,
        )?;
        convo.snapshot_to_storage().await?;
        self.conversations.write().insert(id, convo);
        Ok(id)
    }

    /// Join via a Welcome bundled in a [`MessageEnvelope`] of kind `Welcome`.
    pub async fn join_conversation(
        self: &Arc<Self>,
        welcome_envelope: &MessageEnvelope,
        now_ms: u64,
    ) -> Result<ConversationId> {
        if welcome_envelope.kind != MessageKind::Welcome {
            return Err(Error::Invalid("expected Welcome envelope".into()));
        }
        let convo = Conversation::join(
            &welcome_envelope.payload,
            self.local_device.device_id.clone(),
            self.crypto.clone(),
            self.signing.clone(),
            self.storage.clone(),
            now_ms,
        )?;
        let id = convo.id();
        convo.snapshot_to_storage().await?;
        self.conversations.write().insert(id, convo);
        Ok(id)
    }

    pub fn list_conversations(&self) -> Vec<ConversationMeta> {
        self.conversations
            .read()
            .values()
            .map(|c| c.meta.clone())
            .collect()
    }

    /// Send an application message. Returns once the envelope has been handed to the transport.
    pub async fn send(
        &self,
        conv_id: ConversationId,
        plaintext: Vec<u8>,
        now_ms: u64,
    ) -> Result<MessageEnvelope> {
        let envelope = {
            let mut guard = self.conversations.write();
            let convo = guard
                .get_mut(&conv_id)
                .ok_or_else(|| Error::UnknownConversation(conv_id.as_hex()))?;
            convo.send_application(&plaintext, now_ms)?
        };
        self.transport.send(envelope.clone()).await?;
        Ok(envelope)
    }

    /// Add members. The Commit goes on the wire; the Welcome should be delivered to the new
    /// devices' inboxes (the host transport implements that — typically as a separate addressed
    /// envelope).
    ///
    /// [CR-2] Each entry is `(DeviceId, KeyPackage_bytes)`. The host typically gets the
    /// device_id from the directory at the same time it gets the KeyPackage; we use it to
    /// record a per-conversation `device_id → leaf_index` map so [`Self::revoke_device`]
    /// can later locate the leaf without a fresh directory lookup. The SDK does not
    /// cryptographically verify the host's device-id claim — that's a directory policy
    /// concern.
    //
    // We hold a `parking_lot` read guard across `.await` for `snapshot_to_storage` here. Clippy
    // flags this; we keep it for v0.1 because the alternative is a structural refactor of
    // Conversation::snapshot_to_storage to split sync prep from async writes — see
    // docs/ASSUMPTIONS.md item "lock-during-async-I/O is suboptimal but acceptable for v0.1".
    // The `parking_lot/send_guard` feature (in core/Cargo.toml) makes the guard `Send` so the
    // future is still schedulable across tokio threads.
    #[allow(clippy::await_holding_lock)]
    pub async fn add_members(
        &self,
        conv_id: ConversationId,
        entries: Vec<(DeviceId, Vec<u8>)>,
        now_ms: u64,
    ) -> Result<()> {
        let outcome = {
            let mut guard = self.conversations.write();
            let convo = guard
                .get_mut(&conv_id)
                .ok_or_else(|| Error::UnknownConversation(conv_id.as_hex()))?;
            convo.add_members(entries, now_ms)?
        };
        self.transport.send(outcome.commit).await?;
        self.transport.send(outcome.welcome).await?;
        if let Some(c) = self.conversations.read().get(&conv_id) {
            c.snapshot_to_storage().await?;
        }
        Ok(())
    }

    #[allow(clippy::await_holding_lock)] // see add_members for rationale
    pub async fn remove_members(
        &self,
        conv_id: ConversationId,
        leaf_indexes: Vec<u32>,
        now_ms: u64,
    ) -> Result<()> {
        let envelope = {
            let mut guard = self.conversations.write();
            let convo = guard
                .get_mut(&conv_id)
                .ok_or_else(|| Error::UnknownConversation(conv_id.as_hex()))?;
            convo.remove_members(leaf_indexes, now_ms)?
        };
        self.transport.send(envelope).await?;
        if let Some(c) = self.conversations.read().get(&conv_id) {
            c.snapshot_to_storage().await?;
        }
        Ok(())
    }

    /// Process an inbound envelope coming from the transport's subscribe callback or a sync pull.
    /// Returns `Some` for application traffic, `None` for handshake messages (already merged).
    #[allow(clippy::await_holding_lock)] // see add_members for rationale
    pub async fn process_envelope(
        &self,
        env: &MessageEnvelope,
        now_ms: u64,
    ) -> Result<Option<IncomingMessage>> {
        // Welcome envelopes for unknown conversations are routed to `join_conversation` by the
        // caller. Here we only handle traffic for already-open groups.
        let mut guard = self.conversations.write();
        let convo = match guard.get_mut(&env.conversation_id) {
            Some(c) => c,
            None => return Err(Error::UnknownConversation(env.conversation_id.as_hex())),
        };
        let out = convo.process(env, now_ms)?;
        // Cheap snapshot — only mutates KV the size of the cursor.
        convo.snapshot_to_storage().await?;
        Ok(out)
    }

    /// Catch-up sync: pull missing events for every open conversation since its cursor.
    /// Returns the list of newly-decrypted application messages, in apply order.
    pub async fn sync_conversations(&self, now_ms: u64) -> Result<Vec<IncomingMessage>> {
        let pending: Vec<(ConversationId, SyncCursor)> = self
            .conversations
            .read()
            .iter()
            .map(|(id, c)| (*id, c.cursor.clone()))
            .collect();

        let mut delivered = Vec::new();
        for (conv_id, cursor) in pending {
            loop {
                let batch = self
                    .transport
                    .fetch_since(conv_id, cursor.clone(), 256)
                    .await?;
                if batch.is_empty() {
                    break;
                }
                for env in &batch {
                    if let Some(msg) = self.process_envelope(env, now_ms).await? {
                        delivered.push(msg);
                    }
                }
                if batch.len() < 256 {
                    break;
                } // partial page → caught up
            }
        }
        Ok(delivered)
    }

    /// Rehydrate conversations from storage on startup ([CR-4]).
    ///
    /// Walks the host-side `groups` namespace for meta records, pairs each with its
    /// cursor + device→leaf map, and asks `Conversation::load` to re-attach to the
    /// underlying OpenMLS group state. The MLS state itself was persisted by the
    /// SQLite-backed `PersistentMlsProvider` on the previous run; this method
    /// reconciles the SDK-side caches with what's on disk.
    async fn rehydrate_conversations(self: &Arc<Self>, now_ms: u64) -> Result<()> {
        let metas = self.storage.list_keys("groups", "").await?;
        for path in metas {
            // path looks like "{convId}/meta"
            let Some((id_hex, suffix)) = path.split_once('/') else {
                continue;
            };
            if suffix != "meta" {
                continue;
            }
            let Some(meta_bytes) = self.storage.get("groups", &path).await? else {
                continue;
            };
            let meta: ConversationMeta = match codec::decode(&meta_bytes) {
                Ok(m) => m,
                Err(_) => continue,
            };
            let cursor_bytes = self
                .storage
                .get("cursors", id_hex)
                .await?
                .unwrap_or_default();
            let cursor = if cursor_bytes.is_empty() {
                SyncCursor::default()
            } else {
                SyncCursor::decode(&cursor_bytes).unwrap_or_default()
            };

            // [CR-2] device→leaf map was persisted alongside meta + cursor.
            let device_leaves_bytes = self
                .storage
                .get("device_leaves", id_hex)
                .await?
                .unwrap_or_default();
            let device_leaves: std::collections::BTreeMap<DeviceId, u32> =
                if device_leaves_bytes.is_empty() {
                    std::collections::BTreeMap::new()
                } else {
                    let pairs: Vec<(DeviceId, u32)> =
                        codec::decode(&device_leaves_bytes).unwrap_or_default();
                    pairs.into_iter().collect()
                };

            match Conversation::load(
                meta.id,
                meta.clone(),
                cursor,
                device_leaves,
                self.local_device.device_id.clone(),
                self.crypto.clone(),
                self.signing.clone(),
                self.storage.clone(),
                now_ms,
            ) {
                Ok(Some(convo)) => {
                    tracing::debug!(
                        target: "ping_core::client",
                        convo = %id_hex,
                        epoch = meta.epoch,
                        "rehydrated conversation from disk"
                    );
                    self.conversations.write().insert(meta.id, convo);
                }
                Ok(None) => {
                    tracing::warn!(
                        target: "ping_core::client",
                        convo = %id_hex,
                        "host-side meta present but OpenMLS state missing — skipping"
                    );
                }
                Err(e) => {
                    tracing::warn!(
                        target: "ping_core::client",
                        convo = %id_hex,
                        error = %e,
                        "Conversation::load failed — skipping"
                    );
                }
            }
        }
        Ok(())
    }

    // ------------------- Multi-device API -------------------

    /// Build a [`LinkingTicket`] for a new device. The caller obtains `new_device_kp` from the
    /// new device (e.g., via QR-encoded handshake) and is responsible for sealing the returned
    /// ticket against the new device's ephemeral X25519 pubkey before transmission via
    /// [`ping_link::seal_ticket`].
    ///
    /// [CR-13] `last_app_events` is a host-supplied list of `(conversation_id, app_event_bytes)`
    /// for the new device's "what you missed" UI. The SDK adds its own metas + (currently-
    /// empty) per-conversation MLS state and bundles everything into
    /// [`device::CatchupSnapshot`], CBOR-encoded into the ticket's `catchup_snapshot` field.
    /// Pass an empty `Vec` to suppress catchup data (the new device sees an empty
    /// conversation list until normal sync runs).
    pub async fn build_linking_ticket(
        self: &Arc<Self>,
        new_device_id: DeviceId,
        new_device_kp: Vec<u8>,
        last_app_events: Vec<(ConversationId, Vec<u8>)>,
        now_ms: u64,
    ) -> Result<LinkingTicket> {
        let device_binding_sig = self.identity.sign_device_binding(&new_device_id.0);
        let dg_id = device_group_id_for(self.identity.user_id());

        // [CR-10] DG is eagerly created at init now, but call ensure here too so
        // hosts that bypass `MessagingClient::init` (mocked tests, legacy upgrade
        // paths) keep working.
        self.ensure_device_group(now_ms).await?;

        // Admit the new device to the DeviceGroup.
        let outcome = {
            let mut conversations = self.conversations.write();
            let dg = conversations
                .get_mut(&dg_id)
                .expect("DeviceGroup ensured above");
            // [CR-2] Record the new device's leaf in the DG so future `revoke_device`
            // can find it. The new_device_id we got as a parameter is the inviter's
            // own assertion — same trust model as the rest of `add_members`.
            dg.add_members(vec![(new_device_id.clone(), new_device_kp)], now_ms)?
        };

        // [CR-13] Assemble the catchup snapshot: SDK-known conversation metadata + host-
        // supplied last-known plaintext per conversation. [CR-7] now populates
        // `group_state_bytes` with each group's MLS state so the new device can decrypt
        // historical traffic without re-Welcoming. An empty `group_state_bytes` would
        // mean either a group with no exportable state (shouldn't happen) or an
        // encoder failure (we let those propagate as errors below).
        let catchup_snapshot = if last_app_events.is_empty() && self.conversations.read().is_empty()
        {
            // Cheap path: nothing to snapshot, skip the encode round-trip.
            Vec::new()
        } else {
            let conversation_metas: Vec<CatchupConversationEntry> = self
                .conversations
                .read()
                .values()
                .map(|c| -> Result<CatchupConversationEntry> {
                    // CR-7: per-group state. We deliberately keep the export bytes
                    // inside the (HPKE-sealed-by-CR-3) LinkingTicket; the receiver
                    // calls `import_state_snapshot` with these bytes after `consume_linking_ticket`.
                    let group_bytes = c.export_state_snapshot(now_ms)?.to_vec();
                    Ok(CatchupConversationEntry {
                        conversation_id: c.id(),
                        meta: c.meta().clone(),
                        group_state_bytes: group_bytes,
                    })
                })
                .collect::<Result<_>>()?;
            let last_app_events_per_conv: Vec<CatchupAppEventEntry> = last_app_events
                .into_iter()
                .map(|(conversation_id, app_event_bytes)| CatchupAppEventEntry {
                    conversation_id,
                    app_event_bytes,
                })
                .collect();
            CatchupSnapshot {
                v: CATCHUP_SNAPSHOT_VERSION,
                conversation_metas,
                last_app_events_per_conv,
            }
            .encode()?
        };

        Ok(LinkingTicket {
            v: 1,
            user_id: self.identity.user_id().clone(),
            user_pubkey: self.identity.public_key().to_bytes().to_vec(),
            new_device_id,
            device_binding_sig,
            device_group_welcome: outcome.welcome.payload,
            catchup_snapshot,
        })
    }

    /// Apply a received linking ticket. Joins the user's DeviceGroup; the catch-up snapshot
    /// (if any) is decrypted by the host using the standard per-conversation channel afterwards.
    pub async fn consume_linking_ticket(
        self: &Arc<Self>,
        ticket: &LinkingTicket,
        now_ms: u64,
    ) -> Result<()> {
        // Verify the binding the existing device made for us. (Ed25519 public keys are 32 bytes.)
        let pk_bytes: [u8; 32] = ticket
            .user_pubkey
            .as_slice()
            .try_into()
            .map_err(|_| Error::Identity("user_pubkey must be 32 bytes".into()))?;
        let user_pk = ed25519_dalek::VerifyingKey::from_bytes(&pk_bytes)
            .map_err(|e| Error::Identity(format!("bad user pubkey: {e}")))?;
        Identity::verify_device_binding(
            &user_pk,
            &ticket.user_id,
            &ticket.new_device_id.0,
            &ticket.device_binding_sig,
        )?;
        if ticket.new_device_id != self.local_device.device_id {
            return Err(Error::Invalid(
                "ticket addressed to a different device".into(),
            ));
        }

        let dummy_env = MessageEnvelope::new(
            ConversationId(device_group_id_for(&ticket.user_id).0),
            0,
            MessageKind::Welcome,
            self.local_device.device_id.clone(),
            0,
            crate::clock::Hlc::ZERO,
            ticket.device_group_welcome.clone(),
        );
        self.join_conversation(&dummy_env, now_ms).await?;
        Ok(())
    }

    /// [CR-7] Export the MLS state snapshot for one open conversation.
    ///
    /// Thin pass-through to [`Conversation::export_state_snapshot`]. Returned bytes
    /// are wrapped in `Zeroizing` because they contain past epoch secrets.
    pub fn export_conversation_state_snapshot(
        &self,
        conv_id: ConversationId,
        now_ms: u64,
    ) -> Result<zeroize::Zeroizing<Vec<u8>>> {
        let guard = self.conversations.read();
        let convo = guard
            .get(&conv_id)
            .ok_or_else(|| Error::UnknownConversation(conv_id.as_hex()))?;
        convo.export_state_snapshot(now_ms)
    }

    /// [CR-7] Import a `GroupStateSnapshot` produced by another device's
    /// [`Conversation::export_state_snapshot`].
    ///
    /// Replays the snapshot's entries into this client's OpenMLS provider, then
    /// reconstructs the `Conversation` handle via `MlsGroup::load`. After return,
    /// the conversation is in `list_conversations()` and `send`/`process_envelope`
    /// work against it normally.
    ///
    /// **Scope.** This is for the *same-user* hand-off (linking, recovery). The
    /// snapshot exposes the exporter's view of past epoch secrets for the target
    /// group; only call this when the receiving device has been authenticated to
    /// the same user identity (mnemonic, QR-handshake). Cross-user history transfer
    /// uses HPKE-sealed AppEvent re-shares (umbrella §15.6), not this method.
    ///
    /// **Sanity.** Refuses snapshots whose `group_id` doesn't match the bytes the
    /// receiver intends to claim — guards against host bugs that shuffle snapshots
    /// between groups. Refuses mismatched OpenMLS storage versions outright; no
    /// silent forward/back compatibility.
    pub async fn import_state_snapshot(
        self: &Arc<Self>,
        snapshot_bytes: &[u8],
        now_ms: u64,
    ) -> Result<ConversationId> {
        use crate::device::GroupStateSnapshot;
        let snap = GroupStateSnapshot::decode(snapshot_bytes)
            .map_err(|e| Error::Invalid(format!("snapshot decode: {e}")))?;

        if snap.openmls_storage_version != openmls_traits::storage::CURRENT_VERSION {
            return Err(Error::Invalid(format!(
                "snapshot openmls_storage_version={} not supported (this SDK supports v={})",
                snap.openmls_storage_version,
                openmls_traits::storage::CURRENT_VERSION
            )));
        }

        let conv_id = snap.group_id;

        // Refuse if we already have an active handle for this conv — the host should
        // close it first, otherwise import silently overwrites in-memory state and
        // the existing handle becomes stale.
        if self.conversations.read().contains_key(&conv_id) {
            return Err(Error::Invalid(format!(
                "conversation {} already open; close before importing snapshot",
                conv_id.as_hex()
            )));
        }

        // Replay raw KV pairs into the provider's working set.
        let entries: Vec<(Vec<u8>, Vec<u8>)> =
            snap.entries.into_iter().map(|e| (e.key, e.value)).collect();
        self.crypto
            .import_entries(entries)
            .map_err(|e| Error::Storage(format!("import entries: {e}")))?;

        // Reconstruct the Conversation handle. `Conversation::load` will return
        // `Ok(None)` if OpenMLS still can't find the group — i.e. our snapshot was
        // incomplete or for a different storage version.
        let meta = ConversationMeta {
            id: conv_id,
            name: None,
            epoch: 0, // will be overwritten from the loaded group state in process()
            member_count: 0,
            is_device_group: false, // host can flip this via meta update if needed
            created_at_ms: now_ms,
        };
        let convo = Conversation::load(
            conv_id,
            meta,
            SyncCursor::default(),
            std::collections::BTreeMap::new(),
            self.local_device.device_id.clone(),
            self.crypto.clone(),
            self.signing.clone(),
            self.storage.clone(),
            now_ms,
        )?
        .ok_or_else(|| {
            Error::Invalid(
                "snapshot imported but OpenMLS could not load the group — snapshot may be incomplete or storage version mismatched"
                    .into(),
            )
        })?;

        // Pull the live epoch + member count from the loaded group so the meta we
        // just stubbed is consistent with what we'll observe on subsequent process_envelope.
        let live_epoch = convo.epoch();
        let live_members = convo.group.members().count() as u32;
        let mut convo = convo;
        convo.meta.epoch = live_epoch;
        convo.meta.member_count = live_members;
        convo.snapshot_to_storage().await?;

        self.conversations.write().insert(conv_id, convo);
        Ok(conv_id)
    }

    /// Export a derived secret from one conversation's MLS exporter ([CR-8]).
    ///
    /// Thin pass-through to [`Conversation::export_secret`]. See that method's doc comment
    /// for the contract on `label`, `context`, length validation, and zeroization. The
    /// returned `Zeroizing<Vec<u8>>` is automatically wiped when dropped.
    pub fn export_conversation_secret(
        &self,
        conv_id: ConversationId,
        label: &str,
        context: &[u8],
        length: usize,
    ) -> Result<Zeroizing<Vec<u8>>> {
        let guard = self.conversations.read();
        let convo = guard
            .get(&conv_id)
            .ok_or_else(|| Error::UnknownConversation(conv_id.as_hex()))?;
        convo.export_secret(label, context, length)
    }

    /// Revoke a device by removing its leaf from every conversation where we know its
    /// position ([CR-2]).
    ///
    /// Returns one Commit envelope per conversation the device was a leaf in. The host
    /// broadcasts each envelope to the affected conversation; the SDK has also already
    /// handed them to the transport via `transport.send` (idempotent broadcast is the
    /// host's call).
    ///
    /// **Scope.** The SDK can only resolve leaves it recorded itself — either when it
    /// admitted the device via [`Self::add_members`] or when this device joined as the
    /// target via Welcome. For peer-admitted devices the leaf index isn't locally known;
    /// those conversations are silently skipped. The host can fall back to
    /// `remove_members(leaf_index)` directly using a transport-side directory lookup if
    /// it needs to revoke from those conversations too. See
    /// `docs/architecture/multi-device.md §Device removal` for the broader flow.
    ///
    /// Conversations with no entry for `device_id` produce no envelope; an empty `Vec`
    /// return is a valid outcome (e.g. the device was already revoked, or was never
    /// added by this client).
    #[allow(clippy::await_holding_lock)] // see add_members for rationale
    pub async fn revoke_device(
        &self,
        device_id: DeviceId,
        now_ms: u64,
    ) -> Result<Vec<MessageEnvelope>> {
        // 1. Walk every open conversation and gather (conv_id, leaf_index) pairs where
        //    we know `device_id` controls a leaf. Done under a read lock so we don't hold
        //    the write lock across the per-conversation remove path.
        let targets: Vec<(ConversationId, u32)> = self
            .conversations
            .read()
            .iter()
            .filter_map(|(id, c)| c.leaf_index_of(&device_id).map(|leaf| (*id, leaf)))
            .collect();

        // 2. For each target, emit a remove_members commit. We do this sequentially: each
        //    one is a separate MLS epoch advance on its own group, and they don't share
        //    state, so parallel issuance is safe but adds complexity we don't need for v1.
        let mut envelopes = Vec::with_capacity(targets.len());
        for (conv_id, leaf_index) in targets {
            let envelope = {
                let mut guard = self.conversations.write();
                let convo = guard
                    .get_mut(&conv_id)
                    .ok_or_else(|| Error::UnknownConversation(conv_id.as_hex()))?;
                convo.remove_members(vec![leaf_index], now_ms)?
            };
            self.transport.send(envelope.clone()).await?;
            if let Some(c) = self.conversations.read().get(&conv_id) {
                c.snapshot_to_storage().await?;
            }
            envelopes.push(envelope);
        }

        // 3. Notify the auth-layer server so it can invalidate the
        //    revoked device's KeyPackage pool, mark `auth.devices.revoked_at`,
        //    and refuse any future envelope signed by the revoked device's
        //    JWT. Done AFTER the MLS Commits so peers learn via MLS first
        //    (the canonical path) and the auth layer is the eventual-
        //    consistency cleanup. Transport failures bubble up so callers
        //    can retry — but the MLS-side work has already shipped, so
        //    the device is functionally revoked in every group; only the
        //    auth-layer KeyPackage purge is pending.
        self.transport.revoke_device_remote(device_id).await?;
        Ok(envelopes)
    }
}

fn device_group_id_for(user_id: &UserId) -> ConversationId {
    // Deterministic 16-byte ID derived from the user's id, prefixed so it cannot collide with
    // a randomly-generated ULID in normal use (ULIDs start with a millisecond timestamp).
    let mut bytes = [0u8; 16];
    bytes[0] = 0xFF;
    bytes[1] = 0xDC; // "DeviCe" group sentinel
    let h = codec::sha256(&user_id.0);
    bytes[2..].copy_from_slice(&h[..14]);
    ConversationId(bytes)
}

fn encode_local_device(d: &LocalDevice) -> Result<Vec<u8>> {
    use serde::Serialize;
    #[derive(Serialize)]
    struct Persisted<'a> {
        device_id: &'a DeviceId,
        label: &'a str,
        created_at_ms: u64,
        #[serde(with = "serde_bytes")]
        signing_seed: &'a [u8],
    }
    codec::encode(&Persisted {
        device_id: &d.device_id,
        label: &d.label,
        created_at_ms: d.created_at_ms,
        signing_seed: d.signing.as_bytes(),
    })
}

fn decode_local_device(bytes: &[u8], user_id: UserId) -> Result<LocalDevice> {
    use serde::Deserialize;
    #[derive(Deserialize)]
    struct Persisted {
        device_id: DeviceId,
        label: String,
        created_at_ms: u64,
        #[serde(with = "serde_bytes")]
        signing_seed: Vec<u8>,
    }
    let p: Persisted = codec::decode(bytes)?;
    let seed: [u8; 32] = p
        .signing_seed
        .as_slice()
        .try_into()
        .map_err(|_| Error::Invalid("device signing seed must be 32 bytes".into()))?;
    let signing = ed25519_dalek::SigningKey::from_bytes(&seed);
    Ok(LocalDevice {
        device_id: p.device_id,
        user_id,
        label: p.label,
        signing,
        created_at_ms: p.created_at_ms,
    })
}