From dc39131021a32b500575dc16cc5c49ba1ae016fb Mon Sep 17 00:00:00 2001
From: Bugen Zhao <i@bugenzhao.com>
Date: Thu, 4 Apr 2024 14:27:45 +0800
Subject: [PATCH] refactor(meta): use `axum-embed` to embed dashboard assets
 (#16131)

---
 Cargo.lock                      | 15 +++++++++++++++
 src/meta/dashboard/Cargo.toml   |  1 +
 src/meta/dashboard/src/embed.rs | 33 +++------------------------------
 3 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 69bb6d213bf94..81bc9834b48e2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1459,6 +1459,20 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "axum-embed"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "077959a7f8cf438676af90b483304528eb7e16eadadb7f44e9ada4f9dceb9e62"
+dependencies = [
+ "axum-core 0.4.3",
+ "chrono",
+ "http 1.0.0",
+ "mime_guess",
+ "rust-embed",
+ "tower-service",
+]
+
 [[package]]
 name = "backon"
 version = "0.4.1"
@@ -10105,6 +10119,7 @@ version = "1.7.0-alpha"
 dependencies = [
  "anyhow",
  "axum 0.7.4",
+ "axum-embed",
  "bytes",
  "cargo-emit",
  "dircpy",
diff --git a/src/meta/dashboard/Cargo.toml b/src/meta/dashboard/Cargo.toml
index 121485c904f06..e848c8015cbb1 100644
--- a/src/meta/dashboard/Cargo.toml
+++ b/src/meta/dashboard/Cargo.toml
@@ -11,6 +11,7 @@ repository = { workspace = true }
 [dependencies]
 anyhow = "1"
 axum = "=0.7.4" # TODO: 0.7.5+ does not work with current toolchain
+axum-embed = "0.1"
 bytes = "1"
 hyper = "1"
 mime_guess = "2"
diff --git a/src/meta/dashboard/src/embed.rs b/src/meta/dashboard/src/embed.rs
index e38757bd3bac9..2d60ce1a89704 100644
--- a/src/meta/dashboard/src/embed.rs
+++ b/src/meta/dashboard/src/embed.rs
@@ -12,42 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use axum::http::{header, StatusCode, Uri};
-use axum::response::{IntoResponse as _, Response};
 use axum::Router;
+use axum_embed::ServeEmbed;
 use rust_embed::RustEmbed;
 
-#[derive(RustEmbed)]
+#[derive(RustEmbed, Clone)]
 #[folder = "$OUT_DIR/assets"]
 struct Assets;
 
-// TODO: switch to `axum-embed` for better robustness after bumping to axum 0.7.
-async fn static_handler(uri: Uri) -> Response {
-    let path = {
-        if uri.path().ends_with('/') {
-            // Append `index.html` to directory paths.
-            format!("{}index.html", uri.path())
-        } else {
-            uri.path().to_owned()
-        }
-    };
-    let path = path.trim_start_matches('/');
-
-    match Assets::get(path) {
-        Some(file) => {
-            let mime = file.metadata.mimetype();
-
-            let mut res = file.data.into_response();
-            res.headers_mut()
-                .insert(header::CONTENT_TYPE, mime.parse().unwrap());
-            res
-        }
-
-        None => (StatusCode::NOT_FOUND, "Not Found").into_response(),
-    }
-}
-
 /// Router for embedded assets.
 pub(crate) fn router() -> Router {
-    Router::new().fallback(static_handler)
+    Router::new().nest_service("/", ServeEmbed::<Assets>::new())
 }