quicproquo/crates/quicnprotochat-core/src/group.rs

//! MLS group state machine.
//!
//! # Design
//!
//! [`GroupMember`] wraps an openmls [`MlsGroup`] plus the per-client
//! [`StoreCrypto`] backend. The backend is **persistent** — it holds the
//! in-memory key store that maps init-key references to HPKE private keys.
//! openmls's `new_from_welcome` reads those private keys from the key store to
//! decrypt the Welcome, so the same backend instance must be used from
//! `generate_key_package` through `join_group`.
//!
//! # Wire format
//!
//! All MLS messages are serialised/deserialised using TLS presentation language
//! encoding (`tls_codec`). The resulting byte vectors are what the transport
//! layer (and the Delivery Service) sees.
//!
//! # MLS ciphersuite
//!
//! `MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519` — same as M2.
//!
//! # Ratchet tree
//!
//! `use_ratchet_tree_extension = true` so that the ratchet tree is embedded
//! in Welcome messages. `new_from_welcome` is called with `ratchet_tree = None`;
//! openmls extracts the tree from the Welcome's `GroupInfo` extension.

use std::sync::Arc;

use openmls::prelude::{
    Ciphersuite, Credential, CredentialType, CredentialWithKey, CryptoConfig, GroupId, KeyPackage,
    KeyPackageIn, MlsGroup, MlsGroupConfig, MlsMessageInBody, MlsMessageOut,
    ProcessedMessageContent, ProtocolMessage, ProtocolVersion, TlsDeserializeTrait,
    TlsSerializeTrait,
};
use openmls_traits::OpenMlsCryptoProvider;

use crate::{
    error::CoreError,
    identity::IdentityKeypair,
    keystore::{DiskKeyStore, StoreCrypto},
};

// ── Constants ─────────────────────────────────────────────────────────────────

const CIPHERSUITE: Ciphersuite = Ciphersuite::MLS_128_DHKEMX25519_AES128GCM_SHA256_Ed25519;

// ── GroupMember ───────────────────────────────────────────────────────────────

/// Per-client MLS state: identity keypair, crypto backend, and optional group.
///
/// # Lifecycle
///
/// ```text
/// GroupMember::new(identity)
///   ├─ generate_key_package()   → upload to AS
///   ├─ create_group(group_id)   → become sole member
///   │     └─ add_member(kp)     → invite a peer; returns (commit, welcome)
///   └─ join_group(welcome)      → join after receiving a Welcome
///         ├─ send_message(msg)  → encrypt application data
///         └─ receive_message(b) → decrypt; returns Some(plaintext) or None
/// ```
pub struct GroupMember {
    /// Persistent crypto backend. Holds the in-memory key store with HPKE
    /// private keys created during `generate_key_package`.
    backend: StoreCrypto,
    /// Long-term Ed25519 identity keypair. Also used as the MLS `Signer`.
    identity: Arc<IdentityKeypair>,
    /// Active MLS group, if any.
    group: Option<MlsGroup>,
    /// Shared group configuration (wire format, ratchet tree extension, etc.).
    config: MlsGroupConfig,
}

impl GroupMember {
    /// Create a new `GroupMember` with a fresh crypto backend.
    pub fn new(identity: Arc<IdentityKeypair>) -> Self {
        Self::new_with_state(identity, DiskKeyStore::ephemeral(), None)
    }

    /// Create a `GroupMember` from pre-existing state (identity + optional group + store).
    pub fn new_with_state(
        identity: Arc<IdentityKeypair>,
        key_store: DiskKeyStore,
        group: Option<MlsGroup>,
    ) -> Self {
        let config = MlsGroupConfig::builder()
            .use_ratchet_tree_extension(true)
            .build();

        Self {
            backend: StoreCrypto::new(key_store),
            identity,
            group,
            config,
        }
    }

    // ── KeyPackage ────────────────────────────────────────────────────────────

    /// Generate a fresh single-use MLS KeyPackage.
    ///
    /// The HPKE init private key is stored in `self.backend`'s key store.
    /// **The same `GroupMember` instance must later call `join_group`** so
    /// that `new_from_welcome` can retrieve the private key.
    ///
    /// # Returns
    ///
    /// TLS-encoded KeyPackage bytes, ready for upload to the Authentication
    /// Service.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if openmls fails to create the KeyPackage.
    pub fn generate_key_package(&mut self) -> Result<Vec<u8>, CoreError> {
        let credential_with_key = self.make_credential_with_key()?;

        let key_package = KeyPackage::builder()
            .build(
                CryptoConfig::with_default_version(CIPHERSUITE),
                &self.backend,
                self.identity.as_ref(),
                credential_with_key,
            )
            .map_err(|e| CoreError::Mls(format!("{e:?}")))?;

        key_package
            .tls_serialize_detached()
            .map_err(|e| CoreError::Mls(format!("{e:?}")))
    }

    // ── Group creation ────────────────────────────────────────────────────────

    /// Create a new MLS group with `group_id` as the group identifier.
    ///
    /// The caller becomes the sole member (epoch 0). Use `add_member` to
    /// invite additional members.
    ///
    /// `group_id` can be any non-empty byte string; SHA-256 of a human-readable
    /// name is a good choice.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if the group already exists or openmls fails.
    pub fn create_group(&mut self, group_id: &[u8]) -> Result<(), CoreError> {
        let credential_with_key = self.make_credential_with_key()?;
        let mls_id = GroupId::from_slice(group_id);

        let group = MlsGroup::new_with_group_id(
            &self.backend,
            self.identity.as_ref(),
            &self.config,
            mls_id,
            credential_with_key,
        )
        .map_err(|e| CoreError::Mls(format!("{e:?}")))?;

        self.group = Some(group);
        Ok(())
    }

    // ── Membership ────────────────────────────────────────────────────────────

    /// Add a new member by their TLS-encoded KeyPackage bytes.
    ///
    /// Produces a Commit (to update existing members' state) and a Welcome
    /// (to bootstrap the new member). The caller is responsible for
    /// distributing these:
    ///
    /// - Send `commit_bytes` to all **existing** group members via the DS.
    ///   (In the 2-party case where the creator is the only member, this can
    ///   be discarded — the creator applies it locally via this method.)
    /// - Send `welcome_bytes` to the **new** member via the DS.
    ///
    /// This method also merges the pending Commit into the local group state
    /// (advancing the epoch), so the caller is immediately ready to encrypt.
    ///
    /// # Returns
    ///
    /// `(commit_bytes, welcome_bytes)` — both TLS-encoded MLS messages.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if the KeyPackage is malformed, no active
    /// group exists, or openmls fails.
    pub fn add_member(
        &mut self,
        key_package_bytes: &[u8],
    ) -> Result<(Vec<u8>, Vec<u8>), CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        // Deserialise and validate the peer's KeyPackage. KeyPackage only derives
        // TlsSerialize; KeyPackageIn derives TlsDeserialize and provides validate()
        // which verifies the signature and returns a trusted KeyPackage.
        let key_package: KeyPackage =
            KeyPackageIn::tls_deserialize(&mut key_package_bytes.as_ref())
                .map_err(|e| CoreError::Mls(format!("KeyPackage deserialise: {e:?}")))?
                .validate(self.backend.crypto(), ProtocolVersion::Mls10)
                .map_err(|e| CoreError::Mls(format!("KeyPackage validate: {e:?}")))?;

        // Create the Commit + Welcome. The third return value (GroupInfo) is for
        // external commits and is not needed here.
        let (commit_out, welcome_out, _group_info) = group
            .add_members(&self.backend, self.identity.as_ref(), &[key_package])
            .map_err(|e| CoreError::Mls(format!("add_members: {e:?}")))?;

        // Merge the pending Commit into our own state, advancing the epoch.
        group
            .merge_pending_commit(&self.backend)
            .map_err(|e| CoreError::Mls(format!("merge_pending_commit: {e:?}")))?;

        let commit_bytes = commit_out
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("commit serialise: {e:?}")))?;
        let welcome_bytes = welcome_out
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("welcome serialise: {e:?}")))?;

        Ok((commit_bytes, welcome_bytes))
    }

    /// Join an existing MLS group from a TLS-encoded Welcome message.
    ///
    /// The caller must have previously called [`generate_key_package`] on
    /// **this same instance** so that the HPKE init private key is in the
    /// backend's key store.
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if the Welcome does not match any known
    /// KeyPackage, or openmls validation fails.
    ///
    /// [`generate_key_package`]: Self::generate_key_package
    pub fn join_group(&mut self, welcome_bytes: &[u8]) -> Result<(), CoreError> {
        // Deserialise MlsMessageIn, then extract the inner Welcome.
        let msg_in = openmls::prelude::MlsMessageIn::tls_deserialize(&mut welcome_bytes.as_ref())
            .map_err(|e| CoreError::Mls(format!("Welcome deserialise: {e:?}")))?;

        // into_welcome() is feature-gated in openmls 0.5; extract() is public.
        let welcome = match msg_in.extract() {
            MlsMessageInBody::Welcome(w) => w,
            _ => return Err(CoreError::Mls("expected a Welcome message".into())),
        };

        // ratchet_tree = None because use_ratchet_tree_extension = true embeds
        // the tree inside the Welcome's GroupInfo extension.
        let group = MlsGroup::new_from_welcome(&self.backend, &self.config, welcome, None)
            .map_err(|e| CoreError::Mls(format!("new_from_welcome: {e:?}")))?;

        self.group = Some(group);
        Ok(())
    }

    // ── Application messages ──────────────────────────────────────────────────

    /// Encrypt `plaintext` as an MLS Application message.
    ///
    /// # Returns
    ///
    /// TLS-encoded `MlsMessageOut` bytes (PrivateMessage variant).
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if there is no active group or encryption fails.
    pub fn send_message(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let mls_msg: MlsMessageOut = group
            .create_message(&self.backend, self.identity.as_ref(), plaintext)
            .map_err(|e| CoreError::Mls(format!("create_message: {e:?}")))?;

        mls_msg
            .to_bytes()
            .map_err(|e| CoreError::Mls(format!("message serialise: {e:?}")))
    }

    /// Process an incoming TLS-encoded MLS message.
    ///
    /// # Returns
    ///
    /// - `Ok(Some(plaintext))` for Application messages.
    /// - `Ok(None)` for Commit messages (group state is updated internally).
    ///
    /// # Errors
    ///
    /// Returns [`CoreError::Mls`] if the message is malformed, fails
    /// authentication, or the group state is inconsistent.
    pub fn receive_message(&mut self, bytes: &[u8]) -> Result<Option<Vec<u8>>, CoreError> {
        let group = self
            .group
            .as_mut()
            .ok_or_else(|| CoreError::Mls("no active group".into()))?;

        let msg_in = openmls::prelude::MlsMessageIn::tls_deserialize(&mut bytes.as_ref())
            .map_err(|e| CoreError::Mls(format!("message deserialise: {e:?}")))?;

        // into_protocol_message() is feature-gated; extract() + manual construction is not.
        let protocol_message = match msg_in.extract() {
            MlsMessageInBody::PrivateMessage(m) => ProtocolMessage::PrivateMessage(m),
            MlsMessageInBody::PublicMessage(m) => ProtocolMessage::PublicMessage(m),
            _ => return Err(CoreError::Mls("not a protocol message".into())),
        };

        let processed = group
            .process_message(&self.backend, protocol_message)
            .map_err(|e| CoreError::Mls(format!("process_message: {e:?}")))?;

        match processed.into_content() {
            ProcessedMessageContent::ApplicationMessage(app) => Ok(Some(app.into_bytes())),
            ProcessedMessageContent::StagedCommitMessage(staged) => {
                // Merge the Commit into the local state (epoch advances).
                group
                    .merge_staged_commit(&self.backend, *staged)
                    .map_err(|e| CoreError::Mls(format!("merge_staged_commit: {e:?}")))?;
                Ok(None)
            }
            // Proposals are stored for a later Commit; nothing to return yet.
            ProcessedMessageContent::ProposalMessage(proposal) => {
                group.store_pending_proposal(*proposal);
                Ok(None)
            }
            ProcessedMessageContent::ExternalJoinProposalMessage(proposal) => {
                group.store_pending_proposal(*proposal);
                Ok(None)
            }
        }
    }

    // ── Accessors ─────────────────────────────────────────────────────────────

    /// Return the MLS group ID bytes, or `None` if no group is active.
    pub fn group_id(&self) -> Option<Vec<u8>> {
        self.group
            .as_ref()
            .map(|g| g.group_id().as_slice().to_vec())
    }

    /// Return a reference to the identity keypair.
    pub fn identity(&self) -> &IdentityKeypair {
        &self.identity
    }

    /// Return the private seed of the identity (for persistence).
    pub fn identity_seed(&self) -> [u8; 32] {
        self.identity.seed_bytes()
    }

    /// Return a reference to the underlying crypto backend.
    pub fn backend(&self) -> &StoreCrypto {
        &self.backend
    }

    /// Return a reference to the MLS group, if active.
    pub fn group_ref(&self) -> Option<&MlsGroup> {
        self.group.as_ref()
    }

    /// Return the identity (credential) bytes of all current group members.
    ///
    /// Each entry is the raw credential payload (Ed25519 public key bytes)
    /// extracted from the member's MLS leaf node.
    pub fn member_identities(&self) -> Vec<Vec<u8>> {
        let group = match self.group.as_ref() {
            Some(g) => g,
            None => return Vec::new(),
        };
        group
            .members()
            .map(|m| m.credential.identity().to_vec())
            .collect()
    }

    // ── Private helpers ───────────────────────────────────────────────────────

    fn make_credential_with_key(&self) -> Result<CredentialWithKey, CoreError> {
        let credential = Credential::new(
            self.identity.public_key_bytes().to_vec(),
            CredentialType::Basic,
        )
        .map_err(|e| CoreError::Mls(format!("{e:?}")))?;

        Ok(CredentialWithKey {
            credential,
            signature_key: self.identity.public_key_bytes().to_vec().into(),
        })
    }
}

// ── Unit tests ────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    /// Full two-party MLS round-trip: create group → add member → exchange messages.
    #[test]
    fn two_party_mls_round_trip() {
        let alice_id = Arc::new(IdentityKeypair::generate());
        let bob_id = Arc::new(IdentityKeypair::generate());

        let mut alice = GroupMember::new(Arc::clone(&alice_id));
        let mut bob = GroupMember::new(Arc::clone(&bob_id));

        // Bob generates a KeyPackage (stored in bob's backend key store).
        let bob_kp = bob.generate_key_package().expect("Bob KeyPackage");

        // Alice creates the group.
        alice
            .create_group(b"test-group-m3")
            .expect("Alice create group");

        // Alice adds Bob → (commit, welcome).
        // Alice is the sole existing member, so she merges the commit herself.
        let (_, welcome) = alice.add_member(&bob_kp).expect("Alice add Bob");

        // Bob joins via the Welcome. His backend holds the matching init key.
        bob.join_group(&welcome).expect("Bob join group");

        // Alice → Bob: application message.
        let ct_a = alice.send_message(b"hello bob").expect("Alice send");
        let pt_b = bob
            .receive_message(&ct_a)
            .expect("Bob recv")
            .expect("should be application message");
        assert_eq!(pt_b, b"hello bob");

        // Bob → Alice: reply.
        let ct_b = bob.send_message(b"hello alice").expect("Bob send");
        let pt_a = alice
            .receive_message(&ct_b)
            .expect("Alice recv")
            .expect("should be application message");
        assert_eq!(pt_a, b"hello alice");
    }

    /// `group_id()` returns None before create_group, Some afterwards.
    #[test]
    fn group_id_lifecycle() {
        let id = Arc::new(IdentityKeypair::generate());
        let mut member = GroupMember::new(id);

        assert!(member.group_id().is_none(), "no group before create");
        member.create_group(b"gid").unwrap();
        assert_eq!(
            member.group_id().unwrap(),
            b"gid".as_slice(),
            "group_id must match what was passed"
        );
    }
}