use std::io;
use thiserror::Error;
pub const CHUNK_HEADER_LEN: usize = 8;
pub const MAX_CHUNK_LEN: u32 = 16 * 1024 * 1024;
pub const CHANNEL_ID_ZERO: u32 = 0;
pub const CHUNK_HEADER_BAST: &str = include_str!("../../docs/architecture/chunk-header.bast.json");
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ChunkHeader {
pub channel_id: u32,
pub length: u32,
}
impl ChunkHeader {
pub fn new(channel_id: u32, length: u32) -> Self {
Self { channel_id, length }
}
pub fn is_eof(&self) -> bool {
self.length == 0
}
}
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ChunkError {
#[error("header buffer too short: need {need} bytes, have {have}")]
HeaderTooShort { need: usize, have: usize },
#[error("chunk too large: {length} bytes (max {max})")]
TooLarge { length: u32, max: u32 },
}
pub fn parse_header(buf: &[u8]) -> Result<ChunkHeader, ChunkError> {
if buf.len() < CHUNK_HEADER_LEN {
return Err(ChunkError::HeaderTooShort {
need: CHUNK_HEADER_LEN,
have: buf.len(),
});
}
let channel_id = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
let length = u32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]]);
if length > MAX_CHUNK_LEN {
return Err(ChunkError::TooLarge {
length,
max: MAX_CHUNK_LEN,
});
}
Ok(ChunkHeader { channel_id, length })
}
pub fn write_header(channel_id: u32, length: u32, out: &mut [u8]) -> Result<(), ChunkError> {
if out.len() < CHUNK_HEADER_LEN {
return Err(ChunkError::HeaderTooShort {
need: CHUNK_HEADER_LEN,
have: out.len(),
});
}
let header = &mut out[..CHUNK_HEADER_LEN];
header[0..4].copy_from_slice(&channel_id.to_be_bytes());
header[4..8].copy_from_slice(&length.to_be_bytes());
Ok(())
}
pub async fn read_header<R>(reader: &mut R) -> io::Result<ChunkHeader>
where
R: tokio::io::AsyncRead + Unpin,
{
use tokio::io::AsyncReadExt;
let mut buf = [0u8; CHUNK_HEADER_LEN];
reader.read_exact(&mut buf).await?;
parse_header(&buf).map_err(|e| match e {
ChunkError::HeaderTooShort { .. } => io::Error::other("header too short after read_exact"),
ChunkError::TooLarge { length, max } => io::Error::new(
io::ErrorKind::InvalidData,
format!("chunk too large: {length} bytes (max {max})"),
),
})
}
pub async fn write_chunk<W>(writer: &mut W, channel_id: u32, payload: &[u8]) -> io::Result<()>
where
W: tokio::io::AsyncWrite + Unpin,
{
use tokio::io::AsyncWriteExt;
let mut header = [0u8; CHUNK_HEADER_LEN];
write_header(channel_id, payload.len() as u32, &mut header)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e.to_string()))?;
writer.write_all(&header).await?;
if !payload.is_empty() {
writer.write_all(payload).await?;
}
Ok(())
}
pub async fn write_eof<W>(writer: &mut W, channel_id: u32) -> io::Result<()>
where
W: tokio::io::AsyncWrite + Unpin,
{
write_chunk(writer, channel_id, &[]).await
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::AsyncReadExt;
#[test]
fn parse_header_round_trips_channel_id_and_length() {
let mut buf = [0u8; 8];
write_header(7, 1024, &mut buf).expect("write");
let header = parse_header(&buf).expect("parse");
assert_eq!(header, ChunkHeader::new(7, 1024));
assert!(!header.is_eof());
}
#[test]
fn parse_header_eof_sentinel() {
let mut buf = [0u8; 8];
write_header(3, 0, &mut buf).expect("write");
let header = parse_header(&buf).expect("parse");
assert_eq!(header.length, 0);
assert!(header.is_eof());
}
#[test]
fn parse_header_channel_zero() {
let mut buf = [0u8; 8];
write_header(CHANNEL_ID_ZERO, 512, &mut buf).expect("write");
let header = parse_header(&buf).expect("parse");
assert_eq!(header.channel_id, CHANNEL_ID_ZERO);
}
#[test]
fn parse_header_max_length_accepted() {
let mut buf = [0u8; 8];
write_header(1, MAX_CHUNK_LEN, &mut buf).expect("write");
let header = parse_header(&buf).expect("parse");
assert_eq!(header.length, MAX_CHUNK_LEN);
}
#[test]
fn parse_header_too_large_returns_error() {
let mut buf = [0u8; 8];
write_header(1, MAX_CHUNK_LEN + 1, &mut buf).expect("write");
match parse_header(&buf) {
Err(ChunkError::TooLarge { length, max }) => {
assert_eq!(length, MAX_CHUNK_LEN + 1);
assert_eq!(max, MAX_CHUNK_LEN);
}
other => panic!("expected TooLarge, got {other:?}"),
}
}
#[test]
fn parse_header_short_buffer_returns_error() {
let buf = [0u8; 4];
match parse_header(&buf) {
Err(ChunkError::HeaderTooShort { need, have }) => {
assert_eq!(need, CHUNK_HEADER_LEN);
assert_eq!(have, 4);
}
other => panic!("expected HeaderTooShort, got {other:?}"),
}
}
#[test]
fn write_header_writes_be_bytes() {
let mut buf = [0u8; 8];
write_header(0x01020304, 0x05060708, &mut buf).expect("write");
assert_eq!(buf, [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]);
}
#[test]
fn write_header_short_buffer_returns_error() {
let mut buf = [0u8; 4];
match write_header(1, 0, &mut buf) {
Err(ChunkError::HeaderTooShort { need, have }) => {
assert_eq!(need, CHUNK_HEADER_LEN);
assert_eq!(have, 4);
}
other => panic!("expected HeaderTooShort, got {other:?}"),
}
}
#[test]
fn chunk_header_bast_is_valid_json_and_describes_the_wire_format() {
let doc: serde_json::Value =
serde_json::from_str(CHUNK_HEADER_BAST).expect("BAST doc is valid JSON");
let def = doc
.get("$defs")
.and_then(|d| d.get("ChunkHeader"))
.expect("ChunkHeader def present");
assert_eq!(def.get("kind").and_then(|v| v.as_str()), Some("struct"));
assert_eq!(def.get("endian").and_then(|v| v.as_str()), Some("big"));
let fields = def
.get("fields")
.and_then(|v| v.as_array())
.expect("fields");
assert_eq!(fields.len(), 2);
assert_eq!(
fields[0].get("name").and_then(|v| v.as_str()),
Some("channel_id")
);
assert_eq!(
fields[0].get("kind").and_then(|v| v.as_str()),
Some("uint32")
);
assert_eq!(
fields[1].get("name").and_then(|v| v.as_str()),
Some("length")
);
assert_eq!(
fields[1].get("kind").and_then(|v| v.as_str()),
Some("uint32")
);
}
#[tokio::test]
async fn read_header_round_trips_through_duplex() {
let (mut reader, mut writer) = tokio::io::duplex(64);
write_chunk(&mut writer, 42, b"hello").await.expect("write");
let header = read_header(&mut reader).await.expect("read header");
assert_eq!(header.channel_id, 42);
assert_eq!(header.length, 5);
let mut payload = [0u8; 5];
reader.read_exact(&mut payload).await.expect("read payload");
assert_eq!(&payload, b"hello");
}
#[tokio::test]
async fn write_eof_writes_zero_length_chunk() {
let (mut reader, mut writer) = tokio::io::duplex(64);
write_eof(&mut writer, 7).await.expect("write eof");
let header = read_header(&mut reader).await.expect("read header");
assert_eq!(header.channel_id, 7);
assert_eq!(header.length, 0);
assert!(header.is_eof());
}
#[tokio::test]
async fn write_chunk_empty_payload_writes_eof_sentinel() {
let (mut reader, mut writer) = tokio::io::duplex(64);
write_chunk(&mut writer, 9, &[]).await.expect("write");
let header = read_header(&mut reader).await.expect("read header");
assert_eq!(header.length, 0);
assert!(header.is_eof());
}
#[tokio::test]
async fn read_header_on_closed_stream_returns_unexpected_eof() {
let (mut reader, writer) = tokio::io::duplex(64);
drop(writer);
let mut buf = [0u8; CHUNK_HEADER_LEN];
let result = reader.read_exact(&mut buf).await;
assert!(result.is_err(), "read on closed stream should error");
}
}