4 files changed, 56 insertions, 18 deletions

diff --git a/src/channel/routes.rs b/src/channel/routes.rs
index a92615f..e97c447 100644
--- a/src/channel/routes.rs
+++ b/src/channel/routes.rs
@@ -7,7 +7,13 @@ use axum::{
 };
 
 use super::{app, Channel, Id};
-use crate::{app::App, clock::RequestedAt, error::Internal, login::Login, message::app::SendError};
+use crate::{
+    app::App,
+    clock::RequestedAt,
+    error::{Internal, NotFound},
+    login::Login,
+    message::app::SendError,
+};
 
 #[cfg(test)]
 mod test;
@@ -81,9 +87,7 @@ impl IntoResponse for SendErrorResponse {
     fn into_response(self) -> Response {
         let Self(error) = self;
         match error {
-            not_found @ SendError::ChannelNotFound(_) => {
-                (StatusCode::NOT_FOUND, not_found.to_string()).into_response()
-            }
+            not_found @ SendError::ChannelNotFound(_) => NotFound(not_found).into_response(),
             other => Internal::from(other).into_response(),
         }
     }
diff --git a/src/error.rs b/src/error.rs
index 8792a1d..85573d4 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -69,3 +69,15 @@ impl IntoResponse for Unauthorized {
         (StatusCode::UNAUTHORIZED, "unauthorized").into_response()
     }
 }
+
+pub struct NotFound<E>(pub E);
+
+impl<E> IntoResponse for NotFound<E>
+where
+    E: std::error::Error,
+{
+    fn into_response(self) -> Response {
+        let Self(response) = self;
+        (StatusCode::NOT_FOUND, response.to_string()).into_response()
+    }
+}
diff --git a/src/message/routes.rs b/src/message/routes.rs
index 29fe3d7..e21c674 100644
--- a/src/message/routes.rs
+++ b/src/message/routes.rs
@@ -9,7 +9,7 @@ use axum::{
 use crate::{
     app::App,
     clock::RequestedAt,
-    error::Internal,
+    error::{Internal, NotFound},
     login::Login,
     message::{self, app::DeleteError},
 };
@@ -38,7 +38,7 @@ impl IntoResponse for ErrorResponse {
         let Self(error) = self;
         match error {
             not_found @ (DeleteError::ChannelNotFound(_) | DeleteError::NotFound(_)) => {
-                (StatusCode::NOT_FOUND, not_found.to_string()).into_response()
+                NotFound(not_found).into_response()
             }
             other => Internal::from(other).into_response(),
         }
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()
+    }
 }