ping-openmls-sdk-core 0.1.4

//! 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, MlsMessageBodyIn, MlsMessageIn,
        ProcessedMessage, ProtocolMessage, ProtocolVersion,
    },
};
use openmls_basic_credential::SignatureKeyPair;
use openmls_rust_crypto::OpenMlsRustCrypto;
use openmls_traits::OpenMlsProvider;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use ulid::Ulid;

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

const DEFAULT_CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;

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

/// 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<OpenMlsRustCrypto>,
    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>,
}

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
    }
    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<OpenMlsRustCrypto>,
        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(),
        };
        let cfg = MlsGroupCreateConfig::builder()
            .ciphersuite(DEFAULT_CIPHERSUITE)
            .use_ratchet_tree_extension(true)
            .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,
        };
        Ok(Self {
            id,
            meta,
            group,
            crypto,
            signing,
            own_device,
            seq: 0,
            hlc: Hlc::ZERO.tick(now_ms),
            cursor: SyncCursor::default(),
            storage,
        })
    }

    /// Join an existing conversation from a Welcome message.
    pub(crate) fn join(
        welcome_bytes: &[u8],
        own_device: DeviceId,
        crypto: Arc<OpenMlsRustCrypto>,
        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);
        let meta = ConversationMeta {
            id,
            name: None,
            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();
        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,
        })
    }

    /// Encrypt an application message and produce a wire envelope ready for transport.
    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(
            self.id,
            self.epoch(),
            MessageKind::Application,
            self.own_device.clone(),
            self.seq,
            self.hlc,
            bytes,
        );
        // 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.
    pub fn add_members(&mut self, key_packages: Vec<Vec<u8>>, now_ms: u64) -> Result<AddOutcome> {
        let mut kps = Vec::with_capacity(key_packages.len());
        for raw in &key_packages {
            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)?;
            kps.push(kp);
        }

        // The Commit's wire `epoch` field is the *source* epoch (the epoch the Commit was
        // crafted in, matching the epoch embedded in the inner MLS message bytes). The
        // Welcome's `epoch` is the *post-commit* epoch (it carries the new group state).
        // This split is what lets a joiner's sync cursor correctly filter pre-join Commits.
        let pre_commit_epoch = self.epoch();

        let (commit_out, welcome_out, _gi) = self
            .group
            .add_members(self.crypto.as_ref(), self.signing.as_ref(), &kps)
            .map_err(Error::mls)?;

        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;

        self.seq += 1;
        self.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(),
            self.seq,
            self.hlc,
            commit_bytes,
        );

        let welcome_bytes = mls_message_out_bytes(welcome_out)?;
        let welcome_env = MessageEnvelope::new(
            self.id,
            self.meta.epoch,
            MessageKind::Welcome,
            self.own_device.clone(),
            self.seq,
            self.hlc,
            welcome_bytes,
        );

        // Advance the local cursor past our own Commit (at the post-commit epoch, since
        // we've already merged it locally) so catch-up sync doesn't try to re-apply it.
        self.cursor.advance(
            self.meta.epoch,
            self.own_device.clone(),
            self.seq,
            self.hlc,
            now_ms,
        );

        Ok(AddOutcome {
            commit: commit_env,
            welcome: welcome_env,
        })
    }

    pub fn remove_members(
        &mut self,
        leaf_indexes: Vec<u32>,
        now_ms: u64,
    ) -> Result<MessageEnvelope> {
        use openmls::prelude::LeafNodeIndex;
        let leaves: Vec<LeafNodeIndex> = leaf_indexes.into_iter().map(LeafNodeIndex::new).collect();

        // Capture the source epoch before merge — see add_members for the rationale.
        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)?;
        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;

        self.seq += 1;
        self.hlc = self.hlc.tick(now_ms);
        let bytes = mls_message_out_bytes(commit_out)?;
        let env = MessageEnvelope::new(
            self.id,
            pre_commit_epoch,
            MessageKind::Commit,
            self.own_device.clone(),
            self.seq,
            self.hlc,
            bytes,
        );
        // Advance the local cursor past our own Commit (at the post-commit epoch we've just
        // merged into) so catch-up sync doesn't try to re-apply it.
        self.cursor.advance(
            self.meta.epoch,
            self.own_device.clone(),
            self.seq,
            self.hlc,
            now_ms,
        );
        Ok(env)
    }

    /// 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)?;

        let out = match processed.into_content() {
            ProcessedMessageContent::ApplicationMessage(app) => {
                let pt = app.into_bytes();
                Some(IncomingMessage {
                    conversation_id: self.id,
                    sender_device: env.sender_device.clone(),
                    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)
    }

    pub(crate) async fn snapshot_to_storage(&self) -> Result<()> {
        let blob = self
            .group
            .export_secret(self.crypto.as_ref(), "ping-snapshot-marker", &[], 32)
            .ok();
        // OpenMLS persists the group via its own keystore inside `crypto`. We only need to
        // record meta + cursor here; the group itself is recovered by re-opening with the
        // same provider on next launch.
        let _ = blob; // intentionally unused — present for future binary-snapshot path
        let cursor = self.cursor.encode()?;
        self.storage
            .put("cursors", &self.id.as_hex(), cursor)
            .await?;
        let meta = codec::encode(&self.meta)?;
        self.storage
            .put("groups", &format!("{}/meta", self.id.as_hex()), meta)
            .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,
}

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