sa3p-protocol 0.1.3

//! Binary framing and multiplexing primitives for SA3P transport boundaries.
//!
//! Frame layout:
//! `[MAGIC 4B][STREAM_ID 4B][OPCODE 1B][PAYLOAD_LEN 4B][PAYLOAD]`
//!
//! This crate includes incremental decoding (`FrameCodec`) and stream queues
//! (`Multiplexer`) for socket or vsock transports.

use std::collections::{BTreeMap, VecDeque};

use sa3p_parser::Instruction;
use thiserror::Error;

pub const MAGIC: [u8; 4] = *b"SA3P";
const HEADER_LEN: usize = 13;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum Opcode {
    VfsRead = 0x01,
    VfsWriteChunk = 0x02,
    VfsRename = 0x03,
    PtySpawn = 0x04,
    PtyInput = 0x05,
}

impl Opcode {
    pub fn from_byte(byte: u8) -> Option<Self> {
        match byte {
            0x01 => Some(Self::VfsRead),
            0x02 => Some(Self::VfsWriteChunk),
            0x03 => Some(Self::VfsRename),
            0x04 => Some(Self::PtySpawn),
            0x05 => Some(Self::PtyInput),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Frame {
    pub stream_id: u32,
    pub opcode: Opcode,
    pub payload: Vec<u8>,
}

impl Frame {
    pub fn new(stream_id: u32, opcode: Opcode, payload: Vec<u8>) -> Self {
        Self {
            stream_id,
            opcode,
            payload,
        }
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(HEADER_LEN + self.payload.len());
        out.extend_from_slice(&MAGIC);
        out.extend_from_slice(&self.stream_id.to_be_bytes());
        out.push(self.opcode as u8);
        out.extend_from_slice(&(self.payload.len() as u32).to_be_bytes());
        out.extend_from_slice(&self.payload);
        out
    }

    pub fn decode(bytes: &[u8]) -> Result<Self, ProtocolError> {
        let (frame, consumed) = decode_frame(bytes)?;
        if consumed != bytes.len() {
            return Err(ProtocolError::TrailingBytes {
                trailing: bytes.len() - consumed,
            });
        }
        Ok(frame)
    }
}

#[derive(Debug, Error, PartialEq, Eq)]
pub enum ProtocolError {
    #[error("frame too short: need at least {expected} bytes, got {actual}")]
    FrameTooShort { expected: usize, actual: usize },
    #[error("invalid magic: got {found:?}")]
    InvalidMagic { found: [u8; 4] },
    #[error("unknown opcode: 0x{0:02x}")]
    UnknownOpcode(u8),
    #[error("incomplete frame payload: expected {expected} bytes, got {actual}")]
    IncompletePayload { expected: usize, actual: usize },
    #[error("trailing bytes after frame: {trailing}")]
    TrailingBytes { trailing: usize },
}

pub fn decode_frame(bytes: &[u8]) -> Result<(Frame, usize), ProtocolError> {
    if bytes.len() < HEADER_LEN {
        return Err(ProtocolError::FrameTooShort {
            expected: HEADER_LEN,
            actual: bytes.len(),
        });
    }

    let found_magic = [bytes[0], bytes[1], bytes[2], bytes[3]];
    if found_magic != MAGIC {
        return Err(ProtocolError::InvalidMagic { found: found_magic });
    }

    let stream_id = u32::from_be_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
    let opcode_byte = bytes[8];
    let opcode = Opcode::from_byte(opcode_byte).ok_or(ProtocolError::UnknownOpcode(opcode_byte))?;
    let payload_len = u32::from_be_bytes([bytes[9], bytes[10], bytes[11], bytes[12]]) as usize;

    let total_len = HEADER_LEN + payload_len;
    if bytes.len() < total_len {
        return Err(ProtocolError::IncompletePayload {
            expected: payload_len,
            actual: bytes.len().saturating_sub(HEADER_LEN),
        });
    }

    let payload = bytes[HEADER_LEN..total_len].to_vec();
    Ok((Frame::new(stream_id, opcode, payload), total_len))
}

#[derive(Debug, Default)]
pub struct FrameCodec {
    buffer: Vec<u8>,
}

impl FrameCodec {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn push_bytes(&mut self, bytes: &[u8]) -> Result<Vec<Frame>, ProtocolError> {
        self.buffer.extend_from_slice(bytes);
        let mut frames = Vec::new();

        loop {
            if self.buffer.is_empty() {
                break;
            }

            match decode_frame(&self.buffer) {
                Ok((frame, consumed)) => {
                    self.buffer.drain(..consumed);
                    frames.push(frame);
                }
                Err(ProtocolError::FrameTooShort { .. })
                | Err(ProtocolError::IncompletePayload { .. }) => break,
                Err(err) => return Err(err),
            }
        }

        Ok(frames)
    }

    pub fn buffered_len(&self) -> usize {
        self.buffer.len()
    }
}

#[derive(Debug, Default)]
pub struct Multiplexer {
    queues: BTreeMap<u32, VecDeque<Frame>>,
}

impl Multiplexer {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn push(&mut self, frame: Frame) {
        self.queues
            .entry(frame.stream_id)
            .or_default()
            .push_back(frame);
    }

    pub fn pop_next(&mut self, stream_id: u32) -> Option<Frame> {
        let queue = self.queues.get_mut(&stream_id)?;
        let frame = queue.pop_front();
        if queue.is_empty() {
            self.queues.remove(&stream_id);
        }
        frame
    }

    pub fn stream_ids(&self) -> Vec<u32> {
        self.queues.keys().copied().collect()
    }

    pub fn len_for_stream(&self, stream_id: u32) -> usize {
        self.queues.get(&stream_id).map(VecDeque::len).unwrap_or(0)
    }
}

pub fn frame_from_instruction(stream_id: u32, instruction: &Instruction) -> Option<Frame> {
    match instruction {
        Instruction::WriteChunk(bytes) => {
            Some(Frame::new(stream_id, Opcode::VfsWriteChunk, bytes.clone()))
        }
        _ => None,
    }
}

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

    #[test]
    fn frame_round_trip_encode_decode() {
        let frame = Frame::new(7, Opcode::PtySpawn, b"cargo test".to_vec());
        let encoded = frame.encode();
        let decoded = Frame::decode(&encoded).expect("frame should decode");

        assert_eq!(decoded, frame);
    }

    #[test]
    fn decode_rejects_invalid_magic() {
        let bytes = b"BAD!\0\0\0\x01\x02\0\0\0\x00";
        let err = Frame::decode(bytes).expect_err("decode should fail");

        assert!(matches!(err, ProtocolError::InvalidMagic { .. }));
    }

    #[test]
    fn frame_codec_handles_incremental_feeds() {
        let frame = Frame::new(1, Opcode::VfsRead, b"payload".to_vec());
        let bytes = frame.encode();

        let mut codec = FrameCodec::new();
        let first = codec
            .push_bytes(&bytes[..5])
            .expect("partial push should succeed");
        assert!(first.is_empty());
        assert_eq!(codec.buffered_len(), 5);

        let second = codec
            .push_bytes(&bytes[5..])
            .expect("second push should succeed");
        assert_eq!(second, vec![frame]);
        assert_eq!(codec.buffered_len(), 0);
    }

    #[test]
    fn multiplexer_isolates_stream_queues() {
        let mut mux = Multiplexer::new();
        mux.push(Frame::new(1, Opcode::VfsRead, vec![1]));
        mux.push(Frame::new(2, Opcode::PtyInput, vec![2]));
        mux.push(Frame::new(1, Opcode::VfsRename, vec![3]));

        assert_eq!(mux.stream_ids(), vec![1, 2]);
        assert_eq!(mux.len_for_stream(1), 2);

        let first = mux.pop_next(1).expect("stream 1 frame");
        assert_eq!(first.payload, vec![1]);
        let second = mux.pop_next(1).expect("stream 1 second frame");
        assert_eq!(second.payload, vec![3]);
        assert_eq!(mux.len_for_stream(1), 0);
        assert_eq!(mux.len_for_stream(2), 1);
    }

    #[test]
    fn instruction_write_chunk_maps_to_opcode_0x02_frame() {
        let instruction = Instruction::WriteChunk(b"hello".to_vec());
        let frame = frame_from_instruction(9, &instruction).expect("should map");

        assert_eq!(frame.stream_id, 9);
        assert_eq!(frame.opcode as u8, 0x02);
        assert_eq!(frame.payload, b"hello".to_vec());
    }
}