diff --git a/examples/README.md b/examples/README.md index e10355bf..db850ffb 100644 --- a/examples/README.md +++ b/examples/README.md @@ -2,5 +2,7 @@ This folder contains various examples of how to use tower-http. +- `axum-key-value-store`: A key/value store with an HTTP API built with axum. - `warp-key-value-store`: A key/value store with an HTTP API built with warp. - `tonic-key-value-store`: A key/value store with a gRPC API and client built with tonic. +- `custom-middleware-either`: A custom middleware showing how to unify body types using `Either`. diff --git a/examples/custom-middleware-either/Cargo.toml b/examples/custom-middleware-either/Cargo.toml new file mode 100644 index 00000000..df30f32c --- /dev/null +++ b/examples/custom-middleware-either/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "custom-middleware-either" +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" +tokio = { version = "1.32.0", features = ["full"] } +tower = { version = "0.5", features = ["full"] } diff --git a/examples/custom-middleware-either/README.md b/examples/custom-middleware-either/README.md new file mode 100644 index 00000000..a4c507a1 --- /dev/null +++ b/examples/custom-middleware-either/README.md @@ -0,0 +1,17 @@ +# custom-middleware-either + +Because a Tower `Service` requires a single, concrete `Response` type, you run into a type mismatch if your middleware returns an immediate error body in one branch, but defers to the inner service's body in the other. + +This example solves that problem by using `http_body_util::Either>` to gracefully unify the happy-path body and the early-rejection body, all without requiring dynamic allocations (boxing). + +## How it works + +- **Missing Header:** If the `x-api-key` header is missing, the middleware short-circuits the future using `std::future::ready` and immediately bounces the request with a `400 Bad Request` and a custom error body (`Either::Right`). +- **Valid Request:** If the header is present, the request is passed through to the inner service, which returns a `200 OK` (`Either::Left`). + +## Running the example + +``` +RUST_LOG=custom_middleware_either=trace,tower_http=trace \ + cargo run --bin custom-middleware-either +``` diff --git a/examples/custom-middleware-either/src/main.rs b/examples/custom-middleware-either/src/main.rs new file mode 100644 index 00000000..d5edd08a --- /dev/null +++ b/examples/custom-middleware-either/src/main.rs @@ -0,0 +1,91 @@ +use bytes::Bytes; +use http::{Request, Response, StatusCode}; +use http_body_util::{Either, Full}; +use std::{ + future::Future, + pin::Pin, + task::{Context, Poll}, +}; +use tower::{Service, ServiceBuilder, ServiceExt}; + +// We use `Either` here because this middleware might return a completely different +// body type (our error message) than the inner service returns. Both paths must +// resolve to a single concrete response type. +type ResponseBody = Either>; + +#[derive(Clone)] +pub struct RequireHeader { + inner: S, + header_name: &'static str, +} + +impl Service> for RequireHeader +where + S: Service, Response = Response> + Clone + Send + 'static, + S::Future: Send, + S::Error: Send, + ReqBody: Send + 'static, + ResBody: Send + 'static, +{ + type Response = Response>; + type Error = S::Error; + // We have to Box the Future dynamically here. Because we have two totally + // different return paths (the immediate error vs. waiting on the inner service) + // and Rust needs them to resolve to a single concrete type. + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Request) -> Self::Future { + // If they don't have the header, bounce the request immediately + if !req.headers().contains_key(self.header_name) { + let body = Full::from("Missing required header"); + let res = Response::builder() + .status(StatusCode::BAD_REQUEST) + // Either::Right explicitly tags our middleware-generated error body + .body(Either::Right(body)) + .unwrap(); + + return Box::pin(std::future::ready(Ok(res))); + } + + let mut inner = self.inner.clone(); + Box::pin(async move { + let res = inner.call(req).await?; + // Either::Left maps the inner service's successful body into our unified type + Ok(res.map(Either::Left)) + }) + } +} + +impl RequireHeader { + pub fn new(inner: S, header_name: &'static str) -> Self { + Self { inner, header_name } + } +} + +#[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?; + println!("Bad: {}", res_bad.status()); + + let req_good = Request::builder() + .header("x-api-key", "secret") + .body(Full::::default()) + .unwrap(); + let res_good = service.ready().await?.call(req_good).await?; + println!("Good: {}", res_good.status()); + + Ok(()) +}