use core::num::NonZeroU32;
use thiserror::Error;
const TERMINATOR: u8 = b'\n';
const COMPACT_AT: usize = 8 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MaxFrameSize(NonZeroU32);
impl MaxFrameSize {
pub const DEFAULT: Self = Self(match NonZeroU32::new(4 * 1024 * 1024) {
Some(inner) => inner,
None => NonZeroU32::MIN,
});
pub const fn new(bytes: u32) -> Result<Self, MaxFrameSizeError> {
match NonZeroU32::new(bytes) {
Some(inner) => Ok(Self(inner)),
None => Err(MaxFrameSizeError::Zero),
}
}
pub const fn bytes(self) -> u32 {
self.0.get()
}
fn as_usize(self) -> usize {
usize::try_from(self.0.get()).unwrap_or(usize::MAX)
}
}
impl core::fmt::Display for MaxFrameSize {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
write!(f, "{} bytes", self.0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum MaxFrameSizeError {
#[error("a maximum frame size of zero would reject every frame, including an empty one")]
Zero,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum FramingError {
#[error("frame exceeds the configured maximum of {limit}")]
FrameTooLarge {
limit: MaxFrameSize,
at_least: usize,
},
#[error("frame is not valid UTF-8 (first invalid byte at offset {valid_up_to})")]
NotUtf8 {
valid_up_to: usize,
},
}
#[derive(Debug, Clone)]
pub struct FrameDecoder {
limit: MaxFrameSize,
buffer: Vec<u8>,
start: usize,
}
impl FrameDecoder {
pub const fn new(limit: MaxFrameSize) -> Self {
Self {
limit,
buffer: Vec::new(),
start: 0,
}
}
pub const fn limit(&self) -> MaxFrameSize {
self.limit
}
pub const fn buffered(&self) -> usize {
self.buffer.len() - self.start
}
pub const fn is_empty(&self) -> bool {
self.buffered() == 0
}
pub fn push(&mut self, bytes: &[u8]) -> Result<(), FramingError> {
let limit = self.limit.as_usize();
let mut frame_len = self.buffered();
for byte in bytes {
if *byte == TERMINATOR {
frame_len = 0;
} else {
frame_len += 1;
if frame_len > limit {
return Err(FramingError::FrameTooLarge {
limit: self.limit,
at_least: frame_len,
});
}
}
}
self.buffer.extend_from_slice(bytes);
Ok(())
}
pub fn next_frame(&mut self) -> Result<Option<&str>, FramingError> {
self.compact();
let Some(pending) = self.buffer.get(self.start..) else {
return Ok(None);
};
let Some(newline) = pending.iter().position(|byte| *byte == TERMINATOR) else {
return Ok(None);
};
let end = self.start + newline;
let frame = self.buffer.get(self.start..end).unwrap_or_default();
self.start = end + 1;
match core::str::from_utf8(frame) {
Ok(text) => Ok(Some(text)),
Err(error) => Err(FramingError::NotUtf8 {
valid_up_to: error.valid_up_to(),
}),
}
}
fn compact(&mut self) {
if self.start == self.buffer.len() {
self.buffer.clear();
self.start = 0;
} else if self.start >= COMPACT_AT {
self.buffer.drain(..self.start);
self.start = 0;
}
}
}
pub fn encode_frame(
document: &str,
out: &mut Vec<u8>,
limit: MaxFrameSize,
) -> Result<(), FramingError> {
let len = document.len();
if len > limit.as_usize() {
return Err(FramingError::FrameTooLarge {
limit,
at_least: len,
});
}
out.extend_from_slice(document.as_bytes());
out.push(TERMINATOR);
Ok(())
}
#[cfg(test)]
mod tests {
use assert_matches::assert_matches;
use super::{FrameDecoder, FramingError, MaxFrameSize, MaxFrameSizeError, encode_frame};
#[derive(Debug)]
#[expect(dead_code, reason = "the payloads are read by the derived Debug")]
enum Failure {
Framing(FramingError),
Limit(MaxFrameSizeError),
}
impl From<FramingError> for Failure {
fn from(error: FramingError) -> Self {
Self::Framing(error)
}
}
impl From<MaxFrameSizeError> for Failure {
fn from(error: MaxFrameSizeError) -> Self {
Self::Limit(error)
}
}
fn roomy() -> Result<FrameDecoder, MaxFrameSizeError> {
Ok(FrameDecoder::new(MaxFrameSize::new(1024)?))
}
const fn limit_in_usize(limit: MaxFrameSize) -> usize {
limit.bytes() as usize
}
#[test]
fn one_frame_round_trips_through_encode_and_decode() -> Result<(), Failure> {
let document = r#"{"jsonrpc":"2.0","id":1,"result":{}}"#;
let mut wire = Vec::new();
encode_frame(document, &mut wire, MaxFrameSize::new(1024)?)?;
let mut decoder = roomy()?;
decoder.push(&wire)?;
assert_eq!(decoder.next_frame()?, Some(document));
assert_eq!(decoder.next_frame()?, None);
Ok(())
}
#[test]
fn several_frames_in_one_read_come_back_in_order() -> Result<(), Failure> {
let mut decoder = roomy()?;
decoder.push(b"one\ntwo\nthree\n")?;
for expected in ["one", "two", "three"] {
assert_eq!(decoder.next_frame()?, Some(expected));
}
assert_eq!(decoder.next_frame()?, None);
assert!(decoder.is_empty());
Ok(())
}
#[test]
fn several_frames_one_per_read_come_back_in_order() -> Result<(), Failure> {
let mut decoder = roomy()?;
for chunk in [b"one\n".as_slice(), b"two\n", b"three\n"] {
decoder.push(chunk)?;
}
for expected in ["one", "two", "three"] {
assert_eq!(decoder.next_frame()?, Some(expected));
}
assert_eq!(decoder.next_frame()?, None);
Ok(())
}
#[test]
fn a_frame_split_across_reads_is_reassembled() -> Result<(), Failure> {
let document = r#"{"jsonrpc":"2.0","method":"ping"}"#;
let bytes = document.as_bytes();
let split = document.len() / 2;
let mut decoder = roomy()?;
decoder.push(bytes.get(..split).unwrap_or_default())?;
assert_eq!(
decoder.next_frame()?,
None,
"a frame with no terminator is not yet a frame"
);
decoder.push(bytes.get(split..).unwrap_or_default())?;
decoder.push(b"\n")?;
assert_eq!(decoder.next_frame()?, Some(document));
Ok(())
}
#[test]
fn any_single_byte_split_of_a_two_frame_stream_decodes() -> Result<(), Failure> {
let stream = b"{\"a\":1}\n{\"b\":2}\n";
for split in 0..=stream.len() {
let mut decoder = roomy()?;
decoder.push(stream.get(..split).unwrap_or_default())?;
decoder.push(stream.get(split..).unwrap_or_default())?;
assert_eq!(decoder.next_frame()?, Some(r#"{"a":1}"#), "split {split}");
assert_eq!(decoder.next_frame()?, Some(r#"{"b":2}"#), "split {split}");
assert_eq!(decoder.next_frame()?, None, "split {split}");
}
Ok(())
}
#[test]
fn a_multi_byte_character_split_across_reads_is_not_truncated() -> Result<(), Failure> {
let document = r#"{"emoji":"…"}"#;
let bytes = document.as_bytes();
let split = bytes
.windows(3)
.position(|window| window == "…".as_bytes())
.map_or(0, |at| at + 1);
let mut decoder = roomy()?;
decoder.push(bytes.get(..split).unwrap_or_default())?;
assert_eq!(decoder.next_frame()?, None);
decoder.push(bytes.get(split..).unwrap_or_default())?;
decoder.push(b"\n")?;
assert_eq!(decoder.next_frame()?, Some(document));
Ok(())
}
#[test]
fn a_frame_at_exactly_the_limit_is_accepted() -> Result<(), Failure> {
let limit = MaxFrameSize::new(8)?;
let document = "12345678";
let mut wire = Vec::new();
encode_frame(document, &mut wire, limit)?;
let mut decoder = FrameDecoder::new(limit);
decoder.push(&wire)?;
assert_eq!(decoder.next_frame()?, Some(document));
Ok(())
}
#[test]
fn a_frame_one_byte_over_the_limit_is_refused() -> Result<(), Failure> {
let limit = MaxFrameSize::new(8)?;
let mut wire = Vec::new();
assert_matches!(
encode_frame("123456789", &mut wire, limit),
Err(FramingError::FrameTooLarge { limit: got, .. }) if got == limit
);
assert!(wire.is_empty(), "a refused frame must write nothing");
Ok(())
}
#[test]
fn an_unterminated_frame_cannot_grow_the_buffer_past_the_limit() -> Result<(), Failure> {
let limit = MaxFrameSize::new(64)?;
let mut decoder = FrameDecoder::new(limit);
assert_matches!(
decoder.push(&vec![b'x'; 16 * 1024]),
Err(FramingError::FrameTooLarge { limit: got, at_least })
if got == limit && at_least == limit_in_usize(limit) + 1
);
assert!(
decoder.buffered() <= limit_in_usize(limit),
"held {} bytes for a {limit} limit",
decoder.buffered()
);
Ok(())
}
#[test]
fn a_refusal_leaves_the_decoder_exactly_as_it_was() -> Result<(), Failure> {
let limit = MaxFrameSize::new(16)?;
let mut decoder = FrameDecoder::new(limit);
decoder.push(b"short\npartial")?;
let before = decoder.buffered();
assert_matches!(
decoder.push(&[b'y'; 64]),
Err(FramingError::FrameTooLarge { .. })
);
assert_eq!(decoder.buffered(), before, "the refusal buffered bytes");
assert_eq!(decoder.next_frame()?, Some("short"));
Ok(())
}
#[test]
fn a_completed_frame_does_not_count_against_the_next_one() -> Result<(), Failure> {
let limit = MaxFrameSize::new(8)?;
let mut decoder = FrameDecoder::new(limit);
for _ in 0..100 {
decoder.push(b"12345678\n")?;
assert_eq!(decoder.next_frame()?, Some("12345678"));
}
assert!(decoder.is_empty());
Ok(())
}
#[test]
fn invalid_utf8_is_a_typed_error_rather_than_a_lossy_decode() -> Result<(), Failure> {
let mut decoder = roomy()?;
decoder.push(b"{\"a\":\"\xff\"}\n")?;
assert_matches!(
decoder.next_frame(),
Err(FramingError::NotUtf8 { valid_up_to: 6 })
);
decoder.push(b"next\n")?;
assert_eq!(decoder.next_frame()?, Some("next"));
Ok(())
}
#[test]
fn an_empty_frame_is_a_frame_rather_than_a_skipped_line() -> Result<(), Failure> {
let mut decoder = roomy()?;
decoder.push(b"\n\n")?;
assert_eq!(decoder.next_frame()?, Some(""));
assert_eq!(decoder.next_frame()?, Some(""));
assert_eq!(decoder.next_frame()?, None);
Ok(())
}
#[test]
fn a_carriage_return_is_kept_as_content() -> Result<(), Failure> {
let mut decoder = roomy()?;
decoder.push(b"{\"a\":1}\r\n")?;
assert_eq!(decoder.next_frame()?, Some("{\"a\":1}\r"));
Ok(())
}
#[test]
fn a_carriage_return_before_a_newline_counts_against_the_limit() -> Result<(), Failure> {
let limit = MaxFrameSize::new(4)?;
let mut decoder = FrameDecoder::new(limit);
assert_matches!(
decoder.push(b"1234\r"),
Err(FramingError::FrameTooLarge { .. })
);
Ok(())
}
#[test]
fn the_limit_is_configurable_and_readable_back() -> Result<(), Failure> {
let limit = MaxFrameSize::new(4096)?;
let decoder = FrameDecoder::new(limit);
assert_eq!(decoder.limit(), limit);
assert_eq!(decoder.limit().bytes(), 4096);
Ok(())
}
#[test]
fn a_zero_limit_is_refused_at_construction() {
assert_eq!(MaxFrameSize::new(0), Err(MaxFrameSizeError::Zero));
assert!(MaxFrameSize::new(1).is_ok());
}
#[test]
fn the_default_limit_is_an_ordinary_value_rather_than_a_fallback() {
assert_eq!(MaxFrameSize::DEFAULT.bytes(), 4 * 1024 * 1024);
}
#[test]
fn a_long_connection_does_not_retain_every_frame_it_has_seen() -> Result<(), Failure> {
let limit = MaxFrameSize::new(64)?;
let mut decoder = FrameDecoder::new(limit);
for _ in 0..4096 {
decoder.push(b"a-frame\n")?;
assert_eq!(decoder.next_frame()?, Some("a-frame"));
}
assert!(
decoder.buffered() <= limit_in_usize(limit),
"retained {} bytes of history",
decoder.buffered()
);
assert!(decoder.is_empty());
Ok(())
}
}