1 files changed, 34 insertions, 12 deletions

diff --git a/src/ui.rs b/src/ui.rs
index bba01cc..18dd049 100644
--- a/src/ui.rs
+++ b/src/ui.rs
@@ -1,38 +1,60 @@
 use axum::{
     extract::Path,
     http::{header, StatusCode},
-    response::IntoResponse,
+    response::{IntoResponse, Response},
     routing::get,
     Router,
 };
+use mime_guess::Mime;
+use rust_embed::EmbeddedFile;
 
 #[derive(rust_embed::Embed)]
 #[folder = "hi-ui/build"]
 struct Assets;
 
+pub fn router<S>() -> Router<S>
+where
+    S: Clone + Send + Sync + 'static,
+{
+    Router::new()
+        .route("/*path", get(asset))
+        .route("/", get(root))
+}
+
+async fn asset(Path(path): Path<String>) -> Result<Asset, NotFound<String>> {
+    let mime = mime_guess::from_path(&path).first_or_octet_stream();
+
+    Assets::get(&path)
+        .map(|file| Asset(mime, file))
+        .ok_or(NotFound(format!("not found: {path}")))
+}
+
 async fn root() -> impl IntoResponse {
     asset(Path(String::from("index.html"))).await
 }
 
-async fn asset(Path(path): Path<String>) -> impl IntoResponse {
-    let mime = mime_guess::from_path(&path).first_or_octet_stream();
+struct Asset(Mime, EmbeddedFile);
 
-    match Assets::get(&path) {
-        Some(file) => (
+impl IntoResponse for Asset {
+    fn into_response(self) -> Response {
+        let Self(mime, file) = self;
+        (
             StatusCode::OK,
             [(header::CONTENT_TYPE, mime.as_ref())],
             file.data,
         )
-            .into_response(),
-        None => (StatusCode::NOT_FOUND, "").into_response(),
+            .into_response()
     }
 }
 
-pub fn router<S>() -> Router<S>
+struct NotFound<E>(pub E);
+
+impl<E> IntoResponse for NotFound<E>
 where
-    S: Clone + Send + Sync + 'static,
+    E: IntoResponse,
 {
-    Router::new()
-        .route("/*path", get(asset))
-        .route("/", get(root))
+    fn into_response(self) -> Response {
+        let Self(response) = self;
+        (StatusCode::NOT_FOUND, response).into_response()
+    }
 }