use std::pin::Pin;
use std::task::{Context, Poll};
use bytes::{Buf, Bytes};
use hyper::body::{Body, Frame, SizeHint};
use crate::BoxError;
pub(super) struct H3ReqBody {
recv: h3::server::RequestStream<h3_quinn::RecvStream, Bytes>,
remaining: Option<u64>,
}
impl H3ReqBody {
pub(super) fn new(
recv: h3::server::RequestStream<h3_quinn::RecvStream, Bytes>,
content_length: Option<u64>,
) -> Self {
Self {
recv,
remaining: content_length,
}
}
}
impl Body for H3ReqBody {
type Data = Bytes;
type Error = BoxError;
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Bytes>, BoxError>>> {
let this = self.get_mut();
match this.recv.poll_recv_data(cx) {
Poll::Ready(Ok(Some(mut buf))) => {
let bytes = buf.copy_to_bytes(buf.remaining());
if let Some(remaining) = &mut this.remaining {
match remaining.checked_sub(bytes.len() as u64) {
Some(rest) => *remaining = rest,
None => {
return Poll::Ready(Some(Err(
"h3 request body exceeds its declared content-length".into(),
)));
}
}
}
Poll::Ready(Some(Ok(Frame::data(bytes))))
}
Poll::Ready(Ok(None)) => {
if let Some(n) = this.remaining
&& n != 0
{
return Poll::Ready(Some(Err(
"h3 request body shorter than its declared content-length".into(),
)));
}
Poll::Ready(None)
}
Poll::Ready(Err(e)) => Poll::Ready(Some(Err(Box::new(e)))),
Poll::Pending => Poll::Pending,
}
}
fn size_hint(&self) -> SizeHint {
match self.remaining {
Some(n) => SizeHint::with_exact(n),
None => SizeHint::default(),
}
}
}