ping-openmls-sdk-core 0.1.4

//! `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_rust_crypto::OpenMlsRustCrypto;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;

use crate::{
    codec,
    conversation::{Conversation, ConversationId, ConversationMeta},
    device::{DeviceId, DeviceInfo, LinkingTicket, LocalDevice},
    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,
}

pub struct MessagingClient {
    pub(crate) identity: Identity,
    pub(crate) local_device: LocalDevice,
    pub(crate) crypto: Arc<OpenMlsRustCrypto>,
    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>> {
        let crypto = Arc::new(OpenMlsRustCrypto::default());
        let local_device = match cfg.storage.get("device", "local").await? {
            Some(bytes) => decode_local_device(&bytes, cfg.identity.user_id().clone())?,
            None => {
                let dev = 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
            }
        };

        // OpenMLS signing keypair — derived from the device's signing key for now. In a richer
        // build we'd separate the per-epoch leaf key from the device key; for v1 we reuse.
        let signing = Arc::new(
            SignatureKeyPair::new(DEFAULT_CIPHERSUITE.signature_algorithm()).map_err(Error::mls)?,
        );

        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?;
        Ok(client)
    }

    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).
    //
    // 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,
        key_packages: Vec<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(key_packages, 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. The OpenMLS provider has its own
    /// keystore; this method just rebuilds the per-conversation cursor + meta cache.
    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()
            };

            // We do NOT rebuild the MLS group here; OpenMLS owns its keystore. Operations against
            // a "cold" conversation will lazily resurrect it on first use. (In v0.1 we treat the
            // OpenMLS provider as the source of truth and only persist meta+cursor at this layer.)
            // A production build would extend this with a binary group snapshot on every commit
            // for faster cold start; deliberately deferred — see ARCHITECTURE.md.
            tracing::debug!(target: "ping_core::client", convo = %id_hex, epoch = meta.epoch, "rehydrated conversation meta");
            let _ = (meta, cursor); // placeholder until cold-restore lands
        }
        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.
    ///
    /// In v0.1 we focus on building the ticket; HPKE sealing is left to the caller (a thin
    /// helper crate, `ping-link`, will provide it). This keeps the core crate free of HPKE
    /// machinery that's already easy to do at the binding layer.
    pub async fn build_linking_ticket(
        &self,
        new_device_id: DeviceId,
        new_device_kp: 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());

        // Bootstrap-or-fetch + admit, all under one write lock so the borrow lifetimes are
        // straightforward.
        let outcome = {
            use std::collections::hash_map::Entry;
            let mut conversations = self.conversations.write();
            if let Entry::Vacant(e) = conversations.entry(dg_id) {
                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;
                e.insert(new_dg);
            }
            let dg = conversations
                .get_mut(&dg_id)
                .expect("DeviceGroup just inserted or already present");
            dg.add_members(vec![new_device_kp], now_ms)?
        };

        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: Vec::new(), // populated by host: list of conv metas + last 200 msgs
        })
    }

    /// 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(())
    }

    /// Revoke a device by removing its leaf from every conversation we participate in.
    pub async fn revoke_device(&self, _device_id: DeviceId, _now_ms: u64) -> Result<()> {
        // For each conversation, look up the leaf index whose credential identity matches the
        // device's public key, and remove. v0.1 leaves the leaf-by-device-id resolver as a TODO
        // because OpenMLS exposes leaves keyed by `LeafNodeIndex` and we'd need to walk
        // `members()` to find a match — straightforward but boilerplate.
        Err(Error::Invalid(
            "revoke_device: not implemented in v0.1 — see roadmap".into(),
        ))
    }
}

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,
    })
}