diff --git a/http-body-util/src/future.rs b/http-body-util/src/future.rs new file mode 100644 index 0000000..3100ecd --- /dev/null +++ b/http-body-util/src/future.rs @@ -0,0 +1,403 @@ +use http_body::{Body, SizeHint}; +use std::{ + future::Future, + pin::Pin, + task::{Context, Poll}, +}; + +/// A [`Body`] backed by a fallible [`Future`]. +/// +/// This allows an `F`-typed future that will yield either a `B`-typed body, or an error, to be +/// polled as a body. This is particularly useful when you create a body through an asynchronous +/// computation of some sort. +/// +/// For example, sending a body over a oneshot channel or reading its contents from the filesystem. +#[derive(Debug)] +pub struct TryFutureBody { + inner: Inner, +} + +/// The inner state of a [`TryFutureBody`]. +/// +/// A future is polled until it either yields a body, or fails. +/// +/// ```text +/// ┌────────┐ ┌──────┐ +/// │ Future │ --> `poll_frame()`-+------------------------> │ Body │ +/// └────────┘ | `Poll::Ready(Ok(body))` └──────┘ +/// ↑ | | +/// | | | ┌────────┐ +/// +---------------+ +------------------------> │ Failed │ +/// `Poll::Pending` `Poll::Ready(Err(err))` └────────┘ +/// +/// ``` +#[derive(Debug)] +enum Inner { + /// The future is still being polled. + /// + /// When the body is in this state, the inner future has not yet resolved. When this body is + /// polled, this inner future will be polled. + Future(F), + /// The body has been yielded and is being polled. + /// + /// When the body is in this state, the future has already yielded a body that can now be read. + Body(B), + /// The future failed to yield a body. + Failed, +} + +// === impl TryFutureBody === + +impl TryFutureBody { + /// Wraps the provided future in a [`TryFutureBody`]. + pub fn new(future: F) -> Self { + Self { + inner: Inner::Future(future), + } + } +} + +impl Body for TryFutureBody +where + F: Future>, + B: http_body::Body, + E: Into, +{ + type Data = B::Data; + type Error = B::Error; + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + use self::proj::InnerProj; + + match self.as_mut().project() { + InnerProj::Failed => Poll::Ready(None), + InnerProj::Body(body) => body.poll_frame(cx), + InnerProj::Future(future) => match future.poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Ok(body)) => { + // We received the body. Put it into place, and then poll ourselves again. + let inner = Inner::Body(body); + self.set(Self { inner }); + self.poll_frame(cx) + } + Poll::Ready(Err(err)) => { + // There is no body. Mark ourselves as finished and return the error. + let inner = Inner::Failed; + self.set(Self { inner }); + Poll::Ready(Some(Err(err.into()))) + } + }, + } + } + + fn is_end_stream(&self) -> bool { + let Self { inner } = self; + match inner { + Inner::Future(_) => false, + Inner::Body(body) => body.is_end_stream(), + Inner::Failed => true, + } + } + + fn size_hint(&self) -> SizeHint { + let Self { inner } = self; + match inner { + Inner::Future(_) => SizeHint::new(), + Inner::Body(body) => body.size_hint(), + Inner::Failed => SizeHint::with_exact(0), + } + } +} + +/// Pinning projection for [`TryFutureBody`]. +/// +/// Similar to [`crate::either::proj`], this submodule includes code derived from the output +/// generated by [pin-project-lite]. +mod proj { + use super::{Inner, TryFutureBody}; + use std::{marker::PhantomData, pin::Pin}; + + /// A projection of a [pinned][std::pin::Pin] [`Inner`]. + pub(super) enum InnerProj<'pin, F, B> + where + TryFutureBody: 'pin, + { + Future(Pin<&'pin mut F>), + Body(Pin<&'pin mut B>), + Failed, + } + + // === impl TryFutureBody === + + impl TryFutureBody { + /// Returns an [`InnerProj<'pin, F, B>`] projection. + /// + /// This is used internally by [`TryFutureBody`] to access its inner future and body. + pub(super) fn project<'pin>(self: Pin<&'pin mut Self>) -> InnerProj<'pin, F, B> { + // Safety: + // + // We never move the inner future, or the inner body, out of the mutable reference + // we receive from `Pin::get_unchecked_mut()`. We project their "pinnedness" forwards + // into a `Pin<&mut F>` or a `Pin<&mut B>`, respectively. If the body is finished, + // there is no data that could be moved out. + // + // - https://doc.rust-lang.org/std/pin/struct.Pin.html#method.get_unchecked_mut + // + // For more information on structural pinning, see: + // + unsafe { + let Self { inner } = self.get_unchecked_mut(); + match inner { + Inner::Future(fut) => InnerProj::Future(Pin::new_unchecked(fut)), + Inner::Body(body) => InnerProj::Body(Pin::new_unchecked(body)), + Inner::Failed => InnerProj::Failed, + } + } + } + } + + #[allow(single_use_lifetimes)] + #[allow(unknown_lints)] + #[allow(clippy::used_underscore_binding)] + #[allow(missing_debug_implementations)] + const _: () = { + #[allow(non_snake_case)] + pub struct __Origin<'__pin, F, B> { + __dummy_lifetime: PhantomData<&'__pin ()>, + _Future: F, + _Body: B, + } + impl<'__pin, F, B> Unpin for TryFutureBody where __Origin<'__pin, F, B>: Unpin {} + + #[allow(unused)] + trait MustNotImplDrop {} + #[allow(drop_bounds)] + impl MustNotImplDrop for T {} + impl MustNotImplDrop for TryFutureBody {} + }; +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Full; + use bytes::Bytes; + use std::{convert::Infallible, future::ready, ops::Not}; + + #[test] + fn full_ready_body() { + let mut body = { + let data = Bytes::from_static(b"hello"); + let body = Full::::from(data); + let fut = ready(Ok::<_, Infallible>(body)); + TryFutureBody::new(fut) + }; + + // Confirm that hints are correct before we poll the future. + { + assert!( + body.is_end_stream().not(), + "stream is not over before future resolves" + ); + let hint = body.size_hint(); + assert_eq!( + hint.lower(), + 0, + "size hint has lower bound of 0 before future resolves" + ); + assert_eq!( + hint.upper(), + None, + "size hint has no upper bound before future resolves" + ); + } + + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + // Now poll the body. The future will resolve, and the underlying body will yield "hello". + { + let body = Pin::new(&mut body); + let Poll::Ready(Some(Ok(frame))) = body.poll_frame(&mut cx) else { + panic!("body should yield a frame when polled"); + }; + let data = frame.into_data().expect("frame should contain data"); + assert_eq!(data, "hello", "underlying body frames are returned"); + } + + // The body will yield None after the inner body has finished. + { + let body = Pin::new(&mut body); + let Poll::Ready(None) = body.poll_frame(&mut cx) else { + panic!("body should `Ready(None)` when polled"); + }; + } + + // Finally, show that the body properly reports that the stream is finished. + { + assert!( + body.is_end_stream(), + "stream is over after body is finished" + ); + let hint = body.size_hint(); + assert_eq!( + hint.upper(), + Some(0), + "size hint is upper bound of 0 after body is finished" + ); + } + } + + /// A [`Body`] that returns an `E`-typed error when polled. + struct ErrorBody { + error: Option, + } + + impl ErrorBody { + fn new(error: E) -> Self { + Self { error: Some(error) } + } + } + + impl Body for ErrorBody { + type Data = Bytes; + type Error = E; + + fn poll_frame( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let Self { error } = self.get_mut(); + let error = error.take().map(Err); + Poll::Ready(error) + } + + fn is_end_stream(&self) -> bool { + self.error.is_none() + } + + fn size_hint(&self) -> SizeHint { + if self.error.is_some() { + // Pretend there is a hint until the error is returned. + SizeHint::with_exact(42) + } else { + SizeHint::with_exact(0) + } + } + } + + /// Show that a body that returns an error will be processed correctly. + #[test] + fn error_body() { + type Error = &'static str; + const ERROR: Error = "houston we have a problem"; + + let mut body = { + let body = ErrorBody::new(ERROR); + let fut = ready(Ok::<_, Error>(body)); + TryFutureBody::new(fut) + }; + + // Confirm that hints are correct before we poll the future. + { + assert!( + body.is_end_stream().not(), + "stream is not over before future resolves" + ); + let hint = body.size_hint(); + assert_eq!( + hint.lower(), + 0, + "size hint has lower bound of 0 before future resolves" + ); + assert_eq!( + hint.upper(), + None, + "size hint has no upper bound before future resolves" + ); + } + + // Now poll the body. The future will resolve, and the underlying body will yield "hello". + { + let body = Pin::new(&mut body); + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + let Poll::Ready(Some(Err(err))) = body.poll_frame(&mut cx) else { + panic!("body should yield an error when polled"); + }; + assert_eq!(err, ERROR, "future errors are returned"); + } + + // Finally, show that the body properly reports that the stream is finished. + { + assert!( + body.is_end_stream(), + "stream is over after body is finished" + ); + let hint = body.size_hint(); + assert_eq!( + hint.upper(), + Some(0), + "size hint is upper bound of 0 after body is finished" + ); + } + } + + /// Show that a future that fails to yield a body will be processed correctly. + #[test] + fn error_future() { + const ERROR: &str = "there is no body"; + + let mut body = { + let fut = ready(Err::(ERROR.to_string())); + TryFutureBody::new(fut) + }; + + // Confirm that hints are correct before we poll the future. + { + assert!( + body.is_end_stream().not(), + "stream is not over before future resolves" + ); + let hint = body.size_hint(); + assert_eq!( + hint.lower(), + 0, + "size hint has lower bound of 0 before future resolves" + ); + assert_eq!( + hint.upper(), + None, + "size hint has no upper bound before future resolves" + ); + } + + // Now poll the body. The future will resolve, and the underlying body will yield "hello". + { + let body = Pin::new(&mut body); + let waker = futures_util::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + let Poll::Ready(Some(Err(err))) = body.poll_frame(&mut cx) else { + panic!("body should yield an error when polled"); + }; + assert_eq!(err, "there is no body", "future errors are returned"); + } + + // Finally, show that the body properly reports that the stream is finished. + { + assert!( + body.is_end_stream(), + "stream is over after body is finished" + ); + let hint = body.size_hint(); + assert_eq!( + hint.upper(), + Some(0), + "size hint is upper bound of 0 after body is finished" + ); + } + } +} diff --git a/http-body-util/src/lib.rs b/http-body-util/src/lib.rs index de4239b..60cc651 100644 --- a/http-body-util/src/lib.rs +++ b/http-body-util/src/lib.rs @@ -13,6 +13,7 @@ pub mod combinators; mod either; mod empty; mod full; +mod future; mod limited; mod stream; @@ -27,6 +28,7 @@ pub use self::collected::Collected; pub use self::either::Either; pub use self::empty::Empty; pub use self::full::Full; +pub use self::future::TryFutureBody; pub use self::limited::{LengthLimitError, Limited}; pub use self::stream::{BodyDataStream, BodyStream, StreamBody};