diff --git a/http-body-util/src/combinators/chain.rs b/http-body-util/src/combinators/chain.rs new file mode 100644 index 0000000..fa17a2d --- /dev/null +++ b/http-body-util/src/combinators/chain.rs @@ -0,0 +1,713 @@ +use http::HeaderMap; +use http_body::{Body, Frame, SizeHint}; +use pin_project_lite::pin_project; +use std::{ + pin::Pin, + task::{Context, Poll}, +}; + +pin_project! { + /// A body that links two bodies together, in a chain. + /// + /// See [`BodyExt::chain()`] for more information. + #[project = ChainProj] + pub struct Chain { + #[pin] + inner: Inner, + } +} + +pin_project! { + #[project = InnerProj] + pub enum Inner { + First { + #[pin] + first: A, + second: Option, + }, + Second { + #[pin] + second: B, + trailers: Option, + }, + Finished, + } +} + +// === impl Chain === + +impl Chain { + /// Returns a "chained" body. + /// + /// The contents of the first provided body will precede the contents of the second body. + pub fn new(first: A, second: B) -> Self { + Self { + inner: Inner::First { + first, + second: Some(second), + }, + } + } +} + +impl Body for Chain +where + A: Body, + B: Body, + A::Error: Into, +{ + type Data = B::Data; + type Error = B::Error; + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let ChainProj { inner } = self.as_mut().project(); + match inner.project() { + InnerProj::First { first, second } => { + // Poll the first body until it yields a frame. + let frame = match first.poll_frame(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(res) => res, + }; + + let trailers: Option = match frame { + // There are no more frames in the first body. + None => None, + // The first body has yielded a frame. + Some(Ok(frame)) => { + match frame.into_trailers() { + // A `TRAILERS` frame was yielded... + Ok(trls) => Some(trls), + // We will return other kinds of frames, and continue to poll the + // first body next time around. + Err(frame) => return Poll::Ready(Some(Ok(frame))), + } + } + Some(Err(err)) => { + // The first body returned an error! We are now finished. + let inner = Inner::Finished; + self.set(Self { inner }); + return Poll::Ready(Some(Err(err.into()))); + } + }; + + // If we are here, the first body is complete. Prepare to poll the second body. + let second = second.take().unwrap(); + let inner = Inner::Second { second, trailers }; + self.set(Self { inner }); + self.poll_frame(cx) + } + InnerProj::Second { + second, + trailers: trailers_first, + } => { + // Poll the second body until it yields a frame. + let frame = match second.poll_frame(cx) { + Poll::Pending => return Poll::Pending, + Poll::Ready(res) => res, + }; + + let trailers_second: Option = match frame { + // There are no more frames in the second body. + None => None, + // The second body has yielded a frame. + Some(Ok(frame)) => match frame.into_trailers() { + // A `TRAILERS` frame was yielded... + Ok(trls) => Some(trls), + // We will return other kinds of frames, and continue to poll the + // second body next time around. + Err(frame) => return Poll::Ready(Some(Ok(frame))), + }, + Some(Err(err)) => { + // The second body returned an error! We are now finished. + let inner = Inner::Finished; + self.set(Self { inner }); + return Poll::Ready(Some(Err(err))); + } + }; + + // If we are here, the second body is complete. We should now return trailers + // that have been accumulated from the two bodies. + let trailers = match (trailers_first.take(), trailers_second) { + (Some(mut a), Some(b)) => { + a.extend(b); + Some(Ok(Frame::trailers(a))) + } + (Some(t), None) | (None, Some(t)) => Some(Ok(Frame::trailers(t))), + (None, None) => None, + }; + + // Mark the body as finished before returning the trailers. + let inner = Inner::Finished; + self.set(Self { inner }); + Poll::Ready(trailers) + } + InnerProj::Finished => Poll::Ready(None), + } + } + + fn is_end_stream(&self) -> bool { + let Self { inner } = self; + match inner { + Inner::First { .. } | Inner::Second { .. } => false, + Inner::Finished => true, + } + } + + fn size_hint(&self) -> SizeHint { + let Self { inner } = self; + match inner { + // If the first body is still being polled, return the sum of the two bodies' hints. + Inner::First { first, second } => { + let first = first.size_hint(); + let second = second.as_ref().map(Body::size_hint).unwrap_or_default(); + first + second + } + // If the second body is now being polled, forward its hint. + Inner::Second { second, .. } => second.size_hint(), + // If this body is finished, return a hint of 0. + Inner::Finished => { + let mut hint = SizeHint::new(); + hint.set_exact(0); + hint + } + } + } +} + +/// Unit tests for [`Chain`]. +#[cfg(test)] +mod chain_tests { + use crate::{BodyExt, Empty, Full}; + use bytes::Bytes; + use http::{HeaderMap, HeaderName, HeaderValue}; + use http_body::{Body, Frame}; + use std::{ + future::ready, + ops::Not, + pin::Pin, + task::{Context, Poll}, + }; + + #[tokio::test] + async fn two_empty() { + let mut body = { + type Data = Bytes; + let first = Empty::::new(); + let second = Empty::::new(); + first.chain(second) + }; + + assert!( + body.is_end_stream().not(), + "body is not finished until polled" + ); + assert_eq!(body.size_hint().lower(), 0, "empty bodies size hint is 0"); + assert_eq!( + body.size_hint().upper(), + Some(0), + "empty bodies size hint is 0" + ); + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!(body.is_end_stream(), "body is finished after being polled"); + } + + #[tokio::test] + async fn full_then_empty() { + let mut body = { + type Data = Bytes; + let first = Full::new("hello ".into()); + let second = Empty::::new(); + first.chain(second) + }; + + assert!( + body.is_end_stream().not(), + "body is not finished until polled" + ); + assert_eq!(body.size_hint().lower(), "hello ".len() as u64); + assert_eq!(body.size_hint().upper(), Some("hello ".len() as u64)); + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "hello "); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!( + body.is_end_stream().not(), + "body is not finished until second body is polled" + ); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!(body.is_end_stream(), "body is finished after being polled"); + } + + #[tokio::test] + async fn empty_then_full() { + let mut body = { + type Data = Bytes; + let first = Empty::::new(); + let second = Full::new("world!".into()); + first.chain(second) + }; + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + assert!( + body.is_end_stream().not(), + "body is not finished until polled" + ); + assert_eq!(body.size_hint().lower(), "world!".len() as u64); + assert_eq!(body.size_hint().upper(), Some("world!".len() as u64)); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "world!"); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!( + body.is_end_stream().not(), + "body is not finished until second body is finished" + ); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!(body.is_end_stream(), "body is finished after being polled"); + } + + #[tokio::test] + async fn two_bodies_chain() { + let mut body = { + type Data = Bytes; + let first = Full::::new("hello ".into()); + let second = Full::::new("world!".into()); + first.chain(second) + }; + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + assert!( + body.is_end_stream().not(), + "body is not finished until polled" + ); + assert_eq!(body.size_hint().lower(), "hello world!".len() as u64); + assert_eq!(body.size_hint().upper(), Some("hello world!".len() as u64)); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "hello "); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "world!"); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!( + body.is_end_stream().not(), + "body is not finished until second body is finished" + ); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!(body.is_end_stream(), "body is finished after being polled"); + } + + use self::error_body::*; + mod error_body { + use super::*; + + /// A [`Body`] that returns an error. + pub(super) struct ErrorBody { + error: Option, + } + + #[derive(Debug)] + pub(super) struct ErrorBodyError(pub(super) &'static str); + + // === ErrorBody === + + impl ErrorBody { + pub(super) fn new(msg: &'static str) -> Self { + Self { + error: Some(ErrorBodyError(msg)), + } + } + } + + impl Body for ErrorBody { + type Data = Bytes; + type Error = ErrorBodyError; + + fn poll_frame( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let Self { error } = self.get_mut(); + + Poll::Ready(error.take().map(Result::Err)) + } + } + } + + use self::mock_body::*; + mod mock_body { + //! NB: if `Full` was generic over its error, we could remove this. see hyperium/http-body#85. + use super::*; + use std::marker::PhantomData; + + pub(super) struct MockBody { + data: Option, + _marker: PhantomData, + } + + impl MockBody { + pub(super) fn new(data: impl Into) -> Self { + Self { + data: Some(data.into()), + _marker: PhantomData, + } + } + } + + impl Body for MockBody + where + E: Unpin, + { + type Data = Bytes; + type Error = E; + + fn poll_frame( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let Self { data, _marker } = self.get_mut(); + let frame = data.take().map(Frame::data).map(Result::Ok); + Poll::Ready(frame) + } + } + } + + #[tokio::test] + async fn second_body_error() { + let mut body = { + let first = MockBody::::new("hello "); + let second = ErrorBody::new("failure"); + first.chain(second) + }; + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "hello "); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Err(ErrorBodyError("failure")))) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!( + body.is_end_stream(), + "body is finished after returning an error" + ); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + } + + #[tokio::test] + async fn first_body_error() { + let mut body = { + let first = ErrorBody::new("failure"); + let second = MockBody::::new("world!"); + first.chain(second) + }; + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Err(ErrorBodyError("failure")))) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!( + body.is_end_stream(), + "body is finished after returning an error" + ); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + } + + #[tokio::test] + async fn first_body_trailers_are_propagated() { + let trailers = { + let trls = vec![( + HeaderName::from_static("fourty-two"), + HeaderValue::from_static("42"), + )] + .into_iter() + .collect::(); + ready(Some(Ok(trls))) + }; + + let mut body = { + type Data = Bytes; + let first = Full::::new("hello ".into()).with_trailers(trailers); + let second = Full::::new("world!".into()); + first.chain(second) + }; + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "hello "); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "world!"); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let trailers = frame.into_trailers().expect("should yield trailers"); + let value = trailers + .get("fourty-two") + .map(HeaderValue::to_str) + .transpose() + .expect("header is a string") + .expect("header exists"); + assert_eq!(value, "42"); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!( + body.is_end_stream(), + "body is finished after returning trailers" + ); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + } + + #[tokio::test] + async fn second_body_trailers_are_propagated() { + let trailers = { + let trls = vec![( + HeaderName::from_static("fourty-two"), + HeaderValue::from_static("42"), + )] + .into_iter() + .collect::(); + ready(Some(Ok(trls))) + }; + + let mut body = { + type Data = Bytes; + let first = Full::::new("hello ".into()); + let second = Full::::new("world!".into()).with_trailers(trailers); + first.chain(second) + }; + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "hello "); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "world!"); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let trailers = frame.into_trailers().expect("should yield trailers"); + let value = trailers + .get("fourty-two") + .map(HeaderValue::to_str) + .transpose() + .expect("header is a string") + .expect("header exists"); + assert_eq!(value, "42"); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + assert!( + body.is_end_stream(), + "body is finished after returning trailers" + ); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + } + + #[tokio::test] + async fn body_trailers_are_consolidated() { + type Data = Bytes; + + let first = { + let data = Full::::new("hello ".into()); + let trls = vec![ + ( + HeaderName::from_static("fourty-two"), + HeaderValue::from_static("42"), + ), + ( + HeaderName::from_static("both"), + HeaderValue::from_static("alpha"), + ), + ] + .into_iter() + .collect::(); + data.with_trailers(ready(Some(Ok(trls)))) + }; + + let second = { + let data = Full::::new("world!".into()); + let trls = vec![ + ( + HeaderName::from_static("ten"), + HeaderValue::from_static("10"), + ), + ( + HeaderName::from_static("both"), + HeaderValue::from_static("beta"), + ), + ] + .into_iter() + .collect::(); + data.with_trailers(ready(Some(Ok(trls)))) + }; + + let mut body = first.chain(second); + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "hello "); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let frame = frame.into_data().expect("should yield data"); + assert_eq!(frame, "world!"); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(Some(Ok(frame))) => { + let trailers = frame.into_trailers().expect("should yield trailers"); + assert_eq!(trailers.len(), 3); + + let value = trailers + .get("fourty-two") + .map(HeaderValue::to_str) + .transpose() + .expect("header is a string") + .expect("header exists"); + assert_eq!(value, "42"); + + let value = trailers + .get("ten") + .map(HeaderValue::to_str) + .transpose() + .expect("header is a string") + .expect("header exists"); + assert_eq!(value, "10"); + + // In the event of conflicts, the second body's headers take precendent. + let both = trailers + .get_all("both") + .iter() + .map(|s| s.to_str().unwrap().to_string()) + .collect::>(); + assert_eq!(both.len(), 1); + assert_eq!(both, ["beta"]); + } + other => panic!("unexpected poll outcome: {:?}", other), + } + + match Pin::new(&mut body).poll_frame(&mut cx) { + Poll::Ready(None) => {} + other => panic!("unexpected poll outcome: {:?}", other), + } + } +} diff --git a/http-body-util/src/combinators/mod.rs b/http-body-util/src/combinators/mod.rs index c27f516..463b544 100644 --- a/http-body-util/src/combinators/mod.rs +++ b/http-body-util/src/combinators/mod.rs @@ -1,6 +1,7 @@ //! Combinators for the `Body` trait. mod box_body; +mod chain; mod collect; mod frame; mod fuse; @@ -12,6 +13,7 @@ mod with_trailers; pub use self::{ box_body::{BoxBody, UnsyncBoxBody}, + chain::Chain, collect::Collect, frame::Frame, fuse::Fuse, diff --git a/http-body-util/src/lib.rs b/http-body-util/src/lib.rs index de4239b..b2488c5 100644 --- a/http-body-util/src/lib.rs +++ b/http-body-util/src/lib.rs @@ -182,6 +182,66 @@ pub trait BodyExt: http_body::Body { { combinators::Fuse::new(self) } + + /// Takes two bodies and creates a new body "chaining" them together. + /// + /// Similar to [`std::iter::Iterator::chain()`], this method will return a new + /// [`Body`][http_body::Body] that emits the contents of the first body, and then emits the + /// contents of the second body. + /// + /// If the first body returns an error, the body will consider itself finished and will not + /// poll the second body. + /// + /// Trailers yielded by the first body are buffered while the second body is polled, and then + /// merged with any trailers yielded by the second body via [`http::HeaderMap::extend()`]. + /// Header values from the second body take precedence in the event of any conflicts. + /// + /// # Examples + /// + /// ``` + /// # use bytes::Bytes; + /// # use http_body_util::{BodyExt, Full}; + /// # + /// #[tokio::main] + /// async fn main() { + /// let first = Full::new(Bytes::from("hello ")); + /// let second = Full::new(Bytes::from("world!")); + /// let chained = first.chain(second); + /// + /// let collected = chained.collect().await.unwrap(); + /// assert_eq!(collected.to_bytes(), "hello world!"); + /// } + /// ``` + /// + /// ``` + /// # use bytes::Bytes; + /// # use http::{HeaderMap, HeaderName, HeaderValue}; + /// # use http_body_util::{BodyExt, Full, Empty}; + /// # + /// #[tokio::main] + /// async fn main() { + /// let mut trailers = HeaderMap::new(); + /// trailers.insert( + /// HeaderName::from_static("name"), + /// HeaderValue::from_static("value"), + /// ); + /// let trailers = std::future::ready(Some(Ok(trailers))); + /// + /// let first = Full::new(Bytes::from("trailers")); + /// let second = Full::new(Bytes::from(" too!")); + /// let chained = first.with_trailers(trailers).chain(second); + /// + /// let collected = chained.collect().await.unwrap(); + /// assert_eq!(collected.trailers().unwrap()["name"], "value"); + /// assert_eq!(collected.to_bytes(), "trailers too!"); + /// } + /// ``` + fn chain(self, other: B) -> combinators::Chain + where + Self: Sized, + { + combinators::Chain::new(self, other) + } } impl BodyExt for T where T: http_body::Body {}