diff --git a/examples/custom-future-multiple-bodies/Cargo.toml b/examples/custom-future-multiple-bodies/Cargo.toml new file mode 100644 index 00000000..01fe130b --- /dev/null +++ b/examples/custom-future-multiple-bodies/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "custom-future-multiple-bodies" +version = "0.1.0" +authors = ["Tower Maintainers "] +edition = "2021" +publish = false +license = "MIT" + +[dependencies] +bytes = "1.12.0" +http = "1.4.2" +http-body-util = "0.1.3" +hyper = "1.10.1" +pin-project-lite = "0.2.17" +tokio = { version = "1.32.0", features = ["full"] } +tower = { version = "0.5", features = ["full"] } diff --git a/examples/custom-future-multiple-bodies/README.md b/examples/custom-future-multiple-bodies/README.md new file mode 100644 index 00000000..d231c5d3 --- /dev/null +++ b/examples/custom-future-multiple-bodies/README.md @@ -0,0 +1,14 @@ +# Custom future with multiple bodies + +This example serves to demonstrate how to build a service that returns a +response in case of an event (in this instance, a missing HTTP header) while +leaving the inner service's response untouched. + +This requires wrapping the response body if the user wishes to leave the inner +service's body untouched. + +## Running the example + +``` +cargo run -p custom-future-multiple-bodies +``` diff --git a/examples/custom-future-multiple-bodies/src/main.rs b/examples/custom-future-multiple-bodies/src/main.rs new file mode 100644 index 00000000..22632ea9 --- /dev/null +++ b/examples/custom-future-multiple-bodies/src/main.rs @@ -0,0 +1,122 @@ +use bytes::Bytes; +use http::{Request, Response, StatusCode}; +use http_body_util::{Either, Full}; +use pin_project_lite::pin_project; +use std::{ + future::Future, + pin::Pin, + task::{ready, Context, Poll}, +}; +use tower::Service; + +use http_body_util::BodyExt; +use std::error::Error; +use tower::ServiceBuilder; +use tower::ServiceExt; + +// think of Either as an enum that implements Body if both arms implement Body +// this allows us to combine multple bodies +type ResponseBody = Either>; + +// some helper to go along with it +fn map_resp(resp: Response) -> Response> { + resp.map(|body| Either::Left(body)) +} + +fn new_err_resp(status: StatusCode, body: &'static str) -> Response> { + Response::builder() + .status(status) + .body(Either::Right(Full::from(body))) + .unwrap() +} + +#[derive(Clone)] +pub struct RequireHeader { + inner: S, + header_name: &'static str, +} + +impl RequireHeader { + pub fn new(inner: S, header_name: &'static str) -> Self { + Self { inner, header_name } + } +} + +impl Service> for RequireHeader +where + S: Service, Response = Response>, +{ + type Response = Response>; + type Error = S::Error; + type Future = RequireHeaderFuture; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Request) -> Self::Future { + if !req.headers().contains_key(self.header_name) { + return RequireHeaderFuture::HeaderNotFound; + } + RequireHeaderFuture::Ok { + fut: self.inner.call(req), + } + } +} + +pin_project! { + #[project = EnumProj] + pub enum RequireHeaderFuture { + Ok{ #[pin] fut: F }, + HeaderNotFound, + } +} + +impl Future for RequireHeaderFuture +where + F: Future, E>>, +{ + type Output = Result>, E>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match self.project() { + EnumProj::Ok { fut } => { + let res = ready!(fut.poll(cx)); + + // we use our helper to unify the response body types + // its also possible to return a custom error response here if the inner service failed + Poll::Ready(res.map(|resp| map_resp(resp))) + } + EnumProj::HeaderNotFound => { + Poll::Ready(Ok(new_err_resp(StatusCode::BAD_REQUEST, "missing header"))) + } + } + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let inner_service = tower::service_fn(|_req: Request>| async { + Ok::<_, std::convert::Infallible>(Response::new(Full::new(Bytes::from("Hello, World!")))) + }); + + let mut service = ServiceBuilder::new() + .layer_fn(|inner| RequireHeader::new(inner, "x-api-key")) + .service(inner_service); + + let req_bad = Request::builder().body(Full::::default()).unwrap(); + let res_bad = service.ready().await?.call(req_bad).await?; + assert_eq!(res_bad.status(), StatusCode::BAD_REQUEST); + + let body = res_bad.into_body().collect().await.unwrap(); + assert_eq!(body.to_bytes(), "missing header"); + + let req_good = Request::builder() + .header("x-api-key", "secret") + .body(Full::::default()) + .unwrap(); + let res_good = service.ready().await?.call(req_good).await?; + assert_eq!(res_good.status(), StatusCode::OK); + + Ok(()) +}