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

dioxus_cloudflare/
handler.rs
