dioxus_cloudflare/
handler.rs

1//! Request handler that bridges Cloudflare Workers to the Dioxus Axum router.
2//!
3//! Two entry points are available:
4//!
5//! - [`handle`] — simple one-shot dispatch (no middleware)
6//! - [`Handler`] — builder with before/after middleware hooks
7//!
8//! Both store the Worker `Env` and `Request` in thread-local context,
9//! convert `worker::Request` → `http::Request`, dispatch through the
10//! Dioxus Axum router, and stream the response back via `ReadableStream`.
11
12use std::convert::Infallible;
13use std::pin::Pin;
14use std::task::{Context, Poll};
15
16use axum::body::Body;
17use http_body::Body as HttpBody;
18use tower::ServiceExt;
19use worker::{Request, Response, ResponseBuilder};
20
21use crate::context::{set_context, take_cookies};
22
23#[cfg(feature = "ssr")]
24use crate::ssr::{self, IndexHtml};
25
26type BeforeHook = Box<dyn Fn(&Request) -> worker::Result<Option<Response>>>;
27type AfterHook = Box<dyn Fn(&mut Response) -> worker::Result<()>>;
28
29/// Internal state for SSR rendering.
30#[cfg(feature = "ssr")]
31struct SsrState {
32    build_vdom: Box<dyn Fn() -> dioxus::dioxus_core::VirtualDom>,
33    index: IndexHtml,
34}
35
36/// Configurable request handler with before/after middleware hooks.
37///
38/// # Example
39///
40/// ```rust,ignore
41/// use dioxus_cloudflare::Handler;
42/// use worker::*;
43///
44/// #[event(fetch)]
45/// async fn fetch(req: Request, env: Env, _ctx: Context) -> Result<Response> {
46///     unsafe { __wasm_call_ctors(); }
47///
48///     Handler::new()
49///         .before(|req| {
50///             // Short-circuit OPTIONS requests for CORS preflight
51///             if req.method() == worker::Method::Options {
52///                 let mut resp = Response::empty()?;
53///                 resp.headers_mut().set("Access-Control-Allow-Origin", "*")?;
54///                 resp.headers_mut().set("Access-Control-Allow-Methods", "GET,POST,OPTIONS")?;
55///                 return Ok(Some(resp));
56///             }
57///             Ok(None) // continue to Axum dispatch
58///         })
59///         .after(|resp| {
60///             resp.headers_mut().set("Access-Control-Allow-Origin", "*")?;
61///             Ok(())
62///         })
63///         .handle(req, env)
64///         .await
65/// }
66/// ```
67pub struct Handler {
68    before: Vec<BeforeHook>,
69    after: Vec<AfterHook>,
70    #[cfg(feature = "ssr")]
71    ssr: Option<SsrState>,
72}
73
74impl Handler {
75    /// Create a new handler with no middleware hooks.
76    #[must_use]
77    pub fn new() -> Self {
78        Self {
79            before: Vec::new(),
80            after: Vec::new(),
81            #[cfg(feature = "ssr")]
82            ssr: None,
83        }
84    }
85
86    /// Add a before-dispatch hook.
87    ///
88    /// Runs after context is set (so `cf::env()`, `cf::d1()`, etc. work),
89    /// before Axum dispatch. Return `Ok(None)` to continue,
90    /// `Ok(Some(resp))` to short-circuit with a custom response.
91    ///
92    /// Hooks run in the order they were added.
93    #[must_use]
94    pub fn before(
95        mut self,
96        hook: impl Fn(&Request) -> worker::Result<Option<Response>> + 'static,
97    ) -> Self {
98        self.before.push(Box::new(hook));
99        self
100    }
101
102    /// Add an after-dispatch hook.
103    ///
104    /// Runs after cookies are applied, before returning the response.
105    /// Use this to add headers, log, or modify the final response.
106    ///
107    /// After hooks also run on short-circuited responses from before hooks.
108    /// Hooks run in the order they were added.
109    #[must_use]
110    pub fn after(
111        mut self,
112        hook: impl Fn(&mut Response) -> worker::Result<()> + 'static,
113    ) -> Self {
114        self.after.push(Box::new(hook));
115        self
116    }
117
118    /// Enable SSR rendering for non-API requests.
119    ///
120    /// When the Axum router returns 404 and the request accepts `text/html`,
121    /// the handler renders the given component to HTML and returns it.
122    ///
123    /// Uses a default HTML shell (`<!DOCTYPE html>` with `<div id="main">`).
124    /// Call [`with_index_html`](Self::with_index_html) to provide your own
125    /// shell (e.g., one that loads client WASM for SPA takeover).
126    ///
127    /// # Example
128    ///
129    /// ```rust,ignore
130    /// Handler::new()
131    ///     .with_ssr(App)
132    ///     .handle(req, env)
133    ///     .await
134    /// ```
135    #[cfg(feature = "ssr")]
136    #[must_use]
137    pub fn with_ssr(mut self, app: fn() -> dioxus::prelude::Element) -> Self {
138        self.ssr = Some(SsrState {
139            build_vdom: Box::new(move || {
140                dioxus::dioxus_core::VirtualDom::new(app)
141            }),
142            index: IndexHtml::default_shell(),
143        });
144        self
145    }
146
147    /// Provide a custom `index.html` shell for SSR responses.
148    ///
149    /// The HTML must contain an element with `id="main"` — rendered
150    /// component output is inserted at that point.
151    ///
152    /// Must be called **after** [`with_ssr`](Self::with_ssr).
153    ///
154    /// # Errors
155    ///
156    /// Returns `Err` if:
157    /// - SSR has not been configured (call `with_ssr` first)
158    /// - The HTML does not contain `id="main"`
159    ///
160    /// # Example
161    ///
162    /// ```rust,ignore
163    /// Handler::new()
164    ///     .with_ssr(App)
165    ///     .with_index_html(include_str!("../index.html"))?
166    ///     .handle(req, env)
167    ///     .await
168    /// ```
169    #[cfg(feature = "ssr")]
170    pub fn with_index_html(mut self, html: &str) -> Result<Self, String> {
171        let state = self
172            .ssr
173            .as_mut()
174            .ok_or_else(|| "with_index_html requires with_ssr to be called first".to_string())?;
175        state.index = IndexHtml::new(html)?;
176        Ok(self)
177    }
178
179    /// Dispatch the request through registered `#[server]` functions.
180    ///
181    /// This is the async entry point — call it from `#[event(fetch)]`.
182    #[allow(clippy::missing_errors_doc)]
183    pub async fn handle(&self, req: Request, env: worker::Env) -> worker::Result<Response> {
184        // Store env + request in thread-local for cf::env() / cf::req()
185        let req_clone = req
186            .clone()
187            .map_err(|e| worker::Error::RustError(format!("request clone failed: {e}")))?;
188        set_context(env, req_clone);
189
190        // Run before hooks — short-circuit if one returns a response
191        for hook in &self.before {
192            if let Some(resp) = hook(&req)? {
193                return self.finalize(resp);
194            }
195        }
196
197        // Convert worker::Request → http::Request
198        let http_req = worker_req_to_http(req).await?;
199
200        // Save URI before dispatch — needed for SSR fallback path
201        #[cfg(feature = "ssr")]
202        let uri = http_req.uri().clone();
203
204        // Save Accept header for SSR content-type check
205        #[cfg(feature = "ssr")]
206        let accepts_html = http_req
207            .headers()
208            .get(http::header::ACCEPT)
209            .and_then(|v| v.to_str().ok())
210            .is_some_and(|v| v.contains("text/html"));
211
212        // Dispatch through the Dioxus Axum router
213        let http_resp = dispatch(http_req).await?;
214
215        // SSR fallback: if Axum returned 404, SSR is configured, and the
216        // client accepts HTML, render the app component instead.
217        #[cfg(feature = "ssr")]
218        if http_resp.status() == http::StatusCode::NOT_FOUND {
219            if let Some(ref state) = self.ssr {
220                if accepts_html {
221                    let worker_resp =
222                        ssr::render_ssr(&uri, &state.build_vdom, &state.index).await?;
223                    return self.finalize(worker_resp);
224                }
225            }
226        }
227
228        // Convert http::Response → worker::Response (streaming)
229        let worker_resp = http_to_worker_resp(http_resp)?;
230
231        self.finalize(worker_resp)
232    }
233
234    /// Apply queued cookies and run after hooks on the final response.
235    fn finalize(&self, mut resp: Response) -> worker::Result<Response> {
236        // Apply any cookies queued by cf::set_cookie() / cf::clear_cookie()
237        for cookie in take_cookies() {
238            resp.headers_mut()
239                .append("Set-Cookie", &cookie)
240                .map_err(|e| worker::Error::RustError(format!("cookie append failed: {e}")))?;
241        }
242
243        // Run after hooks
244        for hook in &self.after {
245            hook(&mut resp)?;
246        }
247
248        // Context is NOT cleared here — with streaming, the response body may
249        // still be generating after handle() returns. The next request's
250        // set_context() overwrites the old values, which is safe because Workers
251        // handle one request at a time per isolate.
252
253        Ok(resp)
254    }
255}
256
257impl Default for Handler {
258    fn default() -> Self {
259        Self::new()
260    }
261}
262
263/// Handle an incoming Cloudflare Worker request by dispatching it through
264/// the Dioxus server function router.
265///
266/// This is a convenience wrapper around [`Handler::new().handle()`](Handler::handle).
267/// Use [`Handler`] directly if you need before/after middleware hooks.
268///
269/// # Example
270///
271/// ```rust,ignore
272/// use worker::*;
273///
274/// extern "C" { fn __wasm_call_ctors(); }
275///
276/// #[event(fetch)]
277/// async fn fetch(req: Request, env: Env, _ctx: Context) -> Result<Response> {
278///     // Required: initialize inventory for server function registration.
279///     // SAFETY: Called once per Worker cold start. The `inventory` crate
280///     // requires this in WASM to register #[server] functions.
281///     unsafe { __wasm_call_ctors(); }
282///
283///     dioxus_cloudflare::handle(req, env).await
284/// }
285/// ```
286#[allow(clippy::missing_errors_doc)]
287pub async fn handle(req: Request, env: worker::Env) -> worker::Result<Response> {
288    Handler::new().handle(req, env).await
289}
290
291/// Convert a `worker::Request` into an `http::Request<Body>` that Axum
292/// can process.
293async fn worker_req_to_http(mut req: Request) -> worker::Result<http::Request<Body>> {
294    let method = match req.method() {
295        worker::Method::Get => http::Method::GET,
296        worker::Method::Post => http::Method::POST,
297        worker::Method::Put => http::Method::PUT,
298        worker::Method::Delete => http::Method::DELETE,
299        worker::Method::Options => http::Method::OPTIONS,
300        worker::Method::Head => http::Method::HEAD,
301        worker::Method::Patch => http::Method::PATCH,
302        _ => http::Method::GET,
303    };
304
305    let url = req.url()?;
306    let uri: http::Uri = url
307        .as_str()
308        .parse()
309        .map_err(|e| worker::Error::RustError(format!("invalid URI: {e}")))?;
310
311    let mut builder = http::Request::builder().method(method).uri(uri);
312
313    // Copy headers from worker request to http request
314    for (key, value) in req.headers() {
315        builder = builder.header(&key, &value);
316    }
317
318    // Read body as bytes and wrap in axum Body
319    let body_bytes = req.bytes().await?;
320
321    builder
322        .body(Body::from(body_bytes))
323        .map_err(|e| worker::Error::RustError(format!("request build failed: {e}")))
324}
325
326/// Build an Axum router containing all registered `#[server]` functions,
327/// then dispatch the request through it.
328///
329/// Returns the response with the body passed through (not collected into
330/// bytes), allowing streaming responses to flow through to the Worker.
331async fn dispatch(req: http::Request<Body>) -> worker::Result<http::Response<Body>> {
332    // Collect all #[server] functions registered by inventory and add
333    // them as routes to an Axum router.
334    //
335    // ServerFunction::collect() returns all functions that were registered
336    // when __wasm_call_ctors() ran. Each function knows its own path
337    // (e.g., "/api/ping") and HTTP method.
338    let mut router: axum::Router<dioxus_server::FullstackState> = axum::Router::new();
339    for func in dioxus_server::ServerFunction::collect() {
340        router = router.route(func.path(), func.method_router());
341    }
342
343    // Convert Router<FullstackState> → Router<()> using a headless state
344    // (no SSR renderer needed — we only serve API endpoints).
345    let router = router.with_state(dioxus_server::FullstackState::headless());
346
347    // Dispatch the request through the router. Router's Service impl has
348    // error type Infallible — it always produces a response (possibly 404).
349    Ok(router
350        .oneshot(req)
351        .await
352        .unwrap_or_else(|e: Infallible| match e {}))
353}
354
355/// Adapter that converts an `axum::body::Body` into a
356/// `Stream<Item = Result<Vec<u8>, worker::Error>>`.
357///
358/// This satisfies `TryStream<Ok = Vec<u8>, Error = worker::Error>` via the
359/// blanket impl, matching `ResponseBuilder::from_stream()`'s bounds.
360///
361/// `axum::body::Body` is `Unpin` (wraps `Pin<Box<...>>`), so no pin
362/// projection is needed.
363struct AxumBodyStream(Body);
364
365impl futures_core::Stream for AxumBodyStream {
366    type Item = Result<Vec<u8>, worker::Error>;
367
368    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
369        loop {
370            match Pin::new(&mut self.0).poll_frame(cx) {
371                Poll::Ready(Some(Ok(frame))) => {
372                    if let Ok(data) = frame.into_data() {
373                        return Poll::Ready(Some(Ok(data.to_vec())));
374                    }
375                    // Trailers frame — skip, poll for more data
376                }
377                Poll::Ready(Some(Err(e))) => {
378                    return Poll::Ready(Some(Err(worker::Error::RustError(format!(
379                        "body frame error: {e}"
380                    )))));
381                }
382                Poll::Ready(None) => return Poll::Ready(None),
383                Poll::Pending => return Poll::Pending,
384            }
385        }
386    }
387}
388
389/// Convert an `http::Response<Body>` into a `worker::Response` backed by
390/// a `ReadableStream`, enabling streaming response bodies.
391fn http_to_worker_resp(resp: http::Response<Body>) -> worker::Result<Response> {
392    let (parts, body) = resp.into_parts();
393
394    let mut worker_resp = ResponseBuilder::new()
395        .with_status(parts.status.as_u16())
396        .from_stream(AxumBodyStream(body))?;
397
398    // Copy response headers
399    for (key, value) in &parts.headers {
400        if let Ok(v) = value.to_str() {
401            worker_resp
402                .headers_mut()
403                .set(key.as_str(), v)
404                .map_err(|e| worker::Error::RustError(format!("header set failed: {e}")))?;
405        }
406    }
407
408    Ok(worker_resp)
409}
dioxus_cloudflare/handler.rs

dioxus_cloudflare/
handler.rs
