use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use crate::framing::Framing;
use crate::headers;
const CHUNK: usize = 16 * 1024;
const MAX_CHUNK_LINE: usize = 1024;
pub(crate) struct Buffered<'a, R> {
src: &'a mut R,
buf: Vec<u8>,
pos: usize,
}
impl<'a, R: AsyncRead + Unpin> Buffered<'a, R> {
pub(crate) const fn new(src: &'a mut R, leftover: Vec<u8>) -> Self {
Self {
src,
buf: leftover,
pos: 0,
}
}
fn available(&self) -> &[u8] {
&self.buf[self.pos..]
}
fn consume(&mut self, n: usize) {
self.pos += n;
if self.pos == self.buf.len() {
self.buf.clear();
self.pos = 0;
}
}
async fn fill(&mut self) -> std::io::Result<bool> {
let mut next = vec![0u8; CHUNK];
let n = self.src.read(&mut next).await?;
if n == 0 {
return Ok(false);
}
next.truncate(n);
self.buf.drain(..self.pos);
self.pos = 0;
self.buf.extend_from_slice(&next);
Ok(true)
}
async fn read_line(&mut self) -> std::io::Result<Vec<u8>> {
loop {
if let Some(end) = find_crlf(self.available()) {
let line = self.available()[..end].to_vec();
self.consume(end + 2);
return Ok(line);
}
if self.available().len() > MAX_CHUNK_LINE {
return Err(std::io::Error::other("chunk line is too long"));
}
if !self.fill().await? {
return Err(unexpected_eof());
}
}
}
}
pub(crate) async fn forward<R, W>(
src: &mut Buffered<'_, R>,
dst: &mut W,
framing: Framing,
) -> std::io::Result<u64>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
match framing {
Framing::Empty => Ok(0),
Framing::Length(n) => forward_exact(src, dst, n).await,
Framing::Chunked => forward_chunked(src, dst).await,
Framing::UntilClose => forward_to_close(src, dst).await,
}
}
async fn forward_exact<R, W>(
src: &mut Buffered<'_, R>,
dst: &mut W,
mut remaining: u64,
) -> std::io::Result<u64>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
let total = remaining;
while remaining > 0 {
if src.available().is_empty() && !src.fill().await? {
return Err(unexpected_eof());
}
let take = usize::try_from(remaining)
.unwrap_or(usize::MAX)
.min(src.available().len());
dst.write_all(&src.available()[..take]).await?;
dst.flush().await?;
src.consume(take);
remaining -= take as u64;
}
Ok(total)
}
async fn forward_chunked<R, W>(src: &mut Buffered<'_, R>, dst: &mut W) -> std::io::Result<u64>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
let mut total = 0u64;
loop {
let line = src.read_line().await?;
let size_text = line.split(|&b| b == b';').next().unwrap_or(&[]);
let size_text = std::str::from_utf8(size_text)
.map_err(|_| std::io::Error::other("chunk size is not ASCII"))?
.trim();
if size_text.is_empty() || !size_text.bytes().all(|b| b.is_ascii_hexdigit()) {
return Err(std::io::Error::other("chunk size is not hexadecimal"));
}
let size = u64::from_str_radix(size_text, 16)
.map_err(|_| std::io::Error::other("chunk size is not hexadecimal"))?;
dst.write_all(&line).await?;
dst.write_all(b"\r\n").await?;
if size == 0 {
loop {
let trailer = src.read_line().await?;
if trailer.is_empty() {
dst.write_all(b"\r\n").await?;
break;
}
if !dropped(&trailer) {
dst.write_all(&trailer).await?;
dst.write_all(b"\r\n").await?;
}
}
dst.flush().await?;
return Ok(total);
}
total += forward_exact(src, dst, size).await?;
let terminator = src.read_line().await?;
if !terminator.is_empty() {
return Err(std::io::Error::other("chunk data is not CRLF-terminated"));
}
dst.write_all(b"\r\n").await?;
dst.flush().await?;
}
}
async fn forward_to_close<R, W>(src: &mut Buffered<'_, R>, dst: &mut W) -> std::io::Result<u64>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
let mut total = 0u64;
loop {
if src.available().is_empty() && !src.fill().await? {
dst.flush().await?;
return Ok(total);
}
let n = src.available().len();
dst.write_all(src.available()).await?;
dst.flush().await?;
src.consume(n);
total += n as u64;
}
}
fn find_crlf(buf: &[u8]) -> Option<usize> {
buf.windows(2).position(|w| w == b"\r\n")
}
fn dropped(line: &[u8]) -> bool {
if line.contains(&b'\n') || line.contains(&b'\r') {
return true;
}
line.iter().position(|&b| b == b':').is_none_or(|colon| {
std::str::from_utf8(&line[..colon]).is_ok_and(headers::is_forbidden_in_trailer)
})
}
fn unexpected_eof() -> std::io::Error {
std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
"the body ended before its declared length",
)
}
#[cfg(test)]
#[path = "body_tests.rs"]
mod body_tests;