rust-memory-service/src/protocol/wire.rs

//! Wire Format Definitions
//!
//! Binary protocol wire format

use bytes::{Buf, BufMut, BytesMut};
use std::io;

/// Protocol magic header: "MEMA" in ASCII
pub const MAGIC_HEADER: u32 = 0x4D454D41;

/// Current protocol version
pub const PROTOCOL_VERSION: u8 = 1;

/// Frame header size in bytes
pub const HEADER_SIZE: usize = 31;

/// Message types for the protocol
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum MessageType {
    // Memory operations
    Store = 0x01,
    Retrieve = 0x02,
    Update = 0x03,
    Delete = 0x04,
    BatchStore = 0x05,

    // Context operations
    GetContextWindow = 0x10,
    Consolidate = 0x11,

    // Streaming
    Subscribe = 0x20,
    Unsubscribe = 0x21,
    MemoryEvent = 0x22,

    // Management
    ApplyTtl = 0x30,
    DetectConflicts = 0x31,
    ResolveConflict = 0x32,

    // Health check
    Ping = 0x40,
    Pong = 0x41,

    // Responses
    Success = 0x80,
    Error = 0x81,
    Partial = 0x82,
}

impl From<u8> for MessageType {
    fn from(value: u8) -> Self {
        match value {
            0x01 => MessageType::Store,
            0x02 => MessageType::Retrieve,
            0x03 => MessageType::Update,
            0x04 => MessageType::Delete,
            0x05 => MessageType::BatchStore,
            0x10 => MessageType::GetContextWindow,
            0x11 => MessageType::Consolidate,
            0x20 => MessageType::Subscribe,
            0x21 => MessageType::Unsubscribe,
            0x22 => MessageType::MemoryEvent,
            0x30 => MessageType::ApplyTtl,
            0x31 => MessageType::DetectConflicts,
            0x32 => MessageType::ResolveConflict,
            0x40 => MessageType::Ping,
            0x41 => MessageType::Pong,
            0x80 => MessageType::Success,
            0x81 => MessageType::Error,
            0x82 => MessageType::Partial,
            _ => MessageType::Error,
        }
    }
}

/// Protocol flags
pub mod flags {
    /// Body is LZ4 compressed
    pub const COMPRESSED: u8 = 0x01;
    /// CRC32 checksum is present
    pub const CHECKSUM: u8 = 0x02;
    /// Body is encrypted with ChaCha20-Poly1305
    pub const ENCRYPTED: u8 = 0x04;
}

/// Frame header structure
///
/// Binary layout (31 bytes):
/// ```text
/// +----------------+----------------+----------------+----------------+
/// | Magic (4 bytes) | Version (1) | Type (1) | Flags (1) |
/// +----------------+----------------+----------------+----------------+
/// | Message ID (16 bytes)                                           |
/// +----------------+----------------+----------------+----------------+
/// | Body Length (8 bytes)                                           |
/// +----------------+----------------+----------------+----------------+
/// ```
#[derive(Debug, Clone)]
pub struct FrameHeader {
    /// Magic header (should be MAGIC_HEADER)
    pub magic: u32,
    /// Protocol version
    pub version: u8,
    /// Message type
    pub message_type: MessageType,
    /// Flags (compression, checksum, encryption)
    pub flags: u8,
    /// Unique message ID for correlation
    pub message_id: [u8; 16],
    /// Length of the body in bytes
    pub body_length: u64,
}

impl FrameHeader {
    /// Create a new frame header
    pub fn new(message_type: MessageType, body_length: u64) -> Self {
        let mut message_id = [0u8; 16];
        // Generate random message ID
        getrandom(&mut message_id);

        Self {
            magic: MAGIC_HEADER,
            version: PROTOCOL_VERSION,
            message_type,
            flags: 0,
            message_id,
            body_length,
        }
    }

    /// Create a response header for a request
    pub fn response(request: &FrameHeader, message_type: MessageType, body_length: u64) -> Self {
        Self {
            magic: MAGIC_HEADER,
            version: PROTOCOL_VERSION,
            message_type,
            flags: 0,
            message_id: request.message_id,
            body_length,
        }
    }

    /// Write the header to a buffer
    pub fn write_to(&self, buf: &mut BytesMut) {
        buf.put_u32(self.magic);
        buf.put_u8(self.version);
        buf.put_u8(self.message_type as u8);
        buf.put_u8(self.flags);
        buf.put_slice(&self.message_id);
        buf.put_u64(self.body_length);
    }

    /// Read a header from a buffer
    pub fn read_from(buf: &mut BytesMut) -> io::Result<Option<Self>> {
        if buf.len() < HEADER_SIZE {
            return Ok(None);
        }

        let magic = buf.get_u32();
        if magic != MAGIC_HEADER {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Invalid magic header: {:08x}", magic),
            ));
        }

        let version = buf.get_u8();
        let message_type = MessageType::from(buf.get_u8());
        let flags = buf.get_u8();

        let mut message_id = [0u8; 16];
        buf.copy_to_slice(&mut message_id);

        let body_length = buf.get_u64();

        Ok(Some(Self {
            magic,
            version,
            message_type,
            flags,
            message_id,
            body_length,
        }))
    }

    /// Check if body is compressed
    pub fn is_compressed(&self) -> bool {
        self.flags & flags::COMPRESSED != 0
    }

    /// Check if checksum is present
    pub fn has_checksum(&self) -> bool {
        self.flags & flags::CHECKSUM != 0
    }

    /// Check if body is encrypted
    pub fn is_encrypted(&self) -> bool {
        self.flags & flags::ENCRYPTED != 0
    }

    /// Set compression flag
    pub fn set_compressed(&mut self) {
        self.flags |= flags::COMPRESSED;
    }

    /// Set checksum flag
    pub fn set_checksum(&mut self) {
        self.flags |= flags::CHECKSUM;
    }
}

/// Complete frame with header and body
#[derive(Debug, Clone)]
pub struct Frame {
    pub header: FrameHeader,
    pub body: Vec<u8>,
}

impl Frame {
    /// Create a new frame
    pub fn new(message_type: MessageType, body: Vec<u8>) -> Self {
        Self {
            header: FrameHeader::new(message_type, body.len() as u64),
            body,
        }
    }

    /// Create a success response
    pub fn success(request: &FrameHeader, body: Vec<u8>) -> Self {
        Self {
            header: FrameHeader::response(request, MessageType::Success, body.len() as u64),
            body,
        }
    }

    /// Create an error response
    pub fn error(request: &FrameHeader, message: &str) -> Self {
        let body = message.as_bytes().to_vec();
        Self {
            header: FrameHeader::response(request, MessageType::Error, body.len() as u64),
            body,
        }
    }
}

/// Simple random bytes generator
fn getrandom(buf: &mut [u8]) {
    use rand::RngCore;
    rand::thread_rng().fill_bytes(buf);
}

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

    #[test]
    fn test_header_roundtrip() {
        let header = FrameHeader::new(MessageType::Store, 1024);
        
        let mut buf = BytesMut::with_capacity(HEADER_SIZE);
        header.write_to(&mut buf);

        let parsed = FrameHeader::read_from(&mut buf).unwrap().unwrap();

        assert_eq!(parsed.magic, MAGIC_HEADER);
        assert_eq!(parsed.version, PROTOCOL_VERSION);
        assert_eq!(parsed.body_length, 1024);
    }

    #[test]
    fn test_message_type_conversion() {
        assert_eq!(MessageType::from(0x01), MessageType::Store);
        assert_eq!(MessageType::from(0x80), MessageType::Success);
    }
}