armdb/replication/

protocol.rs

use std::io::{self, Read, Write};

/// Maximum allowed frame payload size. Guards `read_frame` against a
/// malformed length prefix that could trigger multi-GiB allocations.
pub const MAX_FRAME_SIZE: usize = 16 * 1024 * 1024; // 16 MiB cap; protects against malformed length prefix

/// Current version of the variable-length replication wire protocol.
pub const VAR_PROTOCOL_VERSION: u16 = 2;

/// Message types for the replication protocol.
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageType {
    SyncRequest = 1,
    ShardInfo = 2,
    EntryBatch = 3,
    CaughtUp = 4,
    Ack = 5,
    Heartbeat = 6,
    Error = 255,
}

impl MessageType {
    fn from_u8(v: u8) -> Option<Self> {
        match v {
            1 => Some(Self::SyncRequest),
            2 => Some(Self::ShardInfo),
            3 => Some(Self::EntryBatch),
            4 => Some(Self::CaughtUp),
            5 => Some(Self::Ack),
            6 => Some(Self::Heartbeat),
            255 => Some(Self::Error),
            _ => None,
        }
    }
}

impl crate::frame_reader::FrameKind for MessageType {
    const MAX_PAYLOAD: usize = MAX_FRAME_SIZE;
    fn from_u8(b: u8) -> Option<Self> {
        // The inherent from_u8 (above) takes precedence over the trait
        // method — no recursion.
        MessageType::from_u8(b)
    }
}

/// Wire frame: [type:u8][len:u32 LE][payload:len bytes]
pub type Frame = crate::frame_reader::GenericFrame<MessageType>;

pub fn write_frame(w: &mut impl Write, frame: &Frame) -> io::Result<()> {
    w.write_all(&[frame.msg_type as u8])?;
    w.write_all(&(frame.payload.len() as u32).to_le_bytes())?;
    w.write_all(&frame.payload)?;
    w.flush()
}

pub fn read_frame(r: &mut impl Read) -> io::Result<Frame> {
    let mut type_buf = [0u8; 1];
    r.read_exact(&mut type_buf)?;
    let msg_type = MessageType::from_u8(type_buf[0]).ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidData,
            format!("unknown message type: {}", type_buf[0]),
        )
    })?;
    let mut len_buf = [0u8; 4];
    r.read_exact(&mut len_buf)?;
    let len = u32::from_le_bytes(len_buf) as usize;
    if len > MAX_FRAME_SIZE {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "frame too large",
        ));
    }
    let mut payload = vec![0u8; len];
    if len > 0 {
        r.read_exact(&mut payload)?;
    }
    Ok(Frame { msg_type, payload })
}

// --- SyncRequest: version:u16, shard_id:u8, from_gsn:u64, key_len:u16 ---
// Payload layout: version (2 LE) + shard_id (1) + from_gsn (8 LE) + key_len (2 LE) = 13 bytes.
// One Engine = one Tree, so a single key_len suffices (C7/C20 cleanup).

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SyncRequest {
    pub protocol_version: u16,
    pub shard_id: u8,
    pub from_gsn: u64,
    pub key_len: u16,
}

impl SyncRequest {
    pub fn encode(&self) -> Frame {
        let mut payload = Vec::with_capacity(13);
        payload.extend_from_slice(&self.protocol_version.to_le_bytes());
        payload.push(self.shard_id);
        payload.extend_from_slice(&self.from_gsn.to_le_bytes());
        payload.extend_from_slice(&self.key_len.to_le_bytes());
        Frame {
            msg_type: MessageType::SyncRequest,
            payload,
        }
    }

    pub fn decode(payload: &[u8]) -> io::Result<Self> {
        if payload.len() != 13 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "invalid SyncRequest length",
            ));
        }
        let protocol_version =
            u16::from_le_bytes(payload[0..2].try_into().expect("validated length"));
        let shard_id = payload[2];
        let from_gsn = u64::from_le_bytes(payload[3..11].try_into().expect("validated length"));
        let key_len = u16::from_le_bytes(payload[11..13].try_into().expect("validated length"));
        Ok(Self {
            protocol_version,
            shard_id,
            from_gsn,
            key_len,
        })
    }
}

// --- ShardInfo: version:u16, shard_count:u8, max_file_size:u64 ---

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ShardInfo {
    pub protocol_version: u16,
    pub shard_count: u8,
    pub max_file_size: u64,
}

impl ShardInfo {
    pub fn encode(&self) -> Frame {
        let mut payload = Vec::with_capacity(11);
        payload.extend_from_slice(&self.protocol_version.to_le_bytes());
        payload.push(self.shard_count);
        payload.extend_from_slice(&self.max_file_size.to_le_bytes());
        Frame {
            msg_type: MessageType::ShardInfo,
            payload,
        }
    }

    pub fn decode(payload: &[u8]) -> io::Result<Self> {
        if payload.len() != 11 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "invalid ShardInfo length",
            ));
        }
        Ok(Self {
            protocol_version: u16::from_le_bytes(
                payload[0..2].try_into().expect("validated length"),
            ),
            shard_count: payload[2],
            max_file_size: u64::from_le_bytes(payload[3..11].try_into().expect("validated length")),
        })
    }
}

// --- EntryBatch: shard_id:u8, count:u32, entries:[entry_len:u32 + key_len:u16 + gsn:u64 + data] ---

pub struct WireEntry {
    pub entry_len: u32,
    pub key_len: u16,
    pub gsn: u64,
    pub data: Vec<u8>,
}

pub struct EntryBatch {
    pub shard_id: u8,
    pub entries: Vec<WireEntry>,
}

impl EntryBatch {
    pub fn encode(&self) -> Frame {
        let mut payload = Vec::with_capacity(5 + self.entries.len() * 64);
        payload.push(self.shard_id);
        payload.extend_from_slice(&(self.entries.len() as u32).to_le_bytes());
        for e in &self.entries {
            payload.extend_from_slice(&e.entry_len.to_le_bytes());
            payload.extend_from_slice(&e.key_len.to_le_bytes());
            payload.extend_from_slice(&e.gsn.to_le_bytes());
            payload.extend_from_slice(&e.data);
        }
        Frame {
            msg_type: MessageType::EntryBatch,
            payload,
        }
    }

    pub fn decode(payload: &[u8]) -> io::Result<Self> {
        if payload.len() < 5 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "EntryBatch too short",
            ));
        }
        let shard_id = payload[0];
        let count = u32::from_le_bytes(payload[1..5].try_into().expect("impossible")) as usize;
        let mut entries = Vec::with_capacity(count);
        let mut off = 5;
        for _ in 0..count {
            if off + 14 > payload.len() {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "EntryBatch truncated",
                ));
            }
            let entry_len =
                u32::from_le_bytes(payload[off..off + 4].try_into().expect("impossible"));
            let key_len =
                u16::from_le_bytes(payload[off + 4..off + 6].try_into().expect("impossible"));
            let gsn =
                u64::from_le_bytes(payload[off + 6..off + 14].try_into().expect("impossible"));
            off += 14;
            if off + entry_len as usize > payload.len() {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "EntryBatch data truncated",
                ));
            }
            let data = payload[off..off + entry_len as usize].to_vec();
            off += entry_len as usize;
            entries.push(WireEntry {
                entry_len,
                key_len,
                gsn,
                data,
            });
        }
        Ok(Self { shard_id, entries })
    }
}

// --- CaughtUp: shard_id:u8, leader_gsn:u64 ---

pub struct CaughtUp {
    pub shard_id: u8,
    pub leader_gsn: u64,
}

impl CaughtUp {
    pub fn encode(&self) -> Frame {
        let mut payload = Vec::with_capacity(9);
        payload.push(self.shard_id);
        payload.extend_from_slice(&self.leader_gsn.to_le_bytes());
        Frame {
            msg_type: MessageType::CaughtUp,
            payload,
        }
    }

    pub fn decode(payload: &[u8]) -> io::Result<Self> {
        if payload.len() < 9 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "CaughtUp too short",
            ));
        }
        Ok(Self {
            shard_id: payload[0],
            leader_gsn: u64::from_le_bytes(payload[1..9].try_into().expect("impossible")),
        })
    }
}

// --- Ack: shard_id:u8, last_gsn:u64 ---

pub struct AckMessage {
    pub shard_id: u8,
    pub last_gsn: u64,
}

impl AckMessage {
    pub fn encode(&self) -> Frame {
        let mut payload = Vec::with_capacity(9);
        payload.push(self.shard_id);
        payload.extend_from_slice(&self.last_gsn.to_le_bytes());
        Frame {
            msg_type: MessageType::Ack,
            payload,
        }
    }

    pub fn decode(payload: &[u8]) -> io::Result<Self> {
        if payload.len() < 9 {
            return Err(io::Error::new(io::ErrorKind::InvalidData, "Ack too short"));
        }
        Ok(Self {
            shard_id: payload[0],
            last_gsn: u64::from_le_bytes(payload[1..9].try_into().expect("impossible")),
        })
    }
}

// --- Heartbeat ---

pub fn encode_heartbeat() -> Frame {
    Frame {
        msg_type: MessageType::Heartbeat,
        payload: Vec::new(),
    }
}

// --- Error ---

/// Stable error categories carried by variable-replication `Error` frames.
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReplicationErrorCode {
    /// A transient I/O failure for which the follower may reconnect and retry.
    RetryableIo = 1,
    /// A configured resource limit prevented the requested catch-up operation.
    ResourceLimit = 2,
    /// The leader detected corrupted replication log data.
    CorruptedLog = 3,
    /// Leader and follower use incompatible variable-replication protocol versions.
    ProtocolMismatch = 4,
    /// The peer sent a request that violates the replication protocol.
    InvalidRequest = 5,
}

impl ReplicationErrorCode {
    fn from_u8(value: u8) -> Option<Self> {
        match value {
            1 => Some(Self::RetryableIo),
            2 => Some(Self::ResourceLimit),
            3 => Some(Self::CorruptedLog),
            4 => Some(Self::ProtocolMismatch),
            5 => Some(Self::InvalidRequest),
            _ => None,
        }
    }
}

/// A validated typed error decoded from a variable-replication `Error` frame.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReplicationWireError {
    /// Machine-readable error category used for retry classification.
    pub code: ReplicationErrorCode,
    /// Human-readable diagnostic detail supplied by the peer.
    pub message: String,
}

/// Encodes a typed variable-replication error frame.
///
/// Messages longer than `u16::MAX` bytes are truncated at a valid UTF-8 boundary.
pub fn encode_error(code: ReplicationErrorCode, message: &str) -> Frame {
    let max = message.len().min(u16::MAX as usize);
    let end = message.floor_char_boundary(max);
    let message = &message.as_bytes()[..end];
    let mut payload = Vec::with_capacity(3 + message.len());
    payload.push(code as u8);
    payload.extend_from_slice(&(message.len() as u16).to_le_bytes());
    payload.extend_from_slice(message);
    Frame {
        msg_type: MessageType::Error,
        payload,
    }
}

/// Decodes and validates a typed variable-replication error payload.
///
/// # Errors
///
/// Returns [`io::ErrorKind::InvalidData`] for unknown codes, short or trailing
/// payload bytes, a truncated declared message, or invalid UTF-8.
pub fn decode_error(payload: &[u8]) -> io::Result<ReplicationWireError> {
    if payload.len() < 3 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "replication error payload too short",
        ));
    }

    let code = ReplicationErrorCode::from_u8(payload[0]).ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidData,
            format!("unknown replication error code: {}", payload[0]),
        )
    })?;
    let message_len = u16::from_le_bytes([payload[1], payload[2]]) as usize;
    let expected_len = 3 + message_len;
    if payload.len() != expected_len {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "invalid replication error payload length",
        ));
    }
    let message = std::str::from_utf8(&payload[3..])
        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?
        .to_owned();

    Ok(ReplicationWireError { code, message })
}

#[cfg(test)]
mod max_frame_size_tests {
    use super::*;
    use std::io::Cursor;

    #[test]
    fn read_frame_rejects_oversized_length() {
        // Manually construct a frame header: type byte + 4-byte LE length too big
        let mut bytes = Vec::new();
        bytes.push(MessageType::SyncRequest as u8);
        bytes.extend_from_slice(&((MAX_FRAME_SIZE + 1) as u32).to_le_bytes());
        let mut cursor = Cursor::new(bytes);
        let err = read_frame(&mut cursor).unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("frame too large"), "got: {msg}");
    }

    #[test]
    fn frame_reader_rejects_garbage_type_byte_immediately() {
        // R3: a single garbage byte, then silence. Early validation must
        // yield InvalidData immediately — the client's StreamFailure maps it
        // to Terminal, and degrading into a retryable timeout is not
        // acceptable.
        let mut fr = crate::frame_reader::FrameReader::<MessageType>::new();
        let err = fr.read_frame(&mut Cursor::new([0u8])).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
        assert!(err.to_string().contains("unknown message type"));
    }

    #[test]
    fn frame_alias_keeps_debug_contract() {
        // R5: Frame is now an alias of GenericFrame<MessageType>; Debug must
        // survive for downstream code using format!("{frame:?}").
        let frame = Frame {
            msg_type: MessageType::Heartbeat,
            payload: Vec::new(),
        };
        assert!(format!("{frame:?}").contains("Heartbeat"));
    }
}

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

    #[test]
    fn sync_request_v2_round_trip() {
        let req = SyncRequest {
            protocol_version: VAR_PROTOCOL_VERSION,
            shard_id: 7,
            from_gsn: 0xDEAD_BEEF_CAFE_BABE,
            key_len: 32,
        };
        let frame = req.encode();
        assert_eq!(frame.payload.len(), 13);
        assert_eq!(SyncRequest::decode(&frame.payload).unwrap(), req);
    }

    #[test]
    fn shard_info_v2_round_trip() {
        let info = ShardInfo {
            protocol_version: VAR_PROTOCOL_VERSION,
            shard_count: 4,
            max_file_size: 64 * 1024 * 1024,
        };
        let frame = info.encode();
        assert_eq!(frame.payload.len(), 11);
        assert_eq!(ShardInfo::decode(&frame.payload).unwrap(), info);
    }

    #[test]
    fn typed_error_round_trips_every_code() {
        for code in [
            ReplicationErrorCode::RetryableIo,
            ReplicationErrorCode::ResourceLimit,
            ReplicationErrorCode::CorruptedLog,
            ReplicationErrorCode::ProtocolMismatch,
            ReplicationErrorCode::InvalidRequest,
        ] {
            let frame = encode_error(code, "detail");
            assert_eq!(
                decode_error(&frame.payload).unwrap(),
                ReplicationWireError {
                    code,
                    message: "detail".to_owned(),
                }
            );
        }
    }

    #[test]
    fn typed_error_rejects_trailing_and_invalid_utf8() {
        let mut trailing = encode_error(ReplicationErrorCode::InvalidRequest, "x").payload;
        trailing.push(0);
        assert_eq!(
            decode_error(&trailing).unwrap_err().kind(),
            io::ErrorKind::InvalidData
        );

        let invalid = [ReplicationErrorCode::InvalidRequest as u8, 1, 0, 0xFF];
        assert_eq!(
            decode_error(&invalid).unwrap_err().kind(),
            io::ErrorKind::InvalidData
        );
    }

    #[test]
    fn typed_error_truncates_on_utf8_boundary() {
        let message = format!("{}€", "a".repeat(u16::MAX as usize - 1));
        let decoded =
            decode_error(&encode_error(ReplicationErrorCode::RetryableIo, &message).payload)
                .unwrap();
        assert!(decoded.message.is_char_boundary(decoded.message.len()));
        assert_eq!(decoded.message.len(), u16::MAX as usize - 1);
    }
}