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

//! Protocol Codec Implementation
//!
//! Tokio codec for encoding/decoding frames with the binary protocol.

use bytes::{Buf, BytesMut};
use tokio_util::codec::{Decoder, Encoder};
use std::io;
use lz4_flex::{compress_prepend_size, decompress_size_prepended};
use crc32fast::Hasher;

use tracing::trace;
use super::wire::{Frame, FrameHeader, HEADER_SIZE, MAGIC_HEADER};

/// Codec for the memory protocol
pub struct MemoryProtocolCodec {
    /// Enable compression for bodies > 1KB
    enable_compression: bool,
    /// Maximum frame size (16MB default)
    max_frame_size: usize,
}

impl MemoryProtocolCodec {
    /// Create a new codec
    pub fn new() -> Self {
        Self {
            enable_compression: true,
            max_frame_size: 16 * 1024 * 1024, // 16MB
        }
    }

    /// Create codec with custom settings
    pub fn with_config(enable_compression: bool, max_frame_size: usize) -> Self {
        Self {
            enable_compression,
            max_frame_size,
        }
    }

    /// Compress body if needed
    fn maybe_compress(&self, body: &[u8]) -> (Vec<u8>, bool) {
        if self.enable_compression && body.len() > 1024 {
            let compressed = compress_prepend_size(body);
            // Only use compression if it actually reduces size
            if compressed.len() < body.len() {
                return (compressed, true);
            }
        }
        (body.to_vec(), false)
    }

    /// Decompress body if needed
    fn maybe_decompress(&self, body: &[u8], is_compressed: bool) -> io::Result<Vec<u8>> {
        if is_compressed {
            decompress_size_prepended(body)
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
        } else {
            Ok(body.to_vec())
        }
    }

    /// Calculate CRC32 checksum
    fn calculate_checksum(data: &[u8]) -> u32 {
        let mut hasher = Hasher::new();
        hasher.update(data);
        hasher.finalize()
    }
}

impl Default for MemoryProtocolCodec {
    fn default() -> Self {
        Self::new()
    }
}

impl Decoder for MemoryProtocolCodec {
    type Item = Frame;
    type Error = io::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        // Need at least header size
        if src.len() < HEADER_SIZE {
            return Ok(None);
        }

        // Peek at header to get body length
        let body_length = {
            let mut peek = src.clone();
            
            // Validate magic
            let magic = peek.get_u32();
            if magic != MAGIC_HEADER {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("Invalid magic header: {:08x}", magic),
                ));
            }

            // Skip to body length field (at offset 23)
            peek.advance(19); // skip version(1) + type(1) + flags(1) + message_id(16)
            peek.get_u64() as usize
        };

        // Check frame size limit
        if body_length > self.max_frame_size {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Frame too large: {} bytes", body_length),
            ));
        }

        // Wait for complete frame
        if src.len() < HEADER_SIZE + body_length {
            src.reserve(HEADER_SIZE + body_length - src.len());
            return Ok(None);
        }

        // Parse header
        let header = FrameHeader::read_from(src)?
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Failed to parse header"))?;

        // Extract body
        let mut body = src.split_to(body_length).to_vec();

        // Verify checksum if present
        if header.has_checksum() {
            if body.len() < 4 {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "Body too short for checksum",
                ));
            }

            let expected_crc = u32::from_le_bytes([body[0], body[1], body[2], body[3]]);
            body = body[4..].to_vec();

            let actual_crc = Self::calculate_checksum(&body);
            if expected_crc != actual_crc {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("CRC mismatch: expected {:08x}, got {:08x}", expected_crc, actual_crc),
                ));
            }
        }

        // Decompress if needed
        body = self.maybe_decompress(&body, header.is_compressed())?;

        trace!("Decoded frame type: {:?}", header.message_type);
        Ok(Some(Frame { header, body }))
    }
}

impl Encoder<Frame> for MemoryProtocolCodec {
    type Error = io::Error;

    fn encode(&mut self, item: Frame, dst: &mut BytesMut) -> Result<(), Self::Error> {
        let mut header = item.header;
        let mut body = item.body;

        trace!("Encoding frame type: {:?}", header.message_type);

        // Compress if beneficial
        let (compressed_body, was_compressed) = self.maybe_compress(&body);
        if was_compressed {
            body = compressed_body;
            header.set_compressed();
        }

        // Calculate checksum
        let checksum = Self::calculate_checksum(&body);
        header.set_checksum();

        // Update body length (checksum + body)
        header.body_length = (4 + body.len()) as u64;

        // Reserve space
        dst.reserve(HEADER_SIZE + 4 + body.len());

        // Write header
        header.write_to(dst);

        // Write checksum (little-endian)
        dst.extend_from_slice(&checksum.to_le_bytes());

        // Write body
        dst.extend_from_slice(&body);

        Ok(())
    }
}

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

    #[test]
    fn test_encode_decode_roundtrip() {
        let mut codec = MemoryProtocolCodec::new();
        
        let original = Frame::new(MessageType::Store, b"test data".to_vec());
        
        let mut buf = BytesMut::new();
        codec.encode(original.clone(), &mut buf).unwrap();
        
        let decoded = codec.decode(&mut buf).unwrap().unwrap();
        
        assert_eq!(decoded.body, b"test data");
        assert_eq!(decoded.header.message_type, MessageType::Store);
    }

    #[test]
    fn test_compression() {
        let mut codec = MemoryProtocolCodec::new();
        
        // Large body that should be compressed
        let body = vec![b'a'; 2048];
        let original = Frame::new(MessageType::Store, body.clone());
        
        let mut buf = BytesMut::new();
        codec.encode(original, &mut buf).unwrap();
        
        // Verify compression happened (buffer should be smaller)
        assert!(buf.len() < 2048 + HEADER_SIZE);
        
        // Decode and verify body matches
        let decoded = codec.decode(&mut buf).unwrap().unwrap();
        assert_eq!(decoded.body, body);
    }
}