From 5c090dcb3e0189f4dea3a259137ae935aeaf49c1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Ml=C3=A1dek?= <david.mladek.cz@gmail.com>
Date: Fri, 12 Sep 2025 14:33:48 +0200
Subject: [PATCH] axum-extra: make `option_layer` guarantee that the output
 body is `axum::body::Body`

---
 axum-extra/CHANGELOG.md      |  1 +
 axum-extra/src/middleware.rs | 62 +++++++++++++++++++++++++++++-------
 2 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/axum-extra/CHANGELOG.md b/axum-extra/CHANGELOG.md
index d0e9d132..3c7c5671 100644
--- a/axum-extra/CHANGELOG.md
+++ b/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
 
diff --git a/axum-extra/src/middleware.rs b/axum-extra/src/middleware.rs
index 0303d484..2407a166 100644
--- a/axum-extra/src/middleware.rs
+++ b/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
-/// 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.
+/// `axum_extra::middleware::option_layer` makes sure that the output `Body` is [`axum::body::Body`].
 ///
 /// [`Layer`]: tower_layer::Layer
-/// [`BoxError`]: tower::BoxError
-/// [`HandleErrorLayer`]: axum::error_handling::HandleErrorLayer
-/// [`Infallible`]: std::convert::Infallible
-pub fn option_layer<L>(layer: Option<L>) -> Either<L, Identity> {
+pub fn option_layer<L>(layer: Option<L>) -> Either<(ResponseAxumBodyLayer, 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,
+        ))));
+    }
+}