axum-extra: make option_layer guarantee that the output body is axum::body::Body

axum-extra/CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning].
 - **breaking:** Remove unused `async-stream` feature, which was accidentally
   introduced as an implicit feature through an optional dependency which was no
   longer being used ([#3298])
+- **breaking:** `option_layer` now maps the `Response` body type to `axum::body::Body` ([#3469])
 
 [#3298]: https://github.com/tokio-rs/axum/pull/3298
 

axum-extra/src/middleware.rs

@@ -1,6 +1,7 @@
 //! Additional middleware utilities.
 
 use crate::either::Either;
+use axum::middleware::ResponseAxumBodyLayer;
 use tower_layer::Identity;
 
 /// Convert an `Option<Layer>` into a [`Layer`].
@@ -26,19 +27,58 @@ use tower_layer::Identity;
 ///
 /// # Difference between this and [`tower::util::option_layer`]
 ///
-/// [`tower::util::option_layer`] always changes the error type to [`BoxError`] which requires
+/// `axum_extra::middleware::option_layer` makes sure that the output `Body` is [`axum::body::Body`].
-/// using [`HandleErrorLayer`] when used with axum, even if the layer you're applying uses
-/// [`Infallible`].
-///
-/// `axum_extra::middleware::option_layer` on the other hand doesn't change the error type so can
-/// be applied directly.
 ///
 /// [`Layer`]: tower_layer::Layer
-/// [`BoxError`]: tower::BoxError
+pub fn option_layer<L>(layer: Option<L>) -> Either<(ResponseAxumBodyLayer, L), Identity> {
-/// [`HandleErrorLayer`]: axum::error_handling::HandleErrorLayer
-/// [`Infallible`]: std::convert::Infallible
-pub fn option_layer<L>(layer: Option<L>) -> Either<L, Identity> {
     layer
-        .map(Either::E1)
+        .map(|layer| Either::E1((ResponseAxumBodyLayer, layer)))
         .unwrap_or_else(|| Either::E2(Identity::new()))
 }
+
+#[cfg(test)]
+mod tests {
+    use std::{
+        convert::Infallible,
+        pin::Pin,
+        task::{Context, Poll},
+    };
+
+    use axum::{body::Body as AxumBody, Router};
+    use bytes::Bytes;
+    use http_body::Body as HttpBody;
+    use tower_http::map_response_body::MapResponseBodyLayer;
+
+    use super::option_layer;
+
+    #[test]
+    fn remap_response_body() {
+        struct BodyWrapper;
+
+        impl BodyWrapper {
+            fn new(_: AxumBody) -> Self {
+                Self
+            }
+        }
+
+        impl HttpBody for BodyWrapper {
+            type Data = Bytes;
+            type Error = Infallible;
+            fn poll_frame(
+                self: Pin<&mut Self>,
+                _cx: &mut Context<'_>,
+            ) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
+                unimplemented!()
+            }
+            fn is_end_stream(&self) -> bool {
+                unimplemented!()
+            }
+            fn size_hint(&self) -> http_body::SizeHint {
+                unimplemented!()
+            }
+        }
+        let _app: Router = Router::new().layer(option_layer(Some(MapResponseBodyLayer::new(
+            BodyWrapper::new,
+        ))));
+    }
+}