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::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};

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

/// 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>,
}

impl Handler {
    /// Create a new handler with no middleware hooks.
    #[must_use]
    pub fn new() -> Self {
        Self {
            before: Vec::new(),
            after: Vec::new(),
        }
    }

    /// 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
    }

    /// 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);

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

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

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

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

        self.finalize(worker_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)
}