quicproquo/crates/quicprochat-p2p/src/mesh_protocol.rs

//! Mesh protocol messages for peer-to-peer communication.
//!
//! This module defines the control messages used for mesh coordination:
//! - KeyPackage request/response for MLS group setup
//! - Future: route requests, capability queries, etc.

use serde::{Deserialize, Serialize};

use crate::address::MeshAddress;

/// Protocol message type discriminator.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[repr(u8)]
pub enum MessageType {
    /// Request a KeyPackage from a node.
    KeyPackageRequest = 0x10,
    /// Response with KeyPackage data.
    KeyPackageResponse = 0x11,
    /// Node has no KeyPackage available.
    KeyPackageUnavailable = 0x12,
}

/// Request a KeyPackage from a peer.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct KeyPackageRequest {
    /// Who is requesting.
    pub requester_addr: MeshAddress,
    /// Whose KeyPackage is requested.
    pub target_addr: MeshAddress,
    /// Optional: specific hash to request (from announce).
    pub hash: Option<[u8; 8]>,
    /// Request ID for correlation.
    pub request_id: u32,
}

impl KeyPackageRequest {
    /// Create a new request.
    pub fn new(requester: MeshAddress, target: MeshAddress) -> Self {
        Self {
            requester_addr: requester,
            target_addr: target,
            hash: None,
            request_id: rand::random(),
        }
    }

    /// Create with specific hash.
    pub fn with_hash(requester: MeshAddress, target: MeshAddress, hash: [u8; 8]) -> Self {
        Self {
            requester_addr: requester,
            target_addr: target,
            hash: Some(hash),
            request_id: rand::random(),
        }
    }

    /// Serialize to CBOR.
    pub fn to_wire(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        buf.push(MessageType::KeyPackageRequest as u8);
        ciborium::into_writer(self, &mut buf).expect("CBOR serialization");
        buf
    }

    /// Deserialize from CBOR (after type byte).
    pub fn from_wire(bytes: &[u8]) -> anyhow::Result<Self> {
        if bytes.is_empty() || bytes[0] != MessageType::KeyPackageRequest as u8 {
            anyhow::bail!("not a KeyPackageRequest");
        }
        let req: Self = ciborium::from_reader(&bytes[1..])?;
        Ok(req)
    }
}

/// Response with KeyPackage data.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct KeyPackageResponse {
    /// Whose KeyPackage this is.
    pub owner_addr: MeshAddress,
    /// The serialized MLS KeyPackage.
    pub keypackage_bytes: Vec<u8>,
    /// Hash of the KeyPackage (for verification).
    pub hash: [u8; 8],
    /// Matching request ID.
    pub request_id: u32,
}

impl KeyPackageResponse {
    /// Create a new response.
    pub fn new(
        owner: MeshAddress,
        keypackage_bytes: Vec<u8>,
        request_id: u32,
    ) -> Self {
        let hash = crate::announce::compute_keypackage_hash(&keypackage_bytes);
        Self {
            owner_addr: owner,
            keypackage_bytes,
            hash,
            request_id,
        }
    }

    /// Serialize to CBOR.
    pub fn to_wire(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        buf.push(MessageType::KeyPackageResponse as u8);
        ciborium::into_writer(self, &mut buf).expect("CBOR serialization");
        buf
    }

    /// Deserialize from CBOR (after type byte).
    pub fn from_wire(bytes: &[u8]) -> anyhow::Result<Self> {
        if bytes.is_empty() || bytes[0] != MessageType::KeyPackageResponse as u8 {
            anyhow::bail!("not a KeyPackageResponse");
        }
        let resp: Self = ciborium::from_reader(&bytes[1..])?;
        Ok(resp)
    }

    /// Verify the hash matches the KeyPackage.
    pub fn verify_hash(&self) -> bool {
        let computed = crate::announce::compute_keypackage_hash(&self.keypackage_bytes);
        computed == self.hash
    }
}

/// Response indicating no KeyPackage available.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct KeyPackageUnavailable {
    /// Whose KeyPackage was requested.
    pub target_addr: MeshAddress,
    /// Matching request ID.
    pub request_id: u32,
}

impl KeyPackageUnavailable {
    /// Create a new unavailable response.
    pub fn new(target: MeshAddress, request_id: u32) -> Self {
        Self {
            target_addr: target,
            request_id,
        }
    }

    /// Serialize to CBOR.
    pub fn to_wire(&self) -> Vec<u8> {
        let mut buf = Vec::new();
        buf.push(MessageType::KeyPackageUnavailable as u8);
        ciborium::into_writer(self, &mut buf).expect("CBOR serialization");
        buf
    }

    /// Deserialize from CBOR (after type byte).
    pub fn from_wire(bytes: &[u8]) -> anyhow::Result<Self> {
        if bytes.is_empty() || bytes[0] != MessageType::KeyPackageUnavailable as u8 {
            anyhow::bail!("not a KeyPackageUnavailable");
        }
        let resp: Self = ciborium::from_reader(&bytes[1..])?;
        Ok(resp)
    }
}

/// Parse the message type from wire bytes.
pub fn parse_message_type(bytes: &[u8]) -> Option<MessageType> {
    if bytes.is_empty() {
        return None;
    }
    match bytes[0] {
        0x10 => Some(MessageType::KeyPackageRequest),
        0x11 => Some(MessageType::KeyPackageResponse),
        0x12 => Some(MessageType::KeyPackageUnavailable),
        _ => None,
    }
}

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

    fn make_address(seed: u8) -> MeshAddress {
        MeshAddress::from_bytes([seed; 16])
    }

    #[test]
    fn request_roundtrip() {
        let req = KeyPackageRequest::new(make_address(1), make_address(2));
        let wire = req.to_wire();
        let restored = KeyPackageRequest::from_wire(&wire).expect("parse");

        assert_eq!(req.requester_addr, restored.requester_addr);
        assert_eq!(req.target_addr, restored.target_addr);
        assert_eq!(req.request_id, restored.request_id);
    }

    #[test]
    fn request_with_hash_roundtrip() {
        let hash = [0xAB; 8];
        let req = KeyPackageRequest::with_hash(make_address(1), make_address(2), hash);
        let wire = req.to_wire();
        let restored = KeyPackageRequest::from_wire(&wire).expect("parse");

        assert_eq!(req.hash, restored.hash);
        assert_eq!(Some(hash), restored.hash);
    }

    #[test]
    fn response_roundtrip() {
        let kp_bytes = vec![0x42; 100];
        let resp = KeyPackageResponse::new(make_address(3), kp_bytes.clone(), 12345);
        let wire = resp.to_wire();
        let restored = KeyPackageResponse::from_wire(&wire).expect("parse");

        assert_eq!(resp.owner_addr, restored.owner_addr);
        assert_eq!(resp.keypackage_bytes, restored.keypackage_bytes);
        assert_eq!(resp.hash, restored.hash);
        assert_eq!(resp.request_id, restored.request_id);
        assert!(restored.verify_hash());
    }

    #[test]
    fn unavailable_roundtrip() {
        let resp = KeyPackageUnavailable::new(make_address(4), 99999);
        let wire = resp.to_wire();
        let restored = KeyPackageUnavailable::from_wire(&wire).expect("parse");

        assert_eq!(resp.target_addr, restored.target_addr);
        assert_eq!(resp.request_id, restored.request_id);
    }

    #[test]
    fn parse_message_type_works() {
        let req = KeyPackageRequest::new(make_address(1), make_address(2));
        let wire = req.to_wire();
        assert_eq!(parse_message_type(&wire), Some(MessageType::KeyPackageRequest));

        let resp = KeyPackageResponse::new(make_address(3), vec![0x42], 1);
        let wire = resp.to_wire();
        assert_eq!(parse_message_type(&wire), Some(MessageType::KeyPackageResponse));

        let unavail = KeyPackageUnavailable::new(make_address(4), 2);
        let wire = unavail.to_wire();
        assert_eq!(parse_message_type(&wire), Some(MessageType::KeyPackageUnavailable));

        assert_eq!(parse_message_type(&[]), None);
        assert_eq!(parse_message_type(&[0xFF]), None);
    }

    #[test]
    fn measure_protocol_overhead() {
        let req = KeyPackageRequest::new(make_address(1), make_address(2));
        let wire = req.to_wire();
        println!("KeyPackageRequest: {} bytes", wire.len());

        let kp_bytes = vec![0x42; 306]; // Typical MLS KeyPackage size
        let resp = KeyPackageResponse::new(make_address(3), kp_bytes.clone(), 12345);
        let wire = resp.to_wire();
        println!("KeyPackageResponse (306B payload): {} bytes", wire.len());
        println!("Response overhead: {} bytes", wire.len() - 306);

        let unavail = KeyPackageUnavailable::new(make_address(4), 99999);
        let wire = unavail.to_wire();
        println!("KeyPackageUnavailable: {} bytes", wire.len());

        // Assertions
        assert!(req.to_wire().len() < 100, "request should be compact");
        assert!(unavail.to_wire().len() < 50, "unavailable should be compact");
    }
}