Add propagate_extensions middleware #425
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
|
Member
There was a problem hiding this comment. We should move to |
||
| //! | ||
| //! # #[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; | ||
|
Member
There was a problem hiding this comment. We have |
||
| 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)] | ||
|
Member
There was a problem hiding this comment. Let's only include the one we need rather than allowing unused (looks like only debug!)? Let's also gate tracing imports/usage behind the |
||
| 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)] | ||
|
Member
There was a problem hiding this comment. We should derive |
||
| 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(); | ||
|
Member
There was a problem hiding this comment. Let's clone it rather than remove it, and document it as such. (If you think it is useful to have an optional flag to remove instead of clone, I'm fine with it, but I think it is the wrong default). That will bring us in line with Extensions anyway require |
||
| 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)] | ||
|
Member
There was a problem hiding this comment. Also need a test for the case where the targeted extension is missing |
||
| 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; | ||
|
Member
There was a problem hiding this comment. We should move to |
||
|
|
||
| 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"); | ||
|
Member
There was a problem hiding this comment. let's use |
||
| 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"); | ||
| } | ||
| } | ||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's pull this out of this PR, no real relationship. Happy to review a separate one if still needed.