tower-rs · bassmanitram · Oct 23, 2023 · Oct 23, 2023 · Oct 23, 2023
diff --git a/tower-http/Cargo.toml b/tower-http/Cargo.toml
@@ -68,6 +68,7 @@ full = [
     "map-response-body",
     "metrics",
     "normalize-path",
+    "propagate-extension",
     "propagate-header",
     "redirect",
     "request-id",
@@ -91,6 +92,7 @@ map-request-body = []
 map-response-body = []
 metrics = ["tokio/time"]
 normalize-path = []
+propagate-extension = []
 propagate-header = []
 redirect = []
 request-id = ["uuid"]

diff --git a/tower-http/src/builder.rs b/tower-http/src/builder.rs
@@ -54,6 +54,16 @@ pub trait ServiceBuilderExt<L>: crate::sealed::Sealed<L> + Sized {
         header: HeaderName,
     ) -> ServiceBuilder<Stack<crate::propagate_header::PropagateHeaderLayer, L>>;
 
+    /// Propagate an extension from the request to the response.
+    ///
+    /// See [`tower_http::propagate_extension`] for more details.
+    ///
+    /// [`tower_http::propagate_extension`]: crate::propagate_extension
+    #[cfg(feature = "propagate-extension")]
+    fn propagate_extension<T>(
+        self
+    ) -> ServiceBuilder<Stack<crate::propagate_extension::PropagateExtensionLayer<T>, L>>;
+
     /// Add some shareable value to [request extensions].
     ///
     /// See [`tower_http::add_extension`] for more details.
@@ -380,6 +390,13 @@ impl<L> ServiceBuilderExt<L> for ServiceBuilder<L> {
         self.layer(crate::propagate_header::PropagateHeaderLayer::new(header))
     }
 
+    #[cfg(feature = "propagate-extension")]
+    fn propagate_extension<X>(
+        self,
+    ) -> ServiceBuilder<Stack<crate::propagate_extension::PropagateExtensionLayer<X>, L>> {
+        self.layer(crate::propagate_extension::PropagateExtensionLayer::<X>::new())
+    }
+
     #[cfg(feature = "add-extension")]
     fn add_extension<T>(
         self,

diff --git a/tower-http/src/compression/future.rs b/tower-http/src/compression/future.rs
@@ -73,6 +73,7 @@ where
                 CompressionBody::new(BodyInner::zstd(WrapBody::new(body, self.quality)))
             }
             #[cfg(feature = "fs")]
+            #[allow(unreachable_patterns)]
             (true, _) => {
                 // This should never happen because the `AcceptEncoding` struct which is used to determine
                 // `self.encoding` will only enable the different compression algorithms if the

diff --git a/tower-http/src/lib.rs b/tower-http/src/lib.rs
@@ -231,6 +231,9 @@ pub mod auth;
 #[cfg(feature = "set-header")]
 pub mod set_header;
 
+#[cfg(feature = "propagate-extension")]
+pub mod propagate_extension;
+
 #[cfg(feature = "propagate-header")]
 pub mod propagate_header;
 

diff --git a/tower-http/src/propagate_extension.rs b/tower-http/src/propagate_extension.rs
@@ -0,0 +1,248 @@
+//! Propagate an extension from the request to the response.
+//!
+//! This middleware is intended to wrap a Request->Response service handler that is _unaware_ of the
+//! extension. Consequently it _removes_ the extension from the request before forwarding the request, and then
+//! inserts it into the response when the response is ready. As a usage example, if you have pre-service mappers 
+//! that need to share state with post-service mappers, you can store the state in the Request extensions, 
+//! and this middleware will ensure that it is available to the post service mappers via the Response extensions.
+//!
+//! # Example
+//!
+//! ```rust
+//! use http::{Request, Response};
+//! use std::convert::Infallible;
+//! use tower::{Service, ServiceExt, ServiceBuilder, service_fn};
+//! use tower_http::add_extension::AddExtensionLayer;
+//! use tower_http::propagate_extension::PropagateExtensionLayer;
+//! use tower_http::ServiceBuilderExt;
+//! use hyper::Body;
+//!
+//! # #[tokio::main]
+//! # async fn main() -> Result<(), Box<dyn std::error::Error>> {
+//! async fn handle(req: Request<Body>) -> Result<Response<Body>, Infallible> {
+//!     // ...
+//!     # Ok(Response::new(Body::empty()))
+//! }
+//!
+//! //
+//! // Note that while the state object must _implement_ Clone, it should never actually
+//! // _be_ cloned due to the manner in which it is used within the middleware.
+//! //
+//! #[derive(Clone)]
+//! struct MyState {
+//!     state_message: String
+//! };
+//!
+//! let my_state = MyState { state_message: "propagated state".to_string() };
+//! 
+//! let mut svc = ServiceBuilder::new()
+//!     .add_extension(my_state) // any other way of adding the extension to the request is OK too
+//!     .propagate_extension::<MyState>()
+//!     .service_fn(handle);
+//!
+//! // Call the service.
+//! let request = Request::builder()
+//!     .body(Body::empty())?;
+//!
+//! let response = svc.ready().await?.call(request).await?;
+//!
+//! assert_eq!(response.extensions().get::<MyState>().unwrap().state_message, "propagated state");
+//! #
+//! # Ok(())
+//! # }
+//! ```
+
+use futures_util::ready;
+use http::{Request, Response};
+use pin_project_lite::pin_project;
+use std::future::Future;
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+	marker::PhantomData,
+};
+use tower_layer::Layer;
+use tower_service::Service;
+
+#[allow(unused_imports)]
+use tracing::{
+	trace,
+	debug,
+	info,
+	warn,
+	error,
+};
+
+/// Layer that applies [`PropagateExtension`] which propagates an extension from the request to the response.
+///
+/// This middleware is intended to wrap a Request->Response service handler that is _unaware_ of the
+/// extension. Consequently it _removes_ the extension from the request before forwarding the request, and then
+/// inserts it into the response when the response is ready. As a usage example, if you have pre-service mappers 
+/// that need to share state with post-service mappers, you can store the state in the Request extensions, 
+/// and this middleware will ensure that it is available to the post service mappers via the Response extensions.
+///
+/// See the [module docs](crate::propagate_extension) for more details.
+#[derive(Clone, Debug)]
+pub struct PropagateExtensionLayer<X> {
+	_phantom: PhantomData<X>
+}
+
+impl<X> PropagateExtensionLayer<X> {
+    /// Create a new [`PropagateExtensionLayer`].
+    pub fn new() -> Self {
+        Self { _phantom: PhantomData }
+    }
+}
+
+impl<S,X> Layer<S> for PropagateExtensionLayer<X> {
+    type Service = PropagateExtension<S,X>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        PropagateExtension::<S,X> {
+            inner,
+			_phantom: PhantomData
+        }
+    }
+}
+
+/// Middleware that propagates extensions from requests to responses.
+///
+/// If the extension is present on the request it'll be removed from the request and
+/// inserted into the response.
+///
+/// See the [module docs](crate::propagate_extension) for more details.
+#[derive(Clone,Debug)]
+pub struct PropagateExtension<S,X> {
+    inner: S,
+	_phantom: PhantomData<X>
+}
+
+impl<S,X> PropagateExtension<S,X> {
+    /// Create a new [`PropagateExtension`] that propagates the given extension type.
+    pub fn new(inner: S) -> Self {
+        Self { inner, _phantom: PhantomData }
+    }
+
+    define_inner_service_accessors!();
+
+    /// Returns a new [`Layer`] that wraps services with a `PropagateExtension` middleware.
+    ///
+    /// [`Layer`]: tower_layer::Layer
+    pub fn layer() -> PropagateExtensionLayer<X> {
+        PropagateExtensionLayer::<X>::new()
+    }
+}
+
+impl<ReqBody, ResBody, S, X> Service<Request<ReqBody>> for PropagateExtension<S,X>
+where
+	X: Sync + Send + 'static,
+    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = ResponseFuture<S::Future,X>;
+
+    #[inline]
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
+        let extension: Option<X> = req.extensions_mut().remove();
+        debug!("Removed state from request extensions. is_some? {}", extension.is_some());
+
+        ResponseFuture {
+            future: self.inner.call(req),
+            extension,
+        }
+    }
+}
+
+pin_project! {
+    /// Response future for [`PropagateExtension`].
+    #[derive(Debug)]
+    pub struct ResponseFuture<F,X> {
+        #[pin]
+        future: F,
+        extension: Option<X>,
+    }
+}
+
+impl<F, ResBody, E, X> Future for ResponseFuture<F,X>
+where
+	X: Sync + Send + 'static,
+    F: Future<Output = Result<Response<ResBody>, E>>, 
+{
+    type Output = F::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        let mut res = ready!(this.future.poll(cx)?);
+
+        if let Some(extension) = this.extension.take() {
+            debug!("Inserting state into response extensions");
+            res.extensions_mut().insert(extension);
+        } else {
+            debug!("No state to insert into response");
+        }
+
+        Poll::Ready(Ok(res))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+	use http::{Request, Response};
+	use std::convert::Infallible;
+	use tower::{Service, ServiceExt, ServiceBuilder};
+	use crate::add_extension::AddExtensionLayer;
+	use crate::builder::ServiceBuilderExt;
+	use hyper::Body;
+
+	async fn handle(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
+		Ok(Response::new(Body::empty()))
+	}
+
+    #[derive(Clone)]
+	struct MyState {
+		state_message: String
+	}
+
+	#[test]
+	fn basic_test() {
+
+		let my_state = MyState { state_message: "propagated state".to_string() };
+
+		let mut svc = ServiceBuilder::new()
+			.layer(AddExtensionLayer::new(my_state)) // any other way of adding the extension to the request is OK too
+			.layer(PropagateExtensionLayer::<MyState>::new())
+			.service_fn(handle);
+
+		let request = Request::builder().body(Body::empty()).expect("Expected an empty body");
+
+		// Call the service.
+		let ready = futures::executor::block_on(svc.ready()).expect("Expected the service to be ready");
+		let response = futures::executor::block_on(ready.call(request)).expect("Expected the service to be successful");
+		assert_eq!(response.extensions().get::<MyState>().unwrap().state_message, "propagated state");
+	}
+
+    #[test]
+	fn test_server_builder_ext() {
+
+		let my_state = MyState { state_message: "propagated state".to_string() };
+
+		let mut svc = ServiceBuilder::new()
+			.add_extension(my_state) // any other way of adding the extension to the request is OK too
+			.propagate_extension::<MyState>()
+			.service_fn(handle);
+
+		let request = Request::builder().body(Body::empty()).expect("Expected an empty body");
+
+		// Call the service.
+		let ready = futures::executor::block_on(svc.ready()).expect("Expected the service to be ready");
+		let response = futures::executor::block_on(ready.call(request)).expect("Expected the service to be successful");
+		assert_eq!(response.extensions().get::<MyState>().unwrap().state_message, "propagated state");
+	}
+}