quicproquo/crates/quicproquo-core/src/app_message.rs

//! Rich application-layer message format for MLS application payloads.
//!
//! The server sees only opaque ciphertext; structure lives in this client-defined
//! plaintext schema. All messages use: version byte (1) + message_type byte + type-specific payload.
//!
//! # Message ID
//!
//! `message_id` is assigned by the sender (16 random bytes) and included in the
//! serialized payload for Chat (and implied for Reply/Reaction/ReadReceipt via ref_msg_id).
//! Recipients can store message_ids to reference them in replies or reactions.

use crate::error::CoreError;
use rand::RngCore;

/// Current schema version.
pub const VERSION: u8 = 1;

/// Message type discriminant (one byte).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum MessageType {
    Chat = 0x01,
    Reply = 0x02,
    Reaction = 0x03,
    ReadReceipt = 0x04,
    Typing = 0x05,
}

impl MessageType {
    fn from_byte(b: u8) -> Option<Self> {
        match b {
            0x01 => Some(MessageType::Chat),
            0x02 => Some(MessageType::Reply),
            0x03 => Some(MessageType::Reaction),
            0x04 => Some(MessageType::ReadReceipt),
            0x05 => Some(MessageType::Typing),
            _ => None,
        }
    }
}

/// Parsed application message (one of the rich types).
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AppMessage {
    /// Plain chat: body (UTF-8). message_id is included so recipients can store and reference it.
    Chat {
        message_id: [u8; 16],
        body: Vec<u8>,
    },
    Reply {
        ref_msg_id: [u8; 16],
        body: Vec<u8>,
    },
    Reaction {
        ref_msg_id: [u8; 16],
        emoji: Vec<u8>,
    },
    ReadReceipt {
        msg_id: [u8; 16],
    },
    Typing {
        /// 0 = stopped, 1 = typing
        active: u8,
    },
}

/// Generate a new 16-byte message ID (e.g. for Chat/Reply so recipients can reference it).
pub fn generate_message_id() -> [u8; 16] {
    let mut id = [0u8; 16];
    rand::rngs::OsRng.fill_bytes(&mut id);
    id
}

// ── Layout (minimal, no Cap'n Proto) ─────────────────────────────────────────
//
// All messages: [version: 1][type: 1][payload...]
//
// Chat:     [msg_id: 16][body_len: 2 BE][body]
// Reply:    [ref_msg_id: 16][body_len: 2 BE][body]
// Reaction: [ref_msg_id: 16][emoji_len: 1][emoji]
// ReadReceipt: [msg_id: 16]
// Typing:   [active: 1]  0 = stopped, 1 = typing

/// Serialize a rich message into the application payload format.
pub fn serialize(msg_type: MessageType, payload: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(2 + payload.len());
    out.push(VERSION);
    out.push(msg_type as u8);
    out.extend_from_slice(payload);
    out
}

/// Serialize a Chat message (generates message_id internally; pass None to generate, or Some(id) when replying with a known id).
pub fn serialize_chat(body: &[u8], message_id: Option<[u8; 16]>) -> Result<Vec<u8>, CoreError> {
    if body.len() > u16::MAX as usize {
        return Err(CoreError::AppMessage("chat body exceeds maximum length (65535 bytes)".into()));
    }
    let id = message_id.unwrap_or_else(generate_message_id);
    let mut payload = Vec::with_capacity(16 + 2 + body.len());
    payload.extend_from_slice(&id);
    payload.extend_from_slice(&(body.len() as u16).to_be_bytes());
    payload.extend_from_slice(body);
    Ok(serialize(MessageType::Chat, &payload))
}

/// Serialize a Reply message.
pub fn serialize_reply(ref_msg_id: [u8; 16], body: &[u8]) -> Result<Vec<u8>, CoreError> {
    if body.len() > u16::MAX as usize {
        return Err(CoreError::AppMessage("reply body exceeds maximum length (65535 bytes)".into()));
    }
    let mut payload = Vec::with_capacity(16 + 2 + body.len());
    payload.extend_from_slice(&ref_msg_id);
    payload.extend_from_slice(&(body.len() as u16).to_be_bytes());
    payload.extend_from_slice(body);
    Ok(serialize(MessageType::Reply, &payload))
}

/// Serialize a Reaction message.
pub fn serialize_reaction(ref_msg_id: [u8; 16], emoji: &[u8]) -> Result<Vec<u8>, CoreError> {
    if emoji.len() > 255 {
        return Err(CoreError::AppMessage("emoji length > 255".into()));
    }
    let mut payload = Vec::with_capacity(16 + 1 + emoji.len());
    payload.extend_from_slice(&ref_msg_id);
    payload.push(emoji.len() as u8);
    payload.extend_from_slice(emoji);
    Ok(serialize(MessageType::Reaction, &payload))
}

/// Serialize a ReadReceipt message.
pub fn serialize_read_receipt(msg_id: [u8; 16]) -> Vec<u8> {
    serialize(MessageType::ReadReceipt, &msg_id)
}

/// Serialize a Typing message (active: 0 = stopped, 1 = typing).
pub fn serialize_typing(active: u8) -> Vec<u8> {
    let payload = [active];
    serialize(MessageType::Typing, &payload)
}

/// Parse bytes into (MessageType, AppMessage). Fails if version/type unknown or payload too short.
pub fn parse(bytes: &[u8]) -> Result<(MessageType, AppMessage), CoreError> {
    if bytes.len() < 2 {
        return Err(CoreError::AppMessage("payload too short (need version + type)".into()));
    }
    let version = bytes[0];
    if version != VERSION {
        return Err(CoreError::AppMessage(format!("unsupported version {version}")));
    }
    let msg_type = MessageType::from_byte(bytes[1])
        .ok_or_else(|| CoreError::AppMessage(format!("unknown message type {}", bytes[1])))?;
    let payload = &bytes[2..];

    let app = match msg_type {
        MessageType::Chat => parse_chat(payload)?,
        MessageType::Reply => parse_reply(payload)?,
        MessageType::Reaction => parse_reaction(payload)?,
        MessageType::ReadReceipt => parse_read_receipt(payload)?,
        MessageType::Typing => parse_typing(payload)?,
    };
    Ok((msg_type, app))
}

fn parse_chat(payload: &[u8]) -> Result<AppMessage, CoreError> {
    if payload.len() < 16 + 2 {
        return Err(CoreError::AppMessage("Chat payload too short".into()));
    }
    let mut message_id = [0u8; 16];
    message_id.copy_from_slice(&payload[..16]);
    let body_len = u16::from_be_bytes([payload[16], payload[17]]) as usize;
    if payload.len() < 18 + body_len {
        return Err(CoreError::AppMessage("Chat body length exceeds payload".into()));
    }
    let body = payload[18..18 + body_len].to_vec();
    Ok(AppMessage::Chat { message_id, body })
}

fn parse_reply(payload: &[u8]) -> Result<AppMessage, CoreError> {
    if payload.len() < 16 + 2 {
        return Err(CoreError::AppMessage("Reply payload too short".into()));
    }
    let mut ref_msg_id = [0u8; 16];
    ref_msg_id.copy_from_slice(&payload[..16]);
    let body_len = u16::from_be_bytes([payload[16], payload[17]]) as usize;
    if payload.len() < 18 + body_len {
        return Err(CoreError::AppMessage("Reply body length exceeds payload".into()));
    }
    let body = payload[18..18 + body_len].to_vec();
    Ok(AppMessage::Reply { ref_msg_id, body })
}

fn parse_reaction(payload: &[u8]) -> Result<AppMessage, CoreError> {
    if payload.len() < 16 + 1 {
        return Err(CoreError::AppMessage("Reaction payload too short".into()));
    }
    let mut ref_msg_id = [0u8; 16];
    ref_msg_id.copy_from_slice(&payload[..16]);
    let emoji_len = payload[16] as usize;
    if payload.len() < 17 + emoji_len {
        return Err(CoreError::AppMessage("Reaction emoji length exceeds payload".into()));
    }
    let emoji = payload[17..17 + emoji_len].to_vec();
    Ok(AppMessage::Reaction { ref_msg_id, emoji })
}

fn parse_read_receipt(payload: &[u8]) -> Result<AppMessage, CoreError> {
    if payload.len() < 16 {
        return Err(CoreError::AppMessage("ReadReceipt payload too short".into()));
    }
    let mut msg_id = [0u8; 16];
    msg_id.copy_from_slice(&payload[..16]);
    Ok(AppMessage::ReadReceipt { msg_id })
}

fn parse_typing(payload: &[u8]) -> Result<AppMessage, CoreError> {
    if payload.is_empty() {
        return Err(CoreError::AppMessage("Typing payload empty".into()));
    }
    Ok(AppMessage::Typing { active: payload[0] })
}

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

    #[test]
    fn roundtrip_chat() {
        let body = b"hello";
        let encoded = serialize_chat(body, None).unwrap();
        let (t, msg) = parse(&encoded).unwrap();
        assert_eq!(t, MessageType::Chat);
        match &msg {
            AppMessage::Chat { message_id: _, body: b } => assert_eq!(b.as_slice(), body),
            _ => panic!("expected Chat"),
        }
    }

    #[test]
    fn roundtrip_reply() {
        let ref_id = [1u8; 16];
        let body = b"reply text";
        let encoded = serialize_reply(ref_id, body).unwrap();
        let (t, msg) = parse(&encoded).unwrap();
        assert_eq!(t, MessageType::Reply);
        match &msg {
            AppMessage::Reply { ref_msg_id, body: b } => {
                assert_eq!(ref_msg_id, &ref_id);
                assert_eq!(b.as_slice(), body);
            }
            _ => panic!("expected Reply"),
        }
    }

    #[test]
    fn roundtrip_typing() {
        let encoded = serialize_typing(1);
        let (t, msg) = parse(&encoded).unwrap();
        assert_eq!(t, MessageType::Typing);
        match &msg {
            AppMessage::Typing { active } => assert_eq!(*active, 1),
            _ => panic!("expected Typing"),
        }
    }

    #[test]
    fn roundtrip_reaction() {
        let ref_id = [2u8; 16];
        let emoji = "\u{1f44d}".as_bytes();
        let encoded = serialize_reaction(ref_id, emoji).unwrap();
        let (t, msg) = parse(&encoded).unwrap();
        assert_eq!(t, MessageType::Reaction);
        match &msg {
            AppMessage::Reaction { ref_msg_id, emoji: e } => {
                assert_eq!(ref_msg_id, &ref_id);
                assert_eq!(e.as_slice(), emoji);
            }
            _ => panic!("expected Reaction"),
        }
    }

    #[test]
    fn roundtrip_read_receipt() {
        let msg_id = [3u8; 16];
        let encoded = serialize_read_receipt(msg_id);
        let (t, msg) = parse(&encoded).unwrap();
        assert_eq!(t, MessageType::ReadReceipt);
        match &msg {
            AppMessage::ReadReceipt { msg_id: id } => assert_eq!(id, &msg_id),
            _ => panic!("expected ReadReceipt"),
        }
    }

    #[test]
    fn parse_empty_fails() {
        assert!(parse(&[]).is_err());
    }

    #[test]
    fn parse_bad_version_fails() {
        assert!(parse(&[99, 0x01]).is_err());
    }

    #[test]
    fn parse_bad_type_fails() {
        assert!(parse(&[1, 0xFF]).is_err());
    }

    #[test]
    fn chat_body_too_long() {
        let body = vec![0u8; 65536]; // exceeds u16::MAX
        assert!(serialize_chat(&body, None).is_err());
    }

    #[test]
    fn reaction_emoji_too_long() {
        let emoji = vec![0u8; 256];
        assert!(serialize_reaction([0; 16], &emoji).is_err());
    }

    #[test]
    fn parse_truncated_chat_payload() {
        // Version + type + only 10 bytes of payload (needs 18 minimum for chat)
        let mut data = vec![1, 0x01];
        data.extend_from_slice(&[0u8; 10]);
        assert!(parse(&data).is_err());
    }
}