Add custom asynchronous asset handlers

packages/desktop/src/desktop_context.rs

@@ -1,11 +1,18 @@
 use std::cell::RefCell;
+use std::future::Future;
+use std::path::Path;
 use std::rc::Rc;
 use std::rc::Weak;
 
 use crate::create_new_window;
 use crate::events::IpcMessage;
+use crate::protocol::AssetFuture;
+use crate::protocol::AssetHandlerId;
+use crate::protocol::AssetHandlerRegistry;
+use crate::protocol::AssetResponse;
 use crate::query::QueryEngine;
 use crate::shortcut::{HotKey, ShortcutId, ShortcutRegistry, ShortcutRegistryError};
+use crate::AssetHandler;
 use crate::Config;
 use crate::WebviewHandler;
 use dioxus_core::ScopeState;
@@ -64,6 +71,8 @@ pub struct DesktopService {
 
     pub(crate) shortcut_manager: ShortcutRegistry,
 
+    pub(crate) asset_handlers: AssetHandlerRegistry,
+
     #[cfg(target_os = "ios")]
     pub(crate) views: Rc<RefCell<Vec<*mut objc::runtime::Object>>>,
 }
@@ -88,6 +97,7 @@ impl DesktopService {
         webviews: WebviewQueue,
         event_handlers: WindowEventHandlers,
         shortcut_manager: ShortcutRegistry,
+        asset_handlers: AssetHandlerRegistry,
     ) -> Self {
         Self {
             webview: Rc::new(webview),
@@ -97,6 +107,7 @@ impl DesktopService {
             pending_windows: webviews,
             event_handlers,
             shortcut_manager,
+            asset_handlers,
             #[cfg(target_os = "ios")]
             views: Default::default(),
         }
@@ -247,6 +258,18 @@ impl DesktopService {
         self.shortcut_manager.remove_all()
     }
 
+    /// Provide a callback to handle asset loading yourself.
+    ///
+    /// See [`use_asset_handle`](crate::use_asset_handle) for a convenient hook.
+    pub async fn register_asset_handler<F: AssetFuture>(&self, f: impl AssetHandler<F>) {
+        self.asset_handlers.register_handler(f).await;
+    }
+
+    /// Removes an asset handler by its identifier.
+    pub async fn remove_asset_handler(&self, id: AssetHandlerId) {
+        self.asset_handlers.remove_handler(id).await;
+    }
+
     /// Push an objc view to the window
     #[cfg(target_os = "ios")]
     pub fn push_view(&self, view: objc_id::ShareId<objc::runtime::Object>) {

packages/desktop/src/lib.rs

@@ -32,6 +32,7 @@ use dioxus_html::{native_bind::NativeFileEngine, FormData, HtmlEvent};
 use element::DesktopElement;
 use eval::init_eval;
 use futures_util::{pin_mut, FutureExt};
+pub use protocol::{use_asset_handler, AssetFuture, AssetHandler, AssetHandlerId, AssetResponse};
 use shortcut::ShortcutRegistry;
 pub use shortcut::{use_global_shortcut, ShortcutHandle, ShortcutId, ShortcutRegistryError};
 use std::cell::Cell;
@@ -393,7 +394,8 @@ fn create_new_window(
     event_handlers: &WindowEventHandlers,
     shortcut_manager: ShortcutRegistry,
 ) -> WebviewHandler {
-    let (webview, web_context) = webview::build(&mut cfg, event_loop, proxy.clone());
+    let (webview, web_context, asset_handlers) =
+        webview::build(&mut cfg, event_loop, proxy.clone());
     let desktop_context = Rc::from(DesktopService::new(
         webview,
         proxy.clone(),
@@ -401,6 +403,7 @@ fn create_new_window(
         queue.clone(),
         event_handlers.clone(),
         shortcut_manager,
+        asset_handlers,
     ));
 
     let cx = dom.base_scope();

packages/desktop/src/protocol.rs

@@ -1,13 +1,25 @@
+use dioxus_core::ScopeState;
 use dioxus_interpreter_js::{COMMON_JS, INTERPRETER_JS};
 use std::{
     borrow::Cow,
+    collections::HashMap,
+    future::Future,
     path::{Path, PathBuf},
+    pin::Pin,
+    rc::Rc,
+    sync::Arc,
+};
+use tokio::{
+    runtime::Handle,
+    sync::{OnceCell, RwLock},
 };
 use wry::{
     http::{status::StatusCode, Request, Response},
     Result,
 };
 
+use crate::{use_window, DesktopContext};
+
 fn module_loader(root_name: &str) -> String {
     let js = INTERPRETER_JS.replace(
         "/*POST_HANDLE_EDITS*/",
@@ -51,12 +63,132 @@ fn module_loader(root_name: &str) -> String {
     )
 }
 
-pub(super) fn desktop_handler(
+/// An arbitrary asset is an HTTP response containing a binary body.
+pub type AssetResponse = Response<Cow<'static, [u8]>>;
+
+/// A future that returns an [`AssetResponse`]. This future may be spawned in a new thread,
+/// so it must be [`Send`], [`Sync`], and `'static`.
+pub trait AssetFuture: Future<Output = Option<AssetResponse>> + Send + Sync + 'static {}
+impl<T: Future<Output = Option<AssetResponse>> + Send + Sync + 'static> AssetFuture for T {}
+
+/// A handler that takes an asset [`Path`] and returns a future that loads the path.
+/// This handler is stashed indefinitely in a context object, so it must be `'static`.
+pub trait AssetHandler<F: AssetFuture>: Fn(&Path) -> F + Send + Sync + 'static {}
+impl<F: AssetFuture, T: Fn(&Path) -> F + Send + Sync + 'static> AssetHandler<F> for T {}
+
+/// An identifier for a registered asset handler, returned by [`AssetHandlerRegistry::register_handler`].
+#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
+pub struct AssetHandlerId(usize);
+
+struct AssetHandlerRegistryInner {
+    handlers: HashMap<
+        AssetHandlerId,
+        Box<dyn Fn(&Path) -> Pin<Box<dyn AssetFuture>> + Send + Sync + 'static>,
+    >,
+    counter: AssetHandlerId,
+}
+
+#[derive(Clone)]
+pub struct AssetHandlerRegistry(Arc<RwLock<AssetHandlerRegistryInner>>);
+
+impl AssetHandlerRegistry {
+    pub fn new() -> Self {
+        AssetHandlerRegistry(Arc::new(RwLock::new(AssetHandlerRegistryInner {
+            handlers: HashMap::new(),
+            counter: AssetHandlerId(0),
+        })))
+    }
+
+    pub async fn register_handler<F: AssetFuture>(
+        &self,
+        f: impl AssetHandler<F>,
+    ) -> AssetHandlerId {
+        let mut registry = self.0.write().await;
+        let id = registry.counter;
+        registry
+            .handlers
+            .insert(id, Box::new(move |path| Box::pin(f(path))));
+        registry.counter.0 += 1;
+        id
+    }
+
+    pub async fn remove_handler(&self, id: AssetHandlerId) -> Option<()> {
+        let mut registry = self.0.write().await;
+        registry.handlers.remove(&id).map(|_| ())
+    }
+
+    pub async fn try_handlers(&self, path: &Path) -> Option<AssetResponse> {
+        let registry = self.0.read().await;
+        for handler in registry.handlers.values() {
+            if let Some(response) = handler(path).await {
+                return Some(response);
+            }
+        }
+        None
+    }
+}
+
+/// A handle to a registered asset handler.
+pub struct AssetHandlerHandle {
+    desktop: DesktopContext,
+    handler_id: Rc<OnceCell<AssetHandlerId>>,
+}
+
+impl AssetHandlerHandle {
+    /// Returns the [`AssetHandlerId`] for this handle.
+    ///
+    /// Because registering an ID is asynchronous, this may return `None` if the
+    /// registration has not completed yet.
+    pub fn handler_id(&self) -> Option<AssetHandlerId> {
+        self.handler_id.get().copied()
+    }
+}
+
+impl Drop for AssetHandlerHandle {
+    fn drop(&mut self) {
+        let cell = Rc::clone(&self.handler_id);
+        let desktop = Rc::clone(&self.desktop);
+        tokio::task::block_in_place(move || {
+            Handle::current().block_on(async move {
+                if let Some(id) = cell.get() {
+                    desktop.asset_handlers.remove_handler(*id).await;
+                }
+            })
+        });
+    }
+}
+
+/// Provide a callback to handle asset loading yourself.
+///
+/// The callback takes a path as requested by the web view, and it should return `Some(response)`
+/// if you want to load the asset, and `None` if you want to fallback on the default behavior.
+pub fn use_asset_handler<F: AssetFuture>(
+    cx: &ScopeState,
+    handler: impl AssetHandler<F>,
+) -> &AssetHandlerHandle {
+    let desktop = Rc::clone(&use_window(cx));
+    cx.use_hook(|| {
+        let handler_id = Rc::new(OnceCell::new());
+        let handler_id_ref = Rc::clone(&handler_id);
+        let desktop_ref = Rc::clone(&desktop);
+        cx.push_future(async move {
+            let id = desktop.asset_handlers.register_handler(handler).await;
+            handler_id.set(id).unwrap();
+        });
+        AssetHandlerHandle {
+            desktop: desktop_ref,
+            handler_id: handler_id_ref,
+        }
+    })
+}
+
+pub(super) async fn desktop_handler(
     request: &Request<Vec<u8>>,
     custom_head: Option<String>,
     custom_index: Option<String>,
     root_name: &str,
-) -> Result<Response<Cow<'static, [u8]>>> {
+    asset_handlers: &AssetHandlerRegistry,
+) -> Result<AssetResponse> {
     // If the request is for the root, we'll serve the index.html file.
     if request.uri().path() == "/" {
         // If a custom index is provided, just defer to that, expecting the user to know what they're doing.
@@ -96,6 +228,13 @@ pub(super) fn desktop_handler(
         .expect("expected URL to be UTF-8 encoded");
     let path = PathBuf::from(&*decoded);
 
+    // If the user provided a custom asset handler, then call it and return the response
+    // if the request was handled.
+    if let Some(response) = asset_handlers.try_handlers(&path).await {
+        return Ok(response);
+    }
+
+    // Else, try to serve a file from the filesystem.
     // If the path is relative, we'll try to serve it from the assets directory.
     let mut asset = get_asset_root()
         .unwrap_or_else(|| Path::new(".").to_path_buf())

packages/desktop/src/webview.rs

@@ -1,5 +1,7 @@
+use std::sync::Arc;
+
 use crate::desktop_context::EventData;
-use crate::protocol;
+use crate::protocol::{self, AssetHandlerRegistry};
 use crate::{desktop_context::UserWindowEvent, Config};
 use tao::event_loop::{EventLoopProxy, EventLoopWindowTarget};
 pub use wry;
@@ -12,7 +14,7 @@ pub fn build(
     cfg: &mut Config,
     event_loop: &EventLoopWindowTarget<UserWindowEvent>,
     proxy: EventLoopProxy<UserWindowEvent>,
-) -> (WebView, WebContext) {
+) -> (WebView, WebContext, AssetHandlerRegistry) {
     let builder = cfg.window.clone();
     let window = builder.with_visible(false).build(event_loop).unwrap();
     let file_handler = cfg.file_drop_handler.take();
@@ -33,6 +35,8 @@ pub fn build(
     }
 
     let mut web_context = WebContext::new(cfg.data_dir.clone());
+    let asset_handlers = AssetHandlerRegistry::new();
+    let asset_handlers_ref = asset_handlers.clone();
 
     let mut webview = WebViewBuilder::new(window)
         .unwrap()
@@ -45,24 +49,30 @@ pub fn build(
                 _ = proxy.send_event(UserWindowEvent(EventData::Ipc(message), window.id()));
             }
         })
-        .with_custom_protocol(
-            String::from("dioxus"),
-            move |r| match protocol::desktop_handler(
-                &r,
-                custom_head.clone(),
-                index_file.clone(),
-                &root_name,
-            ) {
-                Ok(response) => response,
-                Err(err) => {
+        .with_asynchronous_custom_protocol(String::from("dioxus"), move |request, responder| {
+            let custom_head = custom_head.clone();
+            let index_file = index_file.clone();
+            let root_name = root_name.clone();
+            let asset_handlers_ref = asset_handlers_ref.clone();
+            tokio::spawn(async move {
+                let response_res = protocol::desktop_handler(
+                    &request,
+                    custom_head.clone(),
+                    index_file.clone(),
+                    &root_name,
+                    &asset_handlers_ref,
+                )
+                .await;
+                let response = response_res.unwrap_or_else(|err| {
                     tracing::error!("Error: {}", err);
                     Response::builder()
                         .status(500)
                         .body(err.to_string().into_bytes().into())
                         .unwrap()
-                }
-            },
-        )
+                });
+                responder.respond(response);
+            });
+        })
         .with_file_drop_handler(move |window, evet| {
             file_handler
                 .as_ref()
@@ -119,5 +129,5 @@ pub fn build(
         webview = webview.with_devtools(true);
     }
 
-    (webview.build().unwrap(), web_context)
+    (webview.build().unwrap(), web_context, asset_handlers)
 }