From 5762763aacc859d51e10b8cf62cafdfd855fd278 Mon Sep 17 00:00:00 2001
From: masnagam <masnagam@gmail.com>
Date: Mon, 21 Oct 2024 18:43:04 +0900
Subject: [PATCH] Fix regression due to
 a985d8c1340f80ab69b2b0e5de799df98d567732 (#1155)

---
 utoipa-swagger-ui/Cargo.toml  |  2 ++
 utoipa-swagger-ui/src/axum.rs | 65 ++++++++++++++++++++++++++++++++---
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/utoipa-swagger-ui/Cargo.toml b/utoipa-swagger-ui/Cargo.toml
index f872381f..2fe9b9a7 100644
--- a/utoipa-swagger-ui/Cargo.toml
+++ b/utoipa-swagger-ui/Cargo.toml
@@ -34,7 +34,9 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0" }
 
 [dev-dependencies]
+axum-test = "16.2.0"
 similar = "2.5"
+tokio = { version = "1", features = ["macros"] }
 utoipa-swagger-ui = { path = ".", features = ["actix-web", "axum", "rocket"] }
 
 [package.metadata.docs.rs]
diff --git a/utoipa-swagger-ui/src/axum.rs b/utoipa-swagger-ui/src/axum.rs
index 22a31793..0c2c2aa8 100644
--- a/utoipa-swagger-ui/src/axum.rs
+++ b/utoipa-swagger-ui/src/axum.rs
@@ -45,10 +45,27 @@ where
         let handler = routing::get(serve_swagger_ui).layer(Extension(Arc::new(config)));
         let path: &str = swagger_ui.path.as_ref();
 
-        router
-            .route(path, handler.clone())
-            .route(&format!("{}/", path), handler.clone())
-            .route(&format!("{}/*rest", path), handler)
+        if path == "/" {
+            router
+                .route(path, handler.clone())
+                .route(&format!("{}*rest", path), handler)
+        } else {
+            let path = if path.ends_with('/') {
+                &path[..path.len() - 1]
+            } else {
+                path
+            };
+            debug_assert!(!path.is_empty());
+
+            let slash_path = format!("{}/", path);
+            router
+                .route(
+                    path,
+                    routing::get(|| async move { axum::response::Redirect::to(&slash_path) }),
+                )
+                .route(&format!("{}/", path), handler.clone())
+                .route(&format!("{}/*rest", path), handler)
+        }
     }
 }
 
@@ -96,3 +113,43 @@ async fn serve_swagger_ui(
         Err(error) => (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use axum_test::TestServer;
+
+    #[tokio::test]
+    async fn mount_onto_root() {
+        let app = Router::<()>::from(SwaggerUi::new("/"));
+        let server = TestServer::new(app).unwrap();
+        let response = server.get("/").await;
+        response.assert_status_ok();
+        let response = server.get("/swagger-ui.css").await;
+        response.assert_status_ok();
+    }
+
+    #[tokio::test]
+    async fn mount_onto_path_ends_with_slash() {
+        let app = Router::<()>::from(SwaggerUi::new("/swagger-ui/"));
+        let server = TestServer::new(app).unwrap();
+        let response = server.get("/swagger-ui").await;
+        response.assert_status_see_other();
+        let response = server.get("/swagger-ui/").await;
+        response.assert_status_ok();
+        let response = server.get("/swagger-ui/swagger-ui.css").await;
+        response.assert_status_ok();
+    }
+
+    #[tokio::test]
+    async fn mount_onto_path_not_end_with_slash() {
+        let app = Router::<()>::from(SwaggerUi::new("/swagger-ui"));
+        let server = TestServer::new(app).unwrap();
+        let response = server.get("/swagger-ui").await;
+        response.assert_status_see_other();
+        let response = server.get("/swagger-ui/").await;
+        response.assert_status_ok();
+        let response = server.get("/swagger-ui/swagger-ui.css").await;
+        response.assert_status_ok();
+    }
+}