Fix bugs around merging routers with nested fallbacks (#2096)

davidpdrsn · davidpdrsn · commit bb406b196a84 · 2023-07-16T12:43:24.000+02:00
diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md
@@ -7,7 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Unreleased
 
-- None.
+- **fixed:** Fix bugs around merging routers with nested fallbacks ([#2096])
+
+[#2096]: https://github.com/tokio-rs/axum/pull/2096
 
 # 0.6.18 (30. April, 2023)
 
diff --git a/axum/src/routing/mod.rs b/axum/src/routing/mod.rs
@@ -98,6 +98,7 @@ impl<S, B> fmt::Debug for Router<S, B> {
 pub(crate) const NEST_TAIL_PARAM: &str = "__private__axum_nest_tail_param";
 pub(crate) const NEST_TAIL_PARAM_CAPTURE: &str = "/*__private__axum_nest_tail_param";
 pub(crate) const FALLBACK_PARAM: &str = "__private__axum_fallback";
+pub(crate) const FALLBACK_PARAM_PATH: &str = "/*__private__axum_fallback";
 
 impl<S, B> Router<S, B>
 where
@@ -185,9 +186,12 @@ where
     where
         R: Into<Router<S, B>>,
     {
+        const PANIC_MSG: &str =
+            "Failed to merge fallbacks. This is a bug in axum. Please file an issue";
+
         let Router {
             path_router,
-            fallback_router: other_fallback,
+            fallback_router: mut other_fallback,
             default_fallback,
             catch_all_fallback,
         } = other.into();
@@ -198,16 +202,19 @@ where
             // both have the default fallback
             // use the one from other
             (true, true) => {
-                self.fallback_router = other_fallback;
+                self.fallback_router.merge(other_fallback).expect(PANIC_MSG);
             }
             // self has default fallback, other has a custom fallback
             (true, false) => {
-                self.fallback_router = other_fallback;
+                self.fallback_router.merge(other_fallback).expect(PANIC_MSG);
                 self.default_fallback = false;
             }
             // self has a custom fallback, other has a default
-            // nothing to do
-            (false, true) => {}
+            (false, true) => {
+                let fallback_router = std::mem::take(&mut self.fallback_router);
+                other_fallback.merge(fallback_router).expect(PANIC_MSG);
+                self.fallback_router = other_fallback;
+            }
             // both have a custom fallback, not allowed
             (false, false) => {
                 panic!("Cannot merge two `Router`s that both have a fallback")
diff --git a/axum/src/routing/path_router.rs b/axum/src/routing/path_router.rs
@@ -8,7 +8,7 @@ use tower_service::Service;
 
 use super::{
     future::RouteFuture, not_found::NotFound, strip_prefix::StripPrefix, url_params, Endpoint,
-    MethodRouter, Route, RouteId, FALLBACK_PARAM, NEST_TAIL_PARAM,
+    MethodRouter, Route, RouteId, FALLBACK_PARAM_PATH, NEST_TAIL_PARAM,
 };
 
 pub(super) struct PathRouter<S, B, const IS_FALLBACK: bool> {
@@ -30,7 +30,7 @@ where
 
     pub(super) fn set_fallback(&mut self, endpoint: Endpoint<S, B>) {
         self.replace_endpoint("/", endpoint.clone());
-        self.replace_endpoint(&format!("/*{FALLBACK_PARAM}"), endpoint);
+        self.replace_endpoint(FALLBACK_PARAM_PATH, endpoint);
     }
 }
 
@@ -139,10 +139,26 @@ where
                 .route_id_to_path
                 .get(&id)
                 .expect("no path for route id. This is a bug in axum. Please file an issue");
-            match route {
-                Endpoint::MethodRouter(method_router) => self.route(path, method_router)?,
-                Endpoint::Route(route) => self.route_service(path, route)?,
-            };
+
+            if IS_FALLBACK && (&**path == "/" || &**path == FALLBACK_PARAM_PATH) {
+                // when merging two routers it doesn't matter if you do `a.merge(b)` or
+                // `b.merge(a)`. This must also be true for fallbacks.
+                //
+                // However all fallback routers will have routes for `/` and `/*` so when merging
+                // we have to ignore the top level fallbacks on one side otherwise we get
+                // conflicts.
+                //
+                // `Router::merge` makes sure that when merging fallbacks `other` always has the
+                // fallback we want to keep. It panics if both routers have a custom fallback. Thus
+                // it is always okay to ignore one fallback and `Router::merge` also makes sure the
+                // one we can ignore is that of `self`.
+                self.replace_endpoint(path, route);
+            } else {
+                match route {
+                    Endpoint::MethodRouter(method_router) => self.route(path, method_router)?,
+                    Endpoint::Route(route) => self.route_service(path, route)?,
+                }
+            }
         }
 
         Ok(())
diff --git a/axum/src/routing/tests/fallback.rs b/axum/src/routing/tests/fallback.rs
@@ -241,3 +241,122 @@ async fn doesnt_panic_if_used_with_nested_router() {
     let res = client.get("/foobar").send().await;
     assert_eq!(res.status(), StatusCode::OK);
 }
+
+#[crate::test]
+async fn issue_2072() {
+    let nested_routes = Router::new().fallback(inner_fallback);
+
+    let app = Router::new()
+        .nest("/nested", nested_routes)
+        .merge(Router::new());
+
+    let client = TestClient::new(app);
+
+    let res = client.get("/nested/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "inner");
+
+    let res = client.get("/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "");
+}
+
+#[crate::test]
+async fn issue_2072_outer_fallback_before_merge() {
+    let nested_routes = Router::new().fallback(inner_fallback);
+
+    let app = Router::new()
+        .nest("/nested", nested_routes)
+        .fallback(outer_fallback)
+        .merge(Router::new());
+
+    let client = TestClient::new(app);
+
+    let res = client.get("/nested/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "inner");
+
+    let res = client.get("/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "outer");
+}
+
+#[crate::test]
+async fn issue_2072_outer_fallback_after_merge() {
+    let nested_routes = Router::new().fallback(inner_fallback);
+
+    let app = Router::new()
+        .nest("/nested", nested_routes)
+        .merge(Router::new())
+        .fallback(outer_fallback);
+
+    let client = TestClient::new(app);
+
+    let res = client.get("/nested/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "inner");
+
+    let res = client.get("/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "outer");
+}
+
+#[crate::test]
+async fn merge_router_with_fallback_into_nested_router_with_fallback() {
+    let nested_routes = Router::new().fallback(inner_fallback);
+
+    let app = Router::new()
+        .nest("/nested", nested_routes)
+        .merge(Router::new().fallback(outer_fallback));
+
+    let client = TestClient::new(app);
+
+    let res = client.get("/nested/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "inner");
+
+    let res = client.get("/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "outer");
+}
+
+#[crate::test]
+async fn merging_nested_router_with_fallback_into_router_with_fallback() {
+    let nested_routes = Router::new().fallback(inner_fallback);
+
+    let app = Router::new()
+        .fallback(outer_fallback)
+        .merge(Router::new().nest("/nested", nested_routes));
+
+    let client = TestClient::new(app);
+
+    let res = client.get("/nested/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "inner");
+
+    let res = client.get("/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "outer");
+}
+
+#[crate::test]
+async fn merge_empty_into_router_with_fallback() {
+    let app = Router::new().fallback(outer_fallback).merge(Router::new());
+
+    let client = TestClient::new(app);
+
+    let res = client.get("/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "outer");
+}
+
+#[crate::test]
+async fn merge_router_with_fallback_into_empty() {
+    let app = Router::new().merge(Router::new().fallback(outer_fallback));
+
+    let client = TestClient::new(app);
+
+    let res = client.get("/does-not-exist").send().await;
+    assert_eq!(res.status(), StatusCode::NOT_FOUND);
+    assert_eq!(res.text().await, "outer");
+}
