use std::io::{self, Read};
use bytes::{Buf, Bytes, BytesMut};
const DEFAULT_READ_CHUNK: usize = 8 * 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DecodeStep<F> {
Pending,
Frame(F),
}
#[derive(Debug, derive_more::Display, derive_more::From)]
pub enum DecodeError {
#[display("i/o error while decoding a frame: {_0}")]
Io(io::Error),
#[from(ignore)]
#[display("protocol violation while decoding a frame: {_0}")]
Protocol(Box<str>),
}
impl std::error::Error for DecodeError {}
impl DecodeError {
pub fn protocol(message: impl Into<Box<str>>) -> Self {
DecodeError::Protocol(message.into())
}
}
pub trait IncrementalDecoder {
type Frame;
fn step(&mut self, src: &mut impl Read) -> Result<DecodeStep<Self::Frame>, DecodeError>;
fn has_partial(&self) -> bool;
}
pub fn read_frame_blocking<D, R>(
decoder: &mut D,
src: &mut R,
) -> Result<Option<D::Frame>, DecodeError>
where
D: IncrementalDecoder,
R: Read,
{
let mut tracked = EofTrackingReader {
inner: src,
hit_eof: false,
};
loop {
match decoder.step(&mut tracked)? {
DecodeStep::Frame(frame) => return Ok(Some(frame)),
DecodeStep::Pending => {
if tracked.hit_eof {
return if decoder.has_partial() {
Err(DecodeError::protocol("source ended mid-frame (truncated)"))
} else {
Ok(None)
};
}
}
}
}
}
struct EofTrackingReader<'a, R: Read> {
inner: &'a mut R,
hit_eof: bool,
}
impl<R: Read> Read for EofTrackingReader<'_, R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let read = self.inner.read(buf)?;
if !buf.is_empty() && read == 0 {
self.hit_eof = true;
}
Ok(read)
}
}
#[derive(Debug)]
pub struct AccumulatingBuffer {
buf: BytesMut,
read_chunk: usize,
}
impl Default for AccumulatingBuffer {
fn default() -> Self {
Self::new()
}
}
impl AccumulatingBuffer {
#[must_use]
pub fn new() -> Self {
Self {
buf: BytesMut::new(),
read_chunk: DEFAULT_READ_CHUNK,
}
}
#[must_use]
pub fn with_capacity(capacity: usize) -> Self {
Self {
buf: BytesMut::with_capacity(capacity),
read_chunk: DEFAULT_READ_CHUNK,
}
}
#[must_use]
pub fn with_read_chunk(mut self, read_chunk: usize) -> Self {
self.read_chunk = read_chunk.max(1);
self
}
#[must_use]
pub fn len(&self) -> usize {
self.buf.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.buf.is_empty()
}
#[must_use]
pub fn view(&self) -> &[u8] {
&self.buf
}
pub fn fill_from(&mut self, src: &mut impl Read) -> io::Result<usize> {
let start = self.buf.len();
self.buf.resize(start + self.read_chunk, 0);
match src.read(&mut self.buf[start..]) {
Ok(read) => {
self.buf.truncate(start + read);
Ok(read)
}
Err(err) => {
self.buf.truncate(start);
if err.kind() == io::ErrorKind::WouldBlock {
Ok(0)
} else {
Err(err)
}
}
}
}
#[must_use = "the split-off payload is the decoded frame"]
pub fn split_to(&mut self, n: usize) -> Bytes {
assert!(
n <= self.buf.len(),
"split_to({n}) exceeds buffered length {}",
self.buf.len()
);
self.buf.split_to(n).freeze()
}
pub fn advance(&mut self, n: usize) {
assert!(
n <= self.buf.len(),
"advance({n}) exceeds buffered length {}",
self.buf.len()
);
self.buf.advance(n);
}
pub fn clear(&mut self) {
self.buf.clear();
}
}