1
0
mirror of https://github.com/fafhrd91/actix-web synced 2024-11-27 17:52:56 +01:00

various fixes to MessageBody::complete_body (#2519)

This commit is contained in:
Ali MJ Al-Nasrawy 2021-12-17 01:25:10 +03:00 committed by GitHub
parent 156cc20ac8
commit a6d5776481
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 36 additions and 25 deletions

View File

@ -57,27 +57,7 @@ impl MessageBody for BoxBody {
} }
fn take_complete_body(&mut self) -> Bytes { fn take_complete_body(&mut self) -> Bytes {
debug_assert!( self.0.take_complete_body()
self.is_complete_body(),
"boxed type does not allow taking complete body; caller should make sure to \
call `is_complete_body` first",
);
// we do not have DerefMut access to call take_complete_body directly but since
// is_complete_body is true we should expect the entire bytes chunk in one poll_next
let waker = futures_util::task::noop_waker();
let mut cx = Context::from_waker(&waker);
match self.as_pin_mut().poll_next(&mut cx) {
Poll::Ready(Some(Ok(data))) => data,
_ => {
panic!(
"boxed type indicated it allows taking complete body but failed to \
return Bytes when polled",
);
}
}
} }
} }

View File

@ -134,7 +134,7 @@ mod foreign_impls {
impl<B> MessageBody for Box<B> impl<B> MessageBody for Box<B>
where where
B: MessageBody + Unpin, B: MessageBody + Unpin + ?Sized,
{ {
type Error = B::Error; type Error = B::Error;
@ -164,7 +164,7 @@ mod foreign_impls {
impl<B> MessageBody for Pin<Box<B>> impl<B> MessageBody for Pin<Box<B>>
where where
B: MessageBody, B: MessageBody + ?Sized,
{ {
type Error = B::Error; type Error = B::Error;
@ -175,10 +175,10 @@ mod foreign_impls {
#[inline] #[inline]
fn poll_next( fn poll_next(
mut self: Pin<&mut Self>, self: Pin<&mut Self>,
cx: &mut Context<'_>, cx: &mut Context<'_>,
) -> Poll<Option<Result<Bytes, Self::Error>>> { ) -> Poll<Option<Result<Bytes, Self::Error>>> {
self.as_mut().poll_next(cx) self.get_mut().as_mut().poll_next(cx)
} }
#[inline] #[inline]
@ -475,6 +475,16 @@ where
None => Poll::Ready(None), None => Poll::Ready(None),
} }
} }
#[inline]
fn is_complete_body(&self) -> bool {
self.body.is_complete_body()
}
#[inline]
fn take_complete_body(&mut self) -> Bytes {
self.body.take_complete_body()
}
} }
#[cfg(test)] #[cfg(test)]
@ -630,6 +640,27 @@ mod tests {
assert_eq!(Pin::new(&mut data).poll_next(&mut cx), Poll::Ready(None)); assert_eq!(Pin::new(&mut data).poll_next(&mut cx), Poll::Ready(None));
} }
#[test]
fn complete_body_combinators() {
use crate::body::{BoxBody, EitherBody};
let body = Bytes::from_static(b"test");
let body = BoxBody::new(body);
let body = EitherBody::<_, ()>::left(body);
let body = EitherBody::<(), _>::right(body);
let body = Box::new(body);
let body = Box::pin(body);
let mut body = body;
assert!(body.is_complete_body());
assert_eq!(body.take_complete_body(), b"test".as_ref());
// subsequent poll_next returns None
let waker = futures_util::task::noop_waker();
let mut cx = Context::from_waker(&waker);
assert!(Pin::new(&mut body).poll_next(&mut cx).map_err(drop) == Poll::Ready(None));
}
// down-casting used to be done with a method on MessageBody trait // down-casting used to be done with a method on MessageBody trait
// test is kept to demonstrate equivalence of Any trait // test is kept to demonstrate equivalence of Any trait
#[actix_rt::test] #[actix_rt::test]