monocoque-rs-zmtp 0.3.0

use bytes::{Buf, Bytes, BytesMut};
use monocoque_core::buffer::SegmentedBuffer;
use std::io;
use thiserror::Error;

use monocoque_core::config::STAGING_BUF_INITIAL_CAP;

/// ZMTP protocol errors
#[derive(Debug, Error)]
pub enum ZmtpError {
    #[error("Incomplete frame")]
    Incomplete,

    #[error("Protocol violation: reserved bits set")]
    ReservedBits,

    #[error("Protocol violation: frame size too large")]
    SizeTooLarge,

    #[error("Protocol violation")]
    Protocol,

    #[error("Authentication failed")]
    AuthenticationFailed,
}

impl From<ZmtpError> for io::Error {
    fn from(err: ZmtpError) -> Self {
        Self::new(io::ErrorKind::InvalidData, err)
    }
}

impl From<io::Error> for ZmtpError {
    fn from(_err: io::Error) -> Self {
        // Convert IO errors to Protocol errors for now
        Self::Protocol
    }
}

/// Result type alias for ZMTP operations
pub type Result<T> = std::result::Result<T, ZmtpError>;

/// A decoded ZMTP frame
#[derive(Debug, Clone)]
pub struct ZmtpFrame {
    pub flags: u8,
    pub payload: Bytes,
}

impl ZmtpFrame {
    #[inline]
    pub const fn more(&self) -> bool {
        (self.flags & 0x01) != 0
    }

    #[inline]
    pub const fn is_command(&self) -> bool {
        (self.flags & 0x04) != 0
    }
}

/// Stateful ZMTP decoder
///
/// Fast path:
/// - Entire frame present → zero-copy slice
///
/// Slow path:
/// - Fragmented frame → reassemble into `BytesMut`
pub struct ZmtpDecoder {
    // Fragmentation state
    pending_flags: Option<u8>,
    expected_body_len: usize,
    staging: BytesMut,
    /// Maximum allowed frame body size (enforcement of ZMQ_MAXMSGSIZE)
    max_frame_size: usize,
}

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

impl ZmtpDecoder {
    #[must_use]
    pub fn new() -> Self {
        Self {
            pending_flags: None,
            expected_body_len: 0,
            staging: BytesMut::with_capacity(STAGING_BUF_INITIAL_CAP),
            max_frame_size: 64 * 1024 * 1024, // 64MB default (generous but bounded)
        }
    }

    /// Create a decoder with a custom maximum frame size.
    #[must_use]
    pub fn with_max_frame_size(max_frame_size: usize) -> Self {
        Self {
            pending_flags: None,
            expected_body_len: 0,
            staging: BytesMut::with_capacity(STAGING_BUF_INITIAL_CAP),
            max_frame_size,
        }
    }

    /// Update the maximum body length enforced by the decoder.
    #[inline]
    pub fn set_max_body_len(&mut self, max_body_len: Option<usize>) {
        self.max_frame_size = max_body_len.unwrap_or(64 * 1024 * 1024);
    }

    /// Check if more message frames are expected (partial multipart message).
    ///
    /// Returns `true` if the decoder is in the middle of reassembling a frame
    /// or if the last decoded frame had the MORE flag set.
    ///
    /// # ZeroMQ Compatibility
    ///
    /// Corresponds to `ZMQ_RCVMORE` (13) - check if more frames in current message.
    #[inline]
    pub const fn has_more(&self) -> bool {
        // Decoder is expecting more data for current frame
        self.pending_flags.is_some()
    }

    /// Decode a single frame from `src`
    ///
    /// Returns:
    /// - Ok(Some(frame)) → frame decoded
    /// - Ok(None) → need more data
    /// - Err → protocol violation
    pub fn decode(&mut self, src: &mut SegmentedBuffer) -> Result<Option<ZmtpFrame>> {
        // === Reassembly mode ===
        if let Some(flags) = self.pending_flags {
            let needed = self.expected_body_len - self.staging.len();
            let take = needed.min(src.len());
            if let Some(bytes) = src.take_bytes(take) {
                self.staging.extend_from_slice(&bytes);
            }

            if self.staging.len() < self.expected_body_len {
                return Ok(None);
            }

            let payload = self.staging.split().freeze();
            self.pending_flags = None;
            self.expected_body_len = 0;

            return Ok(Some(ZmtpFrame { flags, payload }));
        }

        // === Header parsing ===
        if src.len() < 2 {
            return Ok(None);
        }

        let front = src.front_chunk();
        let mut hdr = [0u8; 9];
        if front.len() >= 2 {
            hdr[0] = front[0];
            hdr[1] = front[1];
        } else if !src.copy_prefix(2, &mut hdr) {
            return Ok(None);
        }

        let flags = hdr[0];

        if (flags & 0x05) == 0x05 {
            return Err(ZmtpError::Protocol);
        }

        // Reserved bits must be zero (bits 3-7)
        if (flags & 0xF8) != 0 {
            return Err(ZmtpError::ReservedBits);
        }

        let is_long = (flags & 0x02) != 0;
        let header_len = if is_long { 9 } else { 2 };

        if src.len() < header_len {
            return Ok(None);
        }

        // === Body length ===
        let body_len = if is_long {
            if front.len() >= 9 {
                hdr.copy_from_slice(&front[..9]);
            } else if !src.copy_prefix(9, &mut hdr) {
                return Ok(None);
            }
            let mut buf = &hdr[1..9];
            let size = buf.get_u64();

            // MSB must be zero in ZMTP 3.x
            if size > 0x7FFF_FFFF_FFFF_FFFF {
                return Err(ZmtpError::SizeTooLarge);
            }

            let body_len = size as usize;
            if body_len > self.max_frame_size {
                return Err(ZmtpError::SizeTooLarge);
            }
            body_len
        } else {
            let body_len = hdr[1] as usize;
            if body_len > self.max_frame_size {
                return Err(ZmtpError::SizeTooLarge);
            }
            body_len
        };

        let total_len = header_len + body_len;

        // === Fast path: entire frame present ===
        if src.len() >= total_len {
            let payload = src.take_bytes_after_available(header_len, body_len);
            return Ok(Some(ZmtpFrame { flags, payload }));
        }

        // === Slow path: fragmentation ===
        src.advance(header_len);
        self.pending_flags = Some(flags);
        self.expected_body_len = body_len;
        self.staging.clear();

        let available = src.len().min(body_len);
        if let Some(bytes) = src.take_bytes(available) {
            self.staging.extend_from_slice(&bytes);
        }

        Ok(None)
    }
}

impl ZmtpFrame {
    /// Create a data frame
    pub const fn data(payload: Bytes, more: bool) -> Self {
        let mut flags = 0;
        if more {
            flags |= 0x01; // MORE
        }
        if payload.len() > 255 {
            flags |= 0x02; // LONG
        }
        Self { flags, payload }
    }

    /// Create a command frame
    pub const fn command(payload: Bytes) -> Self {
        let mut flags = 0x04; // COMMAND
        if payload.len() > 255 {
            flags |= 0x02; // LONG
        }
        Self { flags, payload }
    }

    /// Encode this frame to bytes
    pub fn encode(&self) -> Bytes {
        let body_len = self.payload.len();
        let is_long = body_len >= 256;
        let flags = if is_long {
            self.flags | 0x02
        } else {
            self.flags & !0x02
        };

        let mut out = BytesMut::with_capacity(if is_long { 9 } else { 2 } + body_len);

        out.extend_from_slice(&[flags]);

        if is_long {
            out.extend_from_slice(&(body_len as u64).to_be_bytes());
        } else {
            out.extend_from_slice(&[body_len as u8]);
        }

        out.extend_from_slice(&self.payload);

        out.freeze()
    }
}

/// Append a ZMTP data-frame header (flags + length prefix) to `buf`.
///
/// The vectored write path builds all headers contiguously in one reused buffer
/// and slices them back out, so the body `Bytes` is never copied into a
/// userspace buffer on its way to the kernel. `more` sets the MORE flag
/// (another frame follows in the same multipart message); frames of 256 bytes
/// or more use the long (8-byte) length form. Returns the header length written
/// (2 or 9) so the caller can slice it back without recomputing.
pub fn write_frame_header(buf: &mut BytesMut, body_len: usize, more: bool) -> usize {
    let is_long = body_len >= 256;

    let mut flags = 0u8;
    if more {
        flags |= 0x01; // MORE
    }
    if is_long {
        flags |= 0x02; // LONG
    }

    buf.extend_from_slice(&[flags]);
    if is_long {
        buf.extend_from_slice(&(body_len as u64).to_be_bytes());
        9
    } else {
        buf.extend_from_slice(&[body_len as u8]);
        2
    }
}

/// Encode a multipart message directly into a buffer.
///
/// This is a zero-allocation helper for encoding messages without
/// creating intermediate `ZmtpFrame` objects.
///
/// # Performance
///
/// Reuses the provided `BytesMut` buffer, avoiding allocations on hot path.
pub fn encode_multipart(msg: &[Bytes], buf: &mut BytesMut) {
    if msg.is_empty() {
        return;
    }

    // Fast path: single-frame message (common case)
    if msg.len() == 1 {
        encode_single(&msg[0], buf);
        return;
    }

    let total_len: usize = msg
        .iter()
        .map(|part| (if part.len() >= 256 { 9 } else { 2 }) + part.len())
        .sum();
    buf.reserve(total_len);

    for (i, part) in msg.iter().enumerate() {
        let more = i < msg.len() - 1;
        let is_long = part.len() >= 256;

        let mut flags = 0u8;
        if more {
            flags |= 0x01; // MORE
        }
        if is_long {
            flags |= 0x02; // LONG
        }

        buf.extend_from_slice(&[flags]);

        if is_long {
            buf.extend_from_slice(&(part.len() as u64).to_be_bytes());
        } else {
            buf.extend_from_slice(&[part.len() as u8]);
        }

        buf.extend_from_slice(part);
    }
}

/// Encode a single-frame data message directly into `buf`.
#[inline]
pub fn encode_single(part: &Bytes, buf: &mut BytesMut) {
    let is_long = part.len() >= 256;
    if is_long {
        buf.reserve(9 + part.len());
        buf.extend_from_slice(&[0x02]);
        buf.extend_from_slice(&(part.len() as u64).to_be_bytes());
    } else {
        buf.reserve(2 + part.len());
        buf.extend_from_slice(&[0x00, part.len() as u8]);
    }
    buf.extend_from_slice(part);
}

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

    #[test]
    fn decode_rejects_command_frame_with_more_flag() {
        let mut decoder = ZmtpDecoder::new();
        let mut src = SegmentedBuffer::new();
        src.push(Bytes::from_static(b"\x05\x06\x05READY"));

        assert!(matches!(decoder.decode(&mut src), Err(ZmtpError::Protocol)));
    }

    #[test]
    fn encode_sets_long_flag_for_public_large_frame_payload() {
        let frame = ZmtpFrame {
            flags: 0,
            payload: Bytes::from(vec![0x42; 256]),
        };

        let encoded = frame.encode();

        assert_eq!(encoded[0] & 0x02, 0x02);
    }

    #[test]
    fn encode_multipart_multi_frame_sets_more_and_long_headers() {
        let msg = vec![
            Bytes::from_static(b"ab"),
            Bytes::from(vec![0xCD; 256]),
            Bytes::from_static(b"z"),
        ];
        let mut buf = BytesMut::new();

        encode_multipart(&msg, &mut buf);

        assert_eq!(buf[0], 0x01);
        assert_eq!(buf[1], 2);
        assert_eq!(&buf[2..4], b"ab");

        assert_eq!(buf[4], 0x03);
        assert_eq!(&buf[5..13], 256u64.to_be_bytes().as_slice());
        assert_eq!(&buf[13..269], vec![0xCD; 256].as_slice());

        assert_eq!(buf[269], 0x00);
        assert_eq!(buf[270], 1);
        assert_eq!(buf[271], b'z');
        assert_eq!(buf.len(), 272);
    }

    #[test]
    fn encode_multipart_multi_frame_appends_to_preallocated_buffer() {
        let msg = vec![Bytes::from_static(b"a"), Bytes::from_static(b"bc")];
        let mut buf = BytesMut::with_capacity(16);
        buf.extend_from_slice(b"prefix");

        encode_multipart(&msg, &mut buf);

        assert_eq!(&buf[..6], b"prefix");
        assert_eq!(buf[6], 0x01);
        assert_eq!(buf[7], 1);
        assert_eq!(buf[8], b'a');
        assert_eq!(buf[9], 0x00);
        assert_eq!(buf[10], 2);
        assert_eq!(&buf[11..13], b"bc");
        assert_eq!(buf.len(), 13);
    }

    #[test]
    fn encode_multipart_empty_message_writes_nothing() {
        let mut buf = BytesMut::from(&b"prefix"[..]);

        encode_multipart(&[], &mut buf);

        assert_eq!(&buf[..], b"prefix");
    }
}