use crate::buf::WriteBuf;
use crate::error::{Error, Result};
use crate::framing::{Frame, Framer};
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LengthWidth {
U8,
U16,
U32,
U24,
U64,
}
impl LengthWidth {
#[inline]
#[must_use]
pub const fn header_size(self) -> usize {
match self {
LengthWidth::U8 => 1,
LengthWidth::U16 => 2,
LengthWidth::U24 => 3,
LengthWidth::U32 => 4,
LengthWidth::U64 => 8,
}
}
#[inline]
#[must_use]
pub const fn max_payload(self) -> u64 {
match self {
LengthWidth::U8 => u8::MAX as u64,
LengthWidth::U16 => u16::MAX as u64,
LengthWidth::U24 => 0x00FF_FFFF,
LengthWidth::U32 => u32::MAX as u64,
LengthWidth::U64 => u64::MAX,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Endian {
Big,
Little,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct LengthPrefixed {
width: LengthWidth,
endian: Endian,
max_payload: u64,
adjustment: i64,
}
impl LengthPrefixed {
pub const DEFAULT_MAX_PAYLOAD: u64 = 8 * 1024 * 1024;
#[inline]
#[must_use]
pub const fn new(width: LengthWidth, endian: Endian) -> Self {
Self {
width,
endian,
max_payload: width.max_payload(),
adjustment: 0,
}
}
#[inline]
#[must_use]
pub const fn bounded(width: LengthWidth, endian: Endian) -> Self {
Self::new(width, endian).with_max_payload(Self::DEFAULT_MAX_PAYLOAD)
}
#[inline]
#[must_use]
pub const fn with_max_payload(mut self, max: u64) -> Self {
let cap = self.width.max_payload();
self.max_payload = if max < cap { max } else { cap };
self
}
#[inline]
#[must_use]
pub const fn with_length_adjustment(mut self, adjustment: i64) -> Self {
self.adjustment = adjustment;
self
}
#[inline]
#[must_use]
pub const fn width(&self) -> LengthWidth {
self.width
}
#[inline]
#[must_use]
pub const fn endian(&self) -> Endian {
self.endian
}
#[inline]
#[must_use]
pub const fn max_payload(&self) -> u64 {
self.max_payload
}
#[inline]
#[must_use]
pub const fn length_adjustment(&self) -> i64 {
self.adjustment
}
#[inline]
fn read_prefix(&self, input: &[u8]) -> Option<u64> {
let big = matches!(self.endian, Endian::Big);
Some(match self.width {
LengthWidth::U8 => u64::from(*input.first()?),
LengthWidth::U16 => {
let b: [u8; 2] = *input.first_chunk()?;
u64::from(if big {
u16::from_be_bytes(b)
} else {
u16::from_le_bytes(b)
})
}
LengthWidth::U24 => {
let [a, b, c] = *input.first_chunk::<3>()?;
let word = if big { [0, a, b, c] } else { [a, b, c, 0] };
u64::from(if big {
u32::from_be_bytes(word)
} else {
u32::from_le_bytes(word)
})
}
LengthWidth::U32 => {
let b: [u8; 4] = *input.first_chunk()?;
u64::from(if big {
u32::from_be_bytes(b)
} else {
u32::from_le_bytes(b)
})
}
LengthWidth::U64 => {
let b: [u8; 8] = *input.first_chunk()?;
if big {
u64::from_be_bytes(b)
} else {
u64::from_le_bytes(b)
}
}
})
}
#[inline]
fn prefix_bytes(&self, value: u64) -> ([u8; 8], usize) {
let n = self.width.header_size();
let out = match self.endian {
Endian::Big => (value << (8 * (8 - n))).to_be_bytes(),
Endian::Little => value.to_le_bytes(),
};
(out, n)
}
fn too_large(&self, len: u64) -> Error {
Error::FrameTooLarge {
len: saturate_usize(len),
limit: saturate_usize(self.max_payload),
}
}
}
impl Framer for LengthPrefixed {
fn next_frame<'a>(&self, input: &'a [u8]) -> Result<Option<Frame<'a>>> {
let header_size = self.width.header_size();
let Some(raw) = self.read_prefix(input) else {
return Ok(None);
};
let len = match raw.checked_add_signed(self.adjustment) {
Some(len) => len,
None if self.adjustment < 0 => return Err(Error::InvalidLengthPrefix),
None => return Err(self.too_large(u64::MAX)),
};
if len > self.max_payload {
return Err(self.too_large(len));
}
let consumed = usize::try_from(len)
.ok()
.and_then(|len| len.checked_add(header_size))
.ok_or_else(|| self.too_large(len))?;
let Some(frame) = input.get(header_size..consumed) else {
return Ok(None);
};
Ok(Some(Frame::new(frame, consumed)))
}
fn write_frame(&self, payload: &[u8], out: &mut WriteBuf<'_>) -> Result<()> {
let len = u64::try_from(payload.len()).unwrap_or(u64::MAX);
if len > self.max_payload {
return Err(self.too_large(len));
}
let prefix = if self.adjustment == 0 {
len
} else {
let prefix = i128::from(len) - i128::from(self.adjustment);
if prefix < 0 {
return Err(Error::InvalidLengthPrefix);
}
u64::try_from(prefix).unwrap_or(u64::MAX)
};
if prefix > self.width.max_payload() {
let effective = i128::from(self.width.max_payload()) + i128::from(self.adjustment);
let effective = u64::try_from(effective.max(0)).unwrap_or(u64::MAX);
return Err(Error::FrameTooLarge {
len: saturate_usize(len),
limit: saturate_usize(effective.min(self.max_payload)),
});
}
let (header, header_len) = self.prefix_bytes(prefix);
let total = payload
.len()
.checked_add(header_len)
.ok_or(Error::BufferFull)?;
if total > out.remaining() {
return Err(Error::BufferFull);
}
out.write_bytes(&header[..header_len])?;
out.write_bytes(payload)
}
#[inline]
fn max_frame_len(&self, payload_len: usize) -> Option<usize> {
payload_len.checked_add(self.width.header_size())
}
}
#[inline]
fn saturate_usize(value: u64) -> usize {
usize::try_from(value).unwrap_or(usize::MAX)
}
#[cfg(test)]
mod tests {
use super::*;
const WIDTHS: [LengthWidth; 5] = [
LengthWidth::U8,
LengthWidth::U16,
LengthWidth::U24,
LengthWidth::U32,
LengthWidth::U64,
];
#[test]
fn huge_u32_prefix_never_panics() {
let framer = LengthPrefixed::new(LengthWidth::U32, Endian::Big);
for prefix in [0xFFFF_FFFFu32, 0xFFFF_FFFE, 0xFFFF_FFFD, 0xFFFF_FFFC] {
let mut input = [0u8; 6];
input[..4].copy_from_slice(&prefix.to_be_bytes());
let result = framer.next_frame(&input);
if usize::BITS > 32 {
assert_eq!(result, Ok(None));
} else {
assert!(matches!(result, Err(Error::FrameTooLarge { .. })));
}
}
}
#[test]
fn huge_u64_prefix_never_panics() {
let framer = LengthPrefixed::new(LengthWidth::U64, Endian::Little);
let input = [0xFF; 12];
assert!(matches!(
framer.next_frame(&input),
Err(Error::FrameTooLarge { .. })
));
let framer = framer.with_length_adjustment(i64::MAX);
assert!(matches!(
framer.next_frame(&input),
Err(Error::FrameTooLarge { .. })
));
}
#[test]
fn every_width_and_endian_round_trips() {
for width in WIDTHS {
for endian in [Endian::Big, Endian::Little] {
let framer = LengthPrefixed::new(width, endian);
let mut out = [0u8; 32];
let mut buf = WriteBuf::new(&mut out);
framer.write_frame(b"abc", &mut buf).unwrap();
let n = buf.position();
assert_eq!(n, width.header_size() + 3);
let frame = framer.next_frame(&out[..n]).unwrap().unwrap();
assert_eq!(frame.payload(), b"abc");
assert_eq!(frame.consumed(), n);
assert_eq!(framer.max_frame_len(3), Some(n));
}
}
}
#[test]
fn prefix_byte_order() {
let mut out = [0u8; 16];
let cases: [(LengthWidth, Endian, &[u8]); 4] = [
(LengthWidth::U24, Endian::Big, &[0x00, 0x01, 0x02]),
(LengthWidth::U24, Endian::Little, &[0x02, 0x01, 0x00]),
(
LengthWidth::U64,
Endian::Big,
&[0, 0, 0, 0, 0, 0, 0x01, 0x02],
),
(
LengthWidth::U64,
Endian::Little,
&[0x02, 0x01, 0, 0, 0, 0, 0, 0],
),
];
for (width, endian, header) in cases {
let framer = LengthPrefixed::new(width, endian);
assert_eq!(framer.read_prefix(header), Some(0x0102));
assert_eq!(framer.read_prefix(&header[1..]), None);
let (bytes, n) = framer.prefix_bytes(0x0102);
assert_eq!(&bytes[..n], header);
let mut buf = WriteBuf::new(&mut out);
assert_eq!(framer.write_frame(&[0u8; 2], &mut buf), Ok(()));
}
}
#[test]
fn u24_limits() {
assert_eq!(LengthWidth::U24.max_payload(), (1 << 24) - 1);
let framer = LengthPrefixed::new(LengthWidth::U24, Endian::Big);
assert_eq!(framer.next_frame(&[0xFF, 0xFF, 0xFF, 0x00]), Ok(None));
}
#[test]
fn u8_big_round_trip() {
let framer = LengthPrefixed::new(LengthWidth::U8, Endian::Big);
let mut out = [0u8; 32];
let mut buf = WriteBuf::new(&mut out);
framer.write_frame(b"abc", &mut buf).unwrap();
let n = buf.position();
assert_eq!(&out[..n], &[0x03, b'a', b'b', b'c']);
let frame = framer.next_frame(&out[..n]).unwrap().unwrap();
assert_eq!(frame.payload(), b"abc");
assert_eq!(frame.consumed(), 4);
}
#[test]
fn partial_frame_returns_none() {
let framer = LengthPrefixed::new(LengthWidth::U16, Endian::Big);
let input = &[0x00, 0x0A, 0x01, 0x02, 0x03];
assert_eq!(framer.next_frame(input).unwrap(), None);
}
#[test]
fn header_only_input() {
let framer = LengthPrefixed::new(LengthWidth::U16, Endian::Big);
let frame = framer.next_frame(&[0x00, 0x00]).unwrap().unwrap();
assert_eq!(frame.payload(), b"");
assert_eq!(frame.consumed(), 2);
assert_eq!(framer.next_frame(&[0x00, 0x01]), Ok(None));
assert_eq!(framer.next_frame(&[0x00]), Ok(None));
}
#[test]
fn empty_input_returns_none() {
let framer = LengthPrefixed::new(LengthWidth::U16, Endian::Big);
assert_eq!(framer.next_frame(&[]).unwrap(), None);
}
#[test]
fn exceeds_max_payload() {
let framer = LengthPrefixed::new(LengthWidth::U16, Endian::Big).with_max_payload(4);
let input = &[0x00, 0x05, 0, 0, 0, 0, 0];
assert_eq!(
framer.next_frame(input),
Err(Error::FrameTooLarge { len: 5, limit: 4 })
);
let input = &[0x00, 0x04, 1, 2, 3, 4];
assert_eq!(
framer.next_frame(input).unwrap().unwrap().payload(),
&[1, 2, 3, 4]
);
}
#[test]
fn write_limit_boundaries() {
let framer = LengthPrefixed::new(LengthWidth::U8, Endian::Big).with_max_payload(3);
let mut out = [0u8; 16];
let mut buf = WriteBuf::new(&mut out);
framer.write_frame(&[0u8; 3], &mut buf).unwrap();
assert_eq!(
framer.write_frame(&[0u8; 4], &mut buf),
Err(Error::FrameTooLarge { len: 4, limit: 3 })
);
assert_eq!(buf.position(), 4);
}
#[test]
fn write_is_atomic_on_buffer_full() {
let framer = LengthPrefixed::new(LengthWidth::U32, Endian::Big);
let mut out = [0xEEu8; 6];
let mut buf = WriteBuf::new(&mut out);
assert_eq!(framer.write_frame(b"abc", &mut buf), Err(Error::BufferFull));
assert_eq!(buf.position(), 0);
assert_eq!(out, [0xEE; 6], "no partial header is written");
}
#[test]
fn bounded_constructor() {
let framer = LengthPrefixed::bounded(LengthWidth::U64, Endian::Big);
assert_eq!(framer.max_payload(), LengthPrefixed::DEFAULT_MAX_PAYLOAD);
assert_eq!(framer.width(), LengthWidth::U64);
assert_eq!(framer.endian(), Endian::Big);
assert_eq!(framer.length_adjustment(), 0);
let mut input = [0u8; 8];
input.copy_from_slice(&(LengthPrefixed::DEFAULT_MAX_PAYLOAD + 1).to_be_bytes());
assert!(matches!(
framer.next_frame(&input),
Err(Error::FrameTooLarge { .. })
));
}
#[test]
fn length_adjustment_including_header() {
let framer = LengthPrefixed::new(LengthWidth::U32, Endian::Big).with_length_adjustment(-4);
let mut out = [0u8; 16];
let mut buf = WriteBuf::new(&mut out);
framer.write_frame(b"xyz", &mut buf).unwrap();
assert_eq!(buf.written(), &[0, 0, 0, 7, b'x', b'y', b'z']);
let frame = framer.next_frame(&out[..7]).unwrap().unwrap();
assert_eq!(frame.payload(), b"xyz");
assert_eq!(frame.consumed(), 7);
assert_eq!(
framer.next_frame(&[0, 0, 0, 3]),
Err(Error::InvalidLengthPrefix)
);
}
#[test]
fn length_adjustment_positive() {
let framer = LengthPrefixed::new(LengthWidth::U8, Endian::Big).with_length_adjustment(2);
let frame = framer
.next_frame(&[0x01, b'a', b'b', b'c'])
.unwrap()
.unwrap();
assert_eq!(frame.payload(), b"abc");
let mut out = [0u8; 8];
let mut buf = WriteBuf::new(&mut out);
framer.write_frame(b"abc", &mut buf).unwrap();
assert_eq!(buf.written(), &[0x01, b'a', b'b', b'c']);
assert_eq!(
framer.write_frame(b"a", &mut buf),
Err(Error::InvalidLengthPrefix)
);
}
#[test]
fn length_adjustment_overflowing_the_width() {
let framer = LengthPrefixed::new(LengthWidth::U8, Endian::Big).with_length_adjustment(-1);
let mut out = [0u8; 512];
let mut buf = WriteBuf::new(&mut out);
assert!(matches!(
framer.write_frame(&[0u8; 255], &mut buf),
Err(Error::FrameTooLarge { .. })
));
assert_eq!(
framer.write_frame(&[0u8; 255], &mut buf),
Err(Error::FrameTooLarge {
len: 255,
limit: 254
})
);
framer.write_frame(&[0u8; 254], &mut buf).unwrap();
assert_eq!(buf.written()[0], 255);
let framer = LengthPrefixed::new(LengthWidth::U8, Endian::Big).with_length_adjustment(2);
assert_eq!(framer.max_payload(), 255);
assert_eq!(
framer.next_frame(&[0xFF]),
Err(Error::FrameTooLarge {
len: 257,
limit: 255
})
);
}
#[test]
fn saturate_usize_is_identity_when_it_fits() {
assert_eq!(saturate_usize(42), 42);
if usize::BITS == 32 {
assert_eq!(saturate_usize(u64::MAX), usize::MAX);
}
}
}