bloop-protocol 1.1.0

//! Message framing over async byte streams.
//!
//! Every message on the wire is a frame consisting of a `message_type` byte,
//! the payload length as a little-endian `u32` and the payload bytes. This
//! module is the only place that layout exists; everything above it deals in
//! [`RawMessage`] values.

use std::io;

use thiserror::Error;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

/// Errors that can occur while reading or writing a frame.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum FrameError {
    /// The underlying stream failed.
    #[error(transparent)]
    Io(#[from] io::Error),

    /// The payload exceeds the allowed maximum length.
    #[error("payload length {length} exceeds the maximum of {max} bytes")]
    PayloadTooLarge { length: u64, max: u64 },
}

/// A raw protocol frame with an opcode and an undecoded payload.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RawMessage {
    /// Opcode identifying the message's purpose.
    pub message_type: u8,

    /// Undecoded payload bytes.
    pub payload: Vec<u8>,
}

impl RawMessage {
    /// Creates a raw message from an opcode and payload bytes.
    pub fn new(message_type: u8, payload: Vec<u8>) -> Self {
        Self {
            message_type,
            payload,
        }
    }
}

/// Reads a single frame from the stream.
///
/// `max_payload_len` bounds the payload allocation: a frame declaring a
/// larger payload is rejected before any payload bytes are read, so a hostile
/// or corrupted length field cannot trigger a huge allocation. Choose the
/// bound per direction. Client-to-server payloads are small (authentication
/// is the largest), while server-to-client audio data can span multiple
/// megabytes.
///
/// # Errors
///
/// Returns [`FrameError::PayloadTooLarge`] if the declared payload length
/// exceeds `max_payload_len`, or [`FrameError::Io`] if the stream fails or
/// ends mid-frame.
pub async fn read_frame<R>(reader: &mut R, max_payload_len: u32) -> Result<RawMessage, FrameError>
where
    R: AsyncRead + Unpin,
{
    let message_type = reader.read_u8().await?;
    let length = reader.read_u32_le().await?;

    if length > max_payload_len {
        return Err(FrameError::PayloadTooLarge {
            length: length.into(),
            max: max_payload_len.into(),
        });
    }

    let mut payload = vec![0; length as usize];
    reader.read_exact(&mut payload).await?;

    Ok(RawMessage {
        message_type,
        payload,
    })
}

/// Writes a single frame to the stream and flushes it.
///
/// Flushing per frame matches the protocol's strict request-response flow,
/// where a peer always waits for its counterpart's answer before sending the
/// next message.
///
/// # Errors
///
/// Returns [`FrameError::PayloadTooLarge`] if the payload does not fit the
/// `u32` length field, or [`FrameError::Io`] if the stream fails.
pub async fn write_frame<W>(writer: &mut W, message: &RawMessage) -> Result<(), FrameError>
where
    W: AsyncWrite + Unpin,
{
    let length: u32 =
        message
            .payload
            .len()
            .try_into()
            .map_err(|_| FrameError::PayloadTooLarge {
                length: message.payload.len() as u64,
                max: u32::MAX.into(),
            })?;

    writer.write_u8(message.message_type).await?;
    writer.write_u32_le(length).await?;
    writer.write_all(&message.payload).await?;
    writer.flush().await?;

    Ok(())
}

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

    #[tokio::test]
    async fn frame_round_trips() {
        let (mut client, mut server) = tokio::io::duplex(1024);
        let message = RawMessage::new(0x08, vec![4, 1, 2, 3, 4]);

        write_frame(&mut client, &message).await.unwrap();
        let read = read_frame(&mut server, 1024).await.unwrap();

        assert_eq!(read, message);
    }

    #[tokio::test]
    async fn frame_layout_matches_spec() {
        let (mut client, mut server) = tokio::io::duplex(1024);
        let message = RawMessage::new(0x01, vec![3, 3]);

        write_frame(&mut client, &message).await.unwrap();

        let mut bytes = [0u8; 7];
        server.read_exact(&mut bytes).await.unwrap();

        assert_eq!(bytes, [0x01, 2, 0, 0, 0, 3, 3]);
    }

    #[tokio::test]
    async fn empty_payload_round_trips() {
        let (mut client, mut server) = tokio::io::duplex(1024);
        let message = RawMessage::new(0x05, vec![]);

        write_frame(&mut client, &message).await.unwrap();
        let read = read_frame(&mut server, 1024).await.unwrap();

        assert_eq!(read, message);
    }

    #[tokio::test]
    async fn oversized_payload_is_rejected_before_reading() {
        let (mut client, mut server) = tokio::io::duplex(1024);
        let message = RawMessage::new(0x0b, vec![0; 512]);

        write_frame(&mut client, &message).await.unwrap();
        let error = read_frame(&mut server, 256).await.unwrap_err();

        assert!(matches!(
            error,
            FrameError::PayloadTooLarge {
                length: 512,
                max: 256
            }
        ));
    }

    #[tokio::test]
    async fn payload_at_exact_limit_is_accepted() {
        let (mut client, mut server) = tokio::io::duplex(1024);
        let message = RawMessage::new(0x0b, vec![7; 256]);

        write_frame(&mut client, &message).await.unwrap();
        let read = read_frame(&mut server, 256).await.unwrap();

        assert_eq!(read, message);
    }

    #[tokio::test]
    async fn truncated_frame_fails_with_io_error() {
        let (mut client, mut server) = tokio::io::duplex(1024);

        client.write_all(&[0x08, 10, 0, 0, 0, 1, 2]).await.unwrap();
        drop(client);

        let error = read_frame(&mut server, 1024).await.unwrap_err();
        assert!(matches!(error, FrameError::Io(_)));
    }
}