use tokio::io::{AsyncBufReadExt, AsyncRead, BufReader};
pub async fn read_bounded_delimited<R: AsyncRead + Unpin>(
reader: &mut BufReader<R>,
delimiter: u8,
max_len: usize,
) -> std::io::Result<Option<Vec<u8>>> {
let mut frame = Vec::new();
loop {
let available = reader.fill_buf().await?;
if available.is_empty() {
return if frame.is_empty() {
Ok(None)
} else {
Err(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"partial delimited frame",
))
};
}
let end = available.iter().position(|&b| b == delimiter);
let take = end.map_or(available.len(), |n| n + 1);
if take > max_len.saturating_sub(frame.len()) {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"delimited frame exceeds cap",
));
}
frame.extend_from_slice(&available[..take]);
reader.consume(take);
if end.is_some() {
return Ok(Some(frame));
}
}
}
#[derive(Debug)]
pub enum BoundedLine {
Line(String),
Eof,
TooLong,
}
pub async fn read_bounded_line<R>(
reader: &mut BufReader<R>,
max_len: usize,
) -> std::io::Result<(BoundedLine, usize)>
where
R: AsyncRead + Unpin,
{
let mut buf: Vec<u8> = Vec::new();
let mut consumed = 0usize;
loop {
let (newline_at, available_len) = {
let available = reader.fill_buf().await?;
if available.is_empty() {
return Ok((BoundedLine::Eof, consumed));
}
(available.iter().position(|&b| b == b'\n'), available.len())
};
match newline_at {
Some(idx) => {
if buf.len() + idx + 1 > max_len {
return Ok((BoundedLine::TooLong, consumed));
}
let take = idx + 1;
buf.extend_from_slice(&reader.buffer()[..take]);
reader.consume(take);
consumed += take;
return Ok((
BoundedLine::Line(String::from_utf8_lossy(&buf).into_owned()),
consumed,
));
}
None => {
if buf.len() + available_len > max_len {
return Ok((BoundedLine::TooLong, consumed));
}
buf.extend_from_slice(&reader.buffer()[..available_len]);
reader.consume(available_len);
consumed += available_len;
}
}
}
}