refactor: turn asset cache into tower service

Author: dpc

Cargo.lock

@@ -89,9 +89,11 @@ snafu = { version = "0.8.5", features = ["rust_1_81"] }
 time = "0.3.36"
 tokio = { version = "1.42.0", features = ["macros"] }
 tokio-stream = "*"
+tower = "*"
 tower-cookies = "*"
 tower-http = "0.5"
 tower-sessions = "0.14.0"
+tower-service = "0.3.3"
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 z32 = "1.1.1"

Cargo.toml

@@ -31,7 +31,9 @@ rostra-util = { workspace = true }
 rostra-util-error = { workspace = true }
 tokio = { workspace = true }
 tokio-stream = { workspace = true, features = ["fs"] }
+tower = { workspace = true }
 tower-cookies = { workspace = true }
+tower-service = { workspace = true }
 tower-http = { workspace = true, features = ["cors", "compression-br", "fs"] }
 tower-sessions = { workspace = true }
 tracing = { workspace = true }

crates/rostra-web-ui/Cargo.toml

@@ -7,7 +7,6 @@ use std::path::{self, PathBuf};
 use std::string::String;
 use std::sync::LazyLock;
 
-use axum::extract::Path;
 use bytes::Bytes;
 use futures::stream::{BoxStream, StreamExt};
 use rostra_util_error::WhateverResult;
@@ -17,39 +16,14 @@ use tracing::{debug, info};
 
 use crate::LOG_TARGET;
 
-const HASH_SPLIT_CHAR: char = '.';
-
-/// Maps static asset filenames to their compressed bytes and content type. This
+/// Pre-loaded and pre-compressed static assets. This
 /// is used to serve static assets from the build directory without reading from
 /// disk, as the cache stays in RAM for the life of the server.
-///
-/// This type should be accessed via the `cache` property in `AppState`.
 #[derive(Debug)]
-pub struct AssetCache(HashMap<String, StaticAsset>);
-
-impl AssetCache {
-    /// Attempts to return a static asset from the cache from a cache key. If
-    /// the asset is not found, `None` is returned.
-    pub fn get(&self, key: &str) -> Option<&StaticAsset> {
-        self.0.get(key)
-    }
-
-    /// Helper method to get a static asset from an extracted request path.
-    pub fn get_from_path(&self, path: &Path<String>) -> Option<&StaticAsset> {
-        let key = Self::get_cache_key(path);
-        self.get(&key)
-    }
-
-    fn get_cache_key(path: &str) -> String {
-        let mut parts = path.split(['.', HASH_SPLIT_CHAR]);
+pub struct StaticAssets(HashMap<String, StaticAsset>);
 
-        let basename = parts.next().unwrap_or_default();
-        let ext = parts.next_back().unwrap_or_default();
-
-        format!("{}.{}", basename, ext)
-    }
-
-    pub async fn load_files(root_dir: &path::Path) -> WhateverResult<Self> {
+impl StaticAssets {
+    pub async fn load(root_dir: &path::Path) -> WhateverResult<Self> {
         info!(target: LOG_TARGET, dir=%root_dir.display(), "Loading assets");
         let mut cache = HashMap::default();
 
@@ -111,7 +85,7 @@ impl AssetCache {
                         },
                     )))
                 })
-                .buffered(16)
+                .buffered(32)
                 .filter_map(
                     |res_opt: WhateverResult<std::option::Option<(String, StaticAsset)>>| {
                         ready(res_opt.transpose())
@@ -122,7 +96,7 @@ impl AssetCache {
 
         for asset_res in assets {
             let (filename, asset) = asset_res?;
-            cache.insert(Self::get_cache_key(&filename), asset);
+            cache.insert(filename, asset);
         }
 
         for (key, asset) in &cache {
@@ -132,6 +106,12 @@ impl AssetCache {
 
         Ok(Self(cache))
     }
+
+    /// Attempts to return a static asset from the cache from a cache key. If
+    /// the asset is not found, `None` is returned.
+    pub fn get(&self, key: &str) -> Option<&StaticAsset> {
+        self.0.get(key)
+    }
 }
 
 /// Represents a single static asset from the build directory. Assets are
@@ -197,8 +177,6 @@ fn calculate_etag(data: &[u8]) -> String {
     format!("\"{}\"", hasher.finish())
 }
 
-// async fn read_dir_stream(dir: impl AsRef<path::Path>) -> impl Stream<Item =
-// io::Result<PathBuf>> {
 fn read_dir_stream(dir: PathBuf) -> BoxStream<'static, io::Result<PathBuf>> {
     async_stream::try_stream! {
         let entries = ReadDirStream::new(

crates/rostra-web-ui/src/asset_cache.rs

@@ -0,0 +1,94 @@
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use axum::body::Body;
+use axum::http::header::{ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_TYPE};
+use axum::http::{HeaderMap, HeaderValue, Request, Response, StatusCode};
+use axum::response::IntoResponse;
+use futures::future::BoxFuture;
+use tower_service::Service;
+
+use crate::asset_cache::StaticAssets;
+
+#[derive(Clone)]
+pub struct StaticAssetService {
+    assets: Arc<StaticAssets>,
+}
+
+impl StaticAssetService {
+    pub fn new(assets: Arc<StaticAssets>) -> Self {
+        Self { assets }
+    }
+
+    fn handle_request(&self, req: Request<Body>) -> Response<Body> {
+        let path = req.uri().path().trim_start_matches('/');
+        let Some(asset) = self.assets.get(path) else {
+            dbg!("NOT FOUND", path);
+            return (StatusCode::NOT_FOUND, Body::empty()).into_response();
+        };
+
+        let req_headers = req.headers();
+        let mut resp_headers = HeaderMap::new();
+
+        // Set content type
+        resp_headers.insert(
+            CONTENT_TYPE,
+            HeaderValue::from_static(asset.content_type().unwrap_or("application/octet-stream")),
+        );
+
+        // Handle ETag and conditional request
+        let etag = asset.etag.clone();
+        if let Some(response) = crate::handle_etag(req_headers, &etag, &mut resp_headers) {
+            return response;
+        }
+
+        let accepts_brotli = req_headers
+            .get_all(ACCEPT_ENCODING)
+            .into_iter()
+            .any(|encodings| {
+                let Ok(str) = encodings.to_str() else {
+                    return false;
+                };
+
+                str.split(',').any(|s| s.trim() == "br")
+            });
+
+        let content = match (accepts_brotli, asset.compressed.as_ref()) {
+            (true, Some(compressed)) => {
+                resp_headers.insert(CONTENT_ENCODING, HeaderValue::from_static("br"));
+                compressed.clone()
+            }
+            _ => asset.raw.clone(),
+        };
+
+        (resp_headers, content).into_response()
+    }
+}
+
+impl<B> Service<Request<B>> for StaticAssetService
+where
+    B: Send + 'static,
+{
+    type Response = Response<Body>;
+    type Error = std::convert::Infallible;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), std::convert::Infallible>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(
+        &mut self,
+        req: Request<B>,
+    ) -> BoxFuture<'static, Result<Response<Body>, std::convert::Infallible>> {
+        let service = self.clone();
+        let uri = req.uri().clone();
+
+        Box::pin(async move {
+            // Convert to a Request<Body> by extracting just the URI
+            let new_req = Request::builder().uri(uri).body(Body::empty()).unwrap();
+
+            Ok(service.handle_request(new_req))
+        })
+    }
+}

crates/rostra-web-ui/src/asset_service.rs

@@ -1,4 +1,5 @@
 pub(crate) mod asset_cache;
+mod asset_service;
 mod error;
 mod fragment;
 pub mod html_utils;
@@ -14,10 +15,9 @@ use std::sync::Arc;
 use std::time::Duration;
 use std::{io, result};
 
-use asset_cache::AssetCache;
+use asset_cache::StaticAssets;
 use axum::http::header::{ACCEPT, CONTENT_TYPE};
 use axum::http::{HeaderName, HeaderValue, Method};
-use axum::routing::get;
 use axum::{Router, middleware};
 use error::{IdMismatchSnafu, UnlockError, UnlockResult};
 use listenfd::ListenFd;
@@ -27,7 +27,7 @@ use rostra_client::{ClientHandle, ClientRefError};
 use rostra_core::id::{RostraId, RostraIdSecretKey};
 use rostra_util::is_rostra_dev_mode_set;
 use rostra_util_error::WhateverResult;
-use routes::{cache_control, get_static_asset};
+use routes::cache_control;
 use snafu::{ResultExt as _, Snafu, Whatever, ensure};
 use tokio::net::{TcpListener, TcpSocket};
 use tokio::signal;
@@ -136,7 +136,7 @@ pub type UiStateClientResult<T> = result::Result<T, UiStateClientError>;
 
 pub struct UiState {
     clients: MultiClient,
-    assets: Option<Arc<AssetCache>>,
+    assets: Option<Arc<StaticAssets>>,
     default_profile: Option<RostraId>,
 }
 
@@ -172,7 +172,7 @@ pub struct Server {
     listener: TcpListener,
 
     state: SharedState,
-    assets: Option<Arc<AssetCache>>,
+    assets: Option<Arc<StaticAssets>>,
     opts: Opts,
 }
 
@@ -219,7 +219,7 @@ impl Server {
             None
         } else {
             Some(Arc::new(
-                AssetCache::load_files(&opts.assets_dir)
+                StaticAssets::load(&opts.assets_dir)
                     .await
                     .context(AssetsLoadSnafu)?,
             ))
@@ -275,22 +275,15 @@ impl Server {
         );
         let mut router = Router::new().merge(routes::route_handler(self.state.clone()));
 
-        match self.assets {
-            Some(_assets) => {
-                router = router.nest("/assets", {
-                    let state = self.state.clone();
-                    Router::new()
-                        .route("/{*file}", get(get_static_asset))
-                        .with_state(state)
-                });
+        router = match self.assets.clone() {
+            Some(assets) => {
+                router.nest_service("/assets", asset_service::StaticAssetService::new(assets))
             }
-            _ => {
-                router = router.nest_service(
-                    "/assets",
-                    ServeDir::new(format!("{}/assets", env!("CARGO_MANIFEST_DIR"))),
-                );
-            }
-        }
+            _ => router.nest_service(
+                "/assets",
+                ServeDir::new(format!("{}/assets", env!("CARGO_MANIFEST_DIR"))),
+            ),
+        };
 
         let session_store = MemoryStore::default();
         let session_layer = SessionManagerLayer::new(session_store)