agecli-skill-protocol 0.2.0

//! Read skill messages from a stream.

use super::message::{SkillMessage, SkillMessageType};
use tokio::io::{AsyncRead, AsyncReadExt};

/// Reads binary-framed skill messages from a stream.
pub struct SkillMessageReader<R> {
    reader: R,
}

impl<R: AsyncRead + Unpin> SkillMessageReader<R> {
    pub fn new(reader: R) -> Self {
        Self { reader }
    }

    /// Read the next message from the stream.
    /// Returns `None` on clean EOF.
    pub async fn read_message(&mut self) -> std::io::Result<Option<SkillMessage>> {
        // Read 4-byte message type
        let mut type_buf = [0u8; 4];
        match read_exact_or_eof(&mut self.reader, &mut type_buf).await? {
            0 => return Ok(None), // EOF
            n if n < 4 => {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::UnexpectedEof,
                    "incomplete message type header",
                ));
            }
            _ => {}
        }

        let msg_type_raw = u32::from_be_bytes(type_buf);
        let msg_type = SkillMessageType::from_u32(msg_type_raw).ok_or_else(|| {
            std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("unknown message type: {msg_type_raw}"),
            )
        })?;

        // Read 4-byte payload length
        let mut len_buf = [0u8; 4];
        self.reader.read_exact(&mut len_buf).await?;
        let payload_len = u32::from_be_bytes(len_buf) as usize;

        // Sanity check
        if payload_len > 16 * 1024 * 1024 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                format!("message payload too large: {payload_len}"),
            ));
        }

        // Read payload
        let mut payload = vec![0u8; payload_len];
        if payload_len > 0 {
            self.reader.read_exact(&mut payload).await?;
        }

        Ok(Some(SkillMessage::new(msg_type, payload)))
    }
}

/// Read exactly `buf.len()` bytes, returning 0 on immediate EOF.
async fn read_exact_or_eof<R: AsyncRead + Unpin>(
    reader: &mut R,
    buf: &mut [u8],
) -> std::io::Result<usize> {
    let mut offset = 0;
    while offset < buf.len() {
        let n = reader.read(&mut buf[offset..]).await?;
        if n == 0 {
            return Ok(offset);
        }
        offset += n;
    }
    Ok(offset)
}