dioxus_cloudflare/

handler.rs

//! Request handler that bridges Cloudflare Workers to the Dioxus Axum router.
//!
//! Two entry points are available:
//!
//! - [`handle`] — simple one-shot dispatch (no middleware)
//! - [`Handler`] — builder with before/after middleware hooks
//!
//! Both store the Worker `Env` and `Request` in thread-local context,
//! convert `worker::Request` → `http::Request`, dispatch through the
//! Dioxus Axum router, and stream the response back via `ReadableStream`.

use std::convert::Infallible;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use axum::body::Body;
use http_body::Body as HttpBody;
use tower::ServiceExt;
use worker::{Request, Response, ResponseBuilder};

use crate::context::{set_context, take_cookies};

#[cfg(feature = "ssr")]
use crate::ssr::{self, IndexHtml};

type BeforeHook = Box<dyn Fn(&Request) -> worker::Result<Option<Response>>>;
type AfterHook = Box<dyn Fn(&mut Response) -> worker::Result<()>>;
type WsHandler = Box<dyn Fn(Request) -> Pin<Box<dyn Future<Output = worker::Result<Response>>>>>;

struct WebSocketRoute {
    path_prefix: String,
    handler: WsHandler,
}

/// Internal state for SSR rendering.
#[cfg(feature = "ssr")]
struct SsrState {
    build_vdom: Box<dyn Fn() -> dioxus::dioxus_core::VirtualDom>,
    index: IndexHtml,
    streaming: bool,
}

/// Configurable request handler with before/after middleware hooks.
///
/// # Example
///
/// ```rust,ignore
/// use dioxus_cloudflare::Handler;
/// use worker::*;
///
/// #[event(fetch)]
/// async fn fetch(req: Request, env: Env, _ctx: Context) -> Result<Response> {
///     unsafe { __wasm_call_ctors(); }
///
///     Handler::new()
///         .before(|req| {
///             // Short-circuit OPTIONS requests for CORS preflight
///             if req.method() == worker::Method::Options {
///                 let mut resp = Response::empty()?;
///                 resp.headers_mut().set("Access-Control-Allow-Origin", "*")?;
///                 resp.headers_mut().set("Access-Control-Allow-Methods", "GET,POST,OPTIONS")?;
///                 return Ok(Some(resp));
///             }
///             Ok(None) // continue to Axum dispatch
///         })
///         .after(|resp| {
///             resp.headers_mut().set("Access-Control-Allow-Origin", "*")?;
///             Ok(())
///         })
///         .handle(req, env)
///         .await
/// }
/// ```
pub struct Handler {
    before: Vec<BeforeHook>,
    after: Vec<AfterHook>,
    websocket_routes: Vec<WebSocketRoute>,
    session_config: Option<crate::session::SessionConfig>,
    #[cfg(feature = "ssr")]
    ssr: Option<SsrState>,
}

impl Handler {
    /// Create a new handler with no middleware hooks.
    #[must_use]
    pub fn new() -> Self {
        Self {
            before: Vec::new(),
            after: Vec::new(),
            websocket_routes: Vec::new(),
            session_config: None,
            #[cfg(feature = "ssr")]
            ssr: None,
        }
    }

    /// Add a before-dispatch hook.
    ///
    /// Runs after context is set (so `cf::env()`, `cf::d1()`, etc. work),
    /// before Axum dispatch. Return `Ok(None)` to continue,
    /// `Ok(Some(resp))` to short-circuit with a custom response.
    ///
    /// Hooks run in the order they were added.
    #[must_use]
    pub fn before(
        mut self,
        hook: impl Fn(&Request) -> worker::Result<Option<Response>> + 'static,
    ) -> Self {
        self.before.push(Box::new(hook));
        self
    }

    /// Add an after-dispatch hook.
    ///
    /// Runs after cookies are applied, before returning the response.
    /// Use this to add headers, log, or modify the final response.
    ///
    /// After hooks also run on short-circuited responses from before hooks.
    /// Hooks run in the order they were added.
    #[must_use]
    pub fn after(
        mut self,
        hook: impl Fn(&mut Response) -> worker::Result<()> + 'static,
    ) -> Self {
        self.after.push(Box::new(hook));
        self
    }

    /// Configure session middleware.
    ///
    /// When enabled, `cf::session()` becomes available in server functions.
    /// Session data is persisted to the configured backend (KV or D1) and
    /// a session cookie is managed automatically.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use dioxus_cloudflare::prelude::*;
    ///
    /// Handler::new()
    ///     .session(SessionConfig::kv("SESSIONS"))
    ///     .handle(req, env)
    ///     .await
    /// ```
    #[must_use]
    pub fn session(mut self, config: crate::session::SessionConfig) -> Self {
        self.session_config = Some(config);
        self
    }

    /// Add a WebSocket upgrade route.
    ///
    /// When a request has the `Upgrade: websocket` header and its path starts
    /// with `path`, the handler is called instead of dispatching through Axum.
    /// The handler receives the owned `worker::Request` (to forward to a
    /// Durable Object stub, for example).
    ///
    /// WebSocket routes are checked after context is set and before hooks run,
    /// but before Axum dispatch.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// Handler::new()
    ///     .websocket("/ws", |req| async move {
    ///         let ns = cf::durable_object("WS_DO")?;
    ///         let id = ns.id_from_name("default").cf()?;
    ///         let stub = id.get_stub().cf()?;
    ///         Ok(stub.fetch_with_request(req).await.cf()?)
    ///     })
    ///     .handle(req, env)
    ///     .await
    /// ```
    #[must_use]
    pub fn websocket<F, Fut>(mut self, path: &str, handler: F) -> Self
    where
        F: Fn(Request) -> Fut + 'static,
        Fut: Future<Output = worker::Result<Response>> + 'static,
    {
        self.websocket_routes.push(WebSocketRoute {
            path_prefix: path.to_string(),
            handler: Box::new(move |req| Box::pin(handler(req))),
        });
        self
    }

    /// Enable SSR rendering for non-API requests.
    ///
    /// When the Axum router returns 404 and the request accepts `text/html`,
    /// the handler renders the given component to HTML and returns it.
    ///
    /// Uses a default HTML shell (`<!DOCTYPE html>` with `<div id="main">`).
    /// Call [`with_index_html`](Self::with_index_html) to provide your own
    /// shell (e.g., one that loads client WASM for SPA takeover).
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// Handler::new()
    ///     .with_ssr(App)
    ///     .handle(req, env)
    ///     .await
    /// ```
    #[cfg(feature = "ssr")]
    #[must_use]
    pub fn with_ssr(mut self, app: fn() -> dioxus::prelude::Element) -> Self {
        self.ssr = Some(SsrState {
            build_vdom: Box::new(move || dioxus::dioxus_core::VirtualDom::new(app)),
            index: IndexHtml::default_shell(),
            streaming: false,
        });
        self
    }

    /// Enable **streaming** SSR rendering for non-API requests.
    ///
    /// Like [`with_ssr`](Self::with_ssr), but sends the initial HTML with
    /// suspense fallbacks immediately, then streams resolved content
    /// out-of-order as each suspense boundary completes.
    ///
    /// If no suspense boundaries are pending after the initial render,
    /// this automatically falls back to a single-shot response with no
    /// overhead.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// Handler::new()
    ///     .with_streaming_ssr(App)
    ///     .with_index_html(include_str!("path/to/index.html"))?
    ///     .handle(req, env)
    ///     .await
    /// ```
    #[cfg(feature = "ssr")]
    #[must_use]
    pub fn with_streaming_ssr(mut self, app: fn() -> dioxus::prelude::Element) -> Self {
        self.ssr = Some(SsrState {
            build_vdom: Box::new(move || dioxus::dioxus_core::VirtualDom::new(app)),
            index: IndexHtml::default_shell(),
            streaming: true,
        });
        self
    }

    /// Provide a custom `index.html` shell for SSR responses.
    ///
    /// The HTML must contain an element with `id="main"` — rendered
    /// component output is inserted at that point.
    ///
    /// Must be called **after** [`with_ssr`](Self::with_ssr).
    ///
    /// # Errors
    ///
    /// Returns `Err` if:
    /// - SSR has not been configured (call `with_ssr` first)
    /// - The HTML does not contain `id="main"`
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// Handler::new()
    ///     .with_ssr(App)
    ///     .with_index_html(include_str!("../index.html"))?
    ///     .handle(req, env)
    ///     .await
    /// ```
    #[cfg(feature = "ssr")]
    pub fn with_index_html(mut self, html: &str) -> Result<Self, String> {
        let state = self
            .ssr
            .as_mut()
            .ok_or_else(|| "with_index_html requires with_ssr to be called first".to_string())?;
        state.index = IndexHtml::new(html)?;
        Ok(self)
    }

    /// Dispatch the request through registered `#[server]` functions.
    ///
    /// This is the async entry point — call it from `#[event(fetch)]`.
    #[allow(clippy::missing_errors_doc)]
    pub async fn handle(&self, req: Request, env: worker::Env) -> worker::Result<Response> {
        // Store env + request in thread-local for cf::env() / cf::req()
        let req_clone = req
            .clone()
            .map_err(|e| worker::Error::RustError(format!("request clone failed: {e}")))?;
        set_context(env, req_clone);

        // Store session config if configured
        if let Some(ref config) = self.session_config {
            crate::session::set_session_config(config.clone());
        }

        // Run before hooks — short-circuit if one returns a response
        for hook in &self.before {
            if let Some(resp) = hook(&req)? {
                return self.flush_and_finalize(resp).await;
            }
        }

        // WebSocket upgrade routing — check before Axum dispatch.
        // WebSocket 101 responses have immutable headers, so we return
        // them directly without running finalize() (no cookies/after hooks).
        if req
            .headers()
            .get("Upgrade")
            .ok()
            .flatten()
            .is_some_and(|v| v.eq_ignore_ascii_case("websocket"))
        {
            for route in &self.websocket_routes {
                if req.path().starts_with(&route.path_prefix) {
                    return (route.handler)(req).await;
                }
            }
        }

        // Convert worker::Request → http::Request
        let http_req = worker_req_to_http(req).await?;

        // Save URI before dispatch — needed for SSR fallback path
        #[cfg(feature = "ssr")]
        let uri = http_req.uri().clone();

        // Save Accept header for SSR content-type check
        #[cfg(feature = "ssr")]
        let accepts_html = http_req
            .headers()
            .get(http::header::ACCEPT)
            .and_then(|v| v.to_str().ok())
            .is_some_and(|v| v.contains("text/html"));

        // Dispatch through the Dioxus Axum router
        let http_resp = dispatch(http_req).await?;

        // SSR fallback: if Axum returned 404, SSR is configured, and the
        // client accepts HTML, render the app component instead.
        #[cfg(feature = "ssr")]
        if http_resp.status() == http::StatusCode::NOT_FOUND {
            if let Some(ref state) = self.ssr {
                if accepts_html {
                    let worker_resp = if state.streaming {
                        ssr::render_ssr_streaming(&uri, &state.build_vdom, &state.index)
                            .await?
                    } else {
                        ssr::render_ssr(&uri, &state.build_vdom, &state.index).await?
                    };
                    return self.flush_and_finalize(worker_resp).await;
                }
            }
        }

        // Convert http::Response → worker::Response (streaming)
        let worker_resp = http_to_worker_resp(http_resp)?;

        self.flush_and_finalize(worker_resp).await
    }

    /// Flush session data to the backend, then apply cookies and after hooks.
    async fn flush_and_finalize(&self, resp: Response) -> worker::Result<Response> {
        if self.session_config.is_some() {
            crate::session::flush_session().await?;
        }
        self.finalize(resp)
    }

    /// Apply queued cookies and run after hooks on the final response.
    fn finalize(&self, mut resp: Response) -> worker::Result<Response> {
        // Apply any cookies queued by cf::set_cookie() / cf::clear_cookie()
        for cookie in take_cookies() {
            resp.headers_mut()
                .append("Set-Cookie", &cookie)
                .map_err(|e| worker::Error::RustError(format!("cookie append failed: {e}")))?;
        }

        // Run after hooks
        for hook in &self.after {
            hook(&mut resp)?;
        }

        // Context is NOT cleared here — with streaming, the response body may
        // still be generating after handle() returns. The next request's
        // set_context() overwrites the old values, which is safe because Workers
        // handle one request at a time per isolate.

        Ok(resp)
    }
}

impl Default for Handler {
    fn default() -> Self {
        Self::new()
    }
}

/// Handle an incoming Cloudflare Worker request by dispatching it through
/// the Dioxus server function router.
///
/// This is a convenience wrapper around [`Handler::new().handle()`](Handler::handle).
/// Use [`Handler`] directly if you need before/after middleware hooks.
///
/// # Example
///
/// ```rust,ignore
/// use worker::*;
///
/// extern "C" { fn __wasm_call_ctors(); }
///
/// #[event(fetch)]
/// async fn fetch(req: Request, env: Env, _ctx: Context) -> Result<Response> {
///     // Required: initialize inventory for server function registration.
///     // SAFETY: Called once per Worker cold start. The `inventory` crate
///     // requires this in WASM to register #[server] functions.
///     unsafe { __wasm_call_ctors(); }
///
///     dioxus_cloudflare::handle(req, env).await
/// }
/// ```
#[allow(clippy::missing_errors_doc)]
pub async fn handle(req: Request, env: worker::Env) -> worker::Result<Response> {
    Handler::new().handle(req, env).await
}

/// Convert a `worker::Request` into an `http::Request<Body>` that Axum
/// can process.
async fn worker_req_to_http(mut req: Request) -> worker::Result<http::Request<Body>> {
    let method = match req.method() {
        worker::Method::Get => http::Method::GET,
        worker::Method::Post => http::Method::POST,
        worker::Method::Put => http::Method::PUT,
        worker::Method::Delete => http::Method::DELETE,
        worker::Method::Options => http::Method::OPTIONS,
        worker::Method::Head => http::Method::HEAD,
        worker::Method::Patch => http::Method::PATCH,
        _ => http::Method::GET,
    };

    let url = req.url()?;
    let uri: http::Uri = url
        .as_str()
        .parse()
        .map_err(|e| worker::Error::RustError(format!("invalid URI: {e}")))?;

    let mut builder = http::Request::builder().method(method).uri(uri);

    // Copy headers from worker request to http request
    for (key, value) in req.headers() {
        builder = builder.header(&key, &value);
    }

    // Read body as bytes and wrap in axum Body
    let body_bytes = req.bytes().await?;

    builder
        .body(Body::from(body_bytes))
        .map_err(|e| worker::Error::RustError(format!("request build failed: {e}")))
}

/// Build an Axum router containing all registered `#[server]` functions,
/// then dispatch the request through it.
///
/// Returns the response with the body passed through (not collected into
/// bytes), allowing streaming responses to flow through to the Worker.
async fn dispatch(req: http::Request<Body>) -> worker::Result<http::Response<Body>> {
    // Collect all #[server] functions registered by inventory and add
    // them as routes to an Axum router.
    //
    // ServerFunction::collect() returns all functions that were registered
    // when __wasm_call_ctors() ran. Each function knows its own path
    // (e.g., "/api/ping") and HTTP method.
    let mut router: axum::Router<dioxus_server::FullstackState> = axum::Router::new();
    for func in dioxus_server::ServerFunction::collect() {
        router = router.route(func.path(), func.method_router());
    }

    // Convert Router<FullstackState> → Router<()> using a headless state
    // (no SSR renderer needed — we only serve API endpoints).
    let router = router.with_state(dioxus_server::FullstackState::headless());

    // Dispatch the request through the router. Router's Service impl has
    // error type Infallible — it always produces a response (possibly 404).
    Ok(router
        .oneshot(req)
        .await
        .unwrap_or_else(|e: Infallible| match e {}))
}

/// Adapter that converts an `axum::body::Body` into a
/// `Stream<Item = Result<Vec<u8>, worker::Error>>`.
///
/// This satisfies `TryStream<Ok = Vec<u8>, Error = worker::Error>` via the
/// blanket impl, matching `ResponseBuilder::from_stream()`'s bounds.
///
/// `axum::body::Body` is `Unpin` (wraps `Pin<Box<...>>`), so no pin
/// projection is needed.
struct AxumBodyStream(Body);

impl futures_core::Stream for AxumBodyStream {
    type Item = Result<Vec<u8>, worker::Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            match Pin::new(&mut self.0).poll_frame(cx) {
                Poll::Ready(Some(Ok(frame))) => {
                    if let Ok(data) = frame.into_data() {
                        return Poll::Ready(Some(Ok(data.to_vec())));
                    }
                    // Trailers frame — skip, poll for more data
                }
                Poll::Ready(Some(Err(e))) => {
                    return Poll::Ready(Some(Err(worker::Error::RustError(format!(
                        "body frame error: {e}"
                    )))));
                }
                Poll::Ready(None) => return Poll::Ready(None),
                Poll::Pending => return Poll::Pending,
            }
        }
    }
}

/// Convert an `http::Response<Body>` into a `worker::Response` backed by
/// a `ReadableStream`, enabling streaming response bodies.
fn http_to_worker_resp(resp: http::Response<Body>) -> worker::Result<Response> {
    let (parts, body) = resp.into_parts();

    let mut worker_resp = ResponseBuilder::new()
        .with_status(parts.status.as_u16())
        .from_stream(AxumBodyStream(body))?;

    // Copy response headers
    for (key, value) in &parts.headers {
        if let Ok(v) = value.to_str() {
            worker_resp
                .headers_mut()
                .set(key.as_str(), v)
                .map_err(|e| worker::Error::RustError(format!("header set failed: {e}")))?;
        }
    }

    Ok(worker_resp)
}