Don't allow extracting MatchedPath in fallbacks (#1934)

tokio-rs · Apr 17, 2023 · 2c87d65 · 2c87d65
1 parent e97462d
commit 2c87d65
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 20 deletions.
diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md
@@ -7,7 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Unreleased
 
-- None.
+- **fixed:** Don't allow extracting `MatchedPath` in fallbacks ([#1934])
+- **fixed:** Fix panic if `Router` with something nested at `/` was used as a fallback ([#1934])
+
+[#1934]: https://github.com/tokio-rs/axum/pull/1934
 
 # 0.6.15 (12. April, 2023)
 

diff --git a/axum/src/extract/matched_path.rs b/axum/src/extract/matched_path.rs
@@ -176,6 +176,7 @@ mod tests {
         Router,
     };
     use http::{Request, StatusCode};
+    use hyper::Body;
 
     #[crate::test]
     async fn extracting_on_handler() {
@@ -353,4 +354,19 @@ mod tests {
         let res = client.get("/foo/bar").send().await;
         assert_eq!(res.status(), StatusCode::OK);
     }
+
+    #[crate::test]
+    async fn cant_extract_in_fallback() {
+        async fn handler(path: Option<MatchedPath>, req: Request<Body>) {
+            assert!(path.is_none());
+            assert!(req.extensions().get::<MatchedPath>().is_none());
+        }
+
+        let app = Router::new().fallback(handler);
+
+        let client = TestClient::new(app);
+
+        let res = client.get("/foo/bar").send().await;
+        assert_eq!(res.status(), StatusCode::OK);
+    }
 }
diff --git a/axum/src/routing/mod.rs b/axum/src/routing/mod.rs
@@ -57,8 +57,8 @@ pub(crate) struct RouteId(u32);
 /// The router type for composing handlers and services.
 #[must_use]
 pub struct Router<S = (), B = Body> {
-    path_router: PathRouter<S, B>,
-    fallback_router: PathRouter<S, B>,
+    path_router: PathRouter<S, B, false>,
+    fallback_router: PathRouter<S, B, true>,
     default_fallback: bool,
 }
 
@@ -499,7 +499,7 @@ impl<S, B> fmt::Debug for Endpoint<S, B> {
     }
 }
 
-struct SuperFallback<S, B>(SyncWrapper<PathRouter<S, B>>);
+struct SuperFallback<S, B>(SyncWrapper<PathRouter<S, B, true>>);
 
 #[test]
 #[allow(warnings)]

diff --git a/axum/src/routing/path_router.rs b/axum/src/routing/path_router.rs
@@ -1,4 +1,4 @@
-use crate::body::{Body, HttpBody};
+use crate::body::HttpBody;
 use axum_core::response::IntoResponse;
 use http::Request;
 use matchit::MatchError;
@@ -11,13 +11,13 @@ use super::{
     RouteId, NEST_TAIL_PARAM,
 };
 
-pub(super) struct PathRouter<S = (), B = Body> {
+pub(super) struct PathRouter<S, B, const IS_FALLBACK: bool> {
     routes: HashMap<RouteId, Endpoint<S, B>>,
     node: Arc<Node>,
     prev_route_id: RouteId,
 }
 
-impl<S, B> PathRouter<S, B>
+impl<S, B, const IS_FALLBACK: bool> PathRouter<S, B, IS_FALLBACK>
 where
     B: HttpBody + Send + 'static,
     S: Clone + Send + Sync + 'static,
@@ -107,7 +107,10 @@ where
         Ok(())
     }
 
-    pub(super) fn merge(&mut self, other: PathRouter<S, B>) -> Result<(), Cow<'static, str>> {
+    pub(super) fn merge(
+        &mut self,
+        other: PathRouter<S, B, IS_FALLBACK>,
+    ) -> Result<(), Cow<'static, str>> {
         let PathRouter {
             routes,
             node,
@@ -131,7 +134,7 @@ where
     pub(super) fn nest(
         &mut self,
         path: &str,
-        router: PathRouter<S, B>,
+        router: PathRouter<S, B, IS_FALLBACK>,
     ) -> Result<(), Cow<'static, str>> {
         let prefix = validate_nest_path(path);
 
@@ -193,7 +196,7 @@ where
         Ok(())
     }
 
-    pub(super) fn layer<L, NewReqBody>(self, layer: L) -> PathRouter<S, NewReqBody>
+    pub(super) fn layer<L, NewReqBody>(self, layer: L) -> PathRouter<S, NewReqBody, IS_FALLBACK>
     where
         L: Layer<Route<B>> + Clone + Send + 'static,
         L::Service: Service<Request<NewReqBody>> + Clone + Send + 'static,
@@ -250,7 +253,7 @@ where
         }
     }
 
-    pub(super) fn with_state<S2>(self, state: S) -> PathRouter<S2, B> {
+    pub(super) fn with_state<S2>(self, state: S) -> PathRouter<S2, B, IS_FALLBACK> {
         let routes = self
             .routes
             .into_iter()
@@ -293,12 +296,14 @@ where
             Ok(match_) => {
                 let id = *match_.value;
 
-                #[cfg(feature = "matched-path")]
-                crate::extract::matched_path::set_matched_path_for_request(
-                    id,
-                    &self.node.route_id_to_path,
-                    req.extensions_mut(),
-                );
+                if !IS_FALLBACK {
+                    #[cfg(feature = "matched-path")]
+                    crate::extract::matched_path::set_matched_path_for_request(
+                        id,
+                        &self.node.route_id_to_path,
+                        req.extensions_mut(),
+                    );
+                }
 
                 url_params::insert_url_params(req.extensions_mut(), match_.params);
 
@@ -347,7 +352,7 @@ where
     }
 }
 
-impl<B, S> Default for PathRouter<S, B> {
+impl<B, S, const IS_FALLBACK: bool> Default for PathRouter<S, B, IS_FALLBACK> {
     fn default() -> Self {
         Self {
             routes: Default::default(),
@@ -357,7 +362,7 @@ impl<B, S> Default for PathRouter<S, B> {
     }
 }
 
-impl<S, B> fmt::Debug for PathRouter<S, B> {
+impl<S, B, const IS_FALLBACK: bool> fmt::Debug for PathRouter<S, B, IS_FALLBACK> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("PathRouter")
             .field("routes", &self.routes)
@@ -366,7 +371,7 @@ impl<S, B> fmt::Debug for PathRouter<S, B> {
     }
 }
 
-impl<S, B> Clone for PathRouter<S, B> {
+impl<S, B, const IS_FALLBACK: bool> Clone for PathRouter<S, B, IS_FALLBACK> {
     fn clone(&self) -> Self {
         Self {
             routes: self.routes.clone(),

diff --git a/axum/src/routing/tests/fallback.rs b/axum/src/routing/tests/fallback.rs
@@ -225,3 +225,19 @@ async fn nest_fallback_on_inner() {
     assert_eq!(res.status(), StatusCode::NOT_FOUND);
     assert_eq!(res.text().await, "inner fallback");
 }
+
+// https://github.com/tokio-rs/axum/issues/1931
+#[crate::test]
+async fn doesnt_panic_if_used_with_nested_router() {
+    async fn handler() {}
+
+    let routes_static =
+        Router::new().nest_service("/", crate::routing::get_service(handler.into_service()));
+
+    let routes_all = Router::new().fallback_service(routes_static);
+
+    let client = TestClient::new(routes_all);
+
+    let res = client.get("/foobar").send().await;
+    assert_eq!(res.status(), StatusCode::OK);
+}