From decdd4c948dbfdfc7d31e7d9e06fa6799886252d Mon Sep 17 00:00:00 2001
From: Kai Jewson <kai.jewson@gmail.com>
Date: Sat, 27 Nov 2021 09:20:31 +0000
Subject: [PATCH] Avoid double-boxing bodies where possible (#569)

* Avoid double-boxing bodies where possible

* Extract `try_downcast` to `util` module
---
 axum/src/body/mod.rs    |  4 ++--
 axum/src/routing/mod.rs | 16 +---------------
 axum/src/util.rs        | 19 +++++++++++++++++++
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/axum/src/body/mod.rs b/axum/src/body/mod.rs
index 5705ac57..6b449cf4 100644
--- a/axum/src/body/mod.rs
+++ b/axum/src/body/mod.rs
@@ -1,6 +1,6 @@
 //! HTTP body utilities.
 
-use crate::{BoxError, Error};
+use crate::{util::try_downcast, BoxError, Error};
 
 mod stream_body;
 
@@ -27,7 +27,7 @@ where
     B: http_body::Body<Data = Bytes> + Send + 'static,
     B::Error: Into<BoxError>,
 {
-    body.map_err(Error::new).boxed_unsync()
+    try_downcast(body).unwrap_or_else(|body| body.map_err(Error::new).boxed_unsync())
 }
 
 pub(crate) fn empty() -> BoxBody {
diff --git a/axum/src/routing/mod.rs b/axum/src/routing/mod.rs
index f660a5cd..f384999a 100644
--- a/axum/src/routing/mod.rs
+++ b/axum/src/routing/mod.rs
@@ -8,7 +8,7 @@ use crate::{
         MatchedPath, OriginalUri,
     },
     routing::strip_prefix::StripPrefix,
-    util::{ByteStr, PercentDecodedByteStr},
+    util::{try_downcast, ByteStr, PercentDecodedByteStr},
     BoxError,
 };
 use bytes::Bytes;
@@ -657,20 +657,6 @@ impl<B, E> Fallback<B, E> {
     }
 }
 
-fn try_downcast<T, K>(k: K) -> Result<T, K>
-where
-    T: 'static,
-    K: Send + 'static,
-{
-    use std::any::Any;
-
-    let k = Box::new(k) as Box<dyn Any + Send + 'static>;
-    match k.downcast() {
-        Ok(t) => Ok(*t),
-        Err(other) => Err(*other.downcast().unwrap()),
-    }
-}
-
 enum Endpoint<B> {
     MethodRouter(MethodRouter<B>),
     Route(Route<B>),
diff --git a/axum/src/util.rs b/axum/src/util.rs
index 61419c8e..71c50e77 100644
--- a/axum/src/util.rs
+++ b/axum/src/util.rs
@@ -65,3 +65,22 @@ pin_project! {
         B { #[pin] inner: B },
     }
 }
+
+pub(crate) fn try_downcast<T, K>(k: K) -> Result<T, K>
+where
+    T: 'static,
+    K: Send + 'static,
+{
+    let mut k = Some(k);
+    if let Some(k) = <dyn std::any::Any>::downcast_mut::<Option<T>>(&mut k) {
+        Ok(k.take().unwrap())
+    } else {
+        Err(k.unwrap())
+    }
+}
+
+#[test]
+fn test_try_downcast() {
+    assert_eq!(try_downcast::<i32, _>(5_u32), Err(5_u32));
+    assert_eq!(try_downcast::<i32, _>(5_i32), Ok(5_i32));
+}