feat: allow setting http response header in fullstack

Author: jkelleyrtp

packages/fullstack-core/src/streaming.rs

@@ -2,8 +2,8 @@ use crate::{HttpError, ServerFnError};
 use axum_core::{extract::FromRequest, response::IntoResponse};
 use dioxus_core::{try_consume_context, CapturedError};
 use dioxus_signals::{ReadableExt, Signal, WritableExt};
-use http::request::Parts;
 use http::StatusCode;
+use http::{request::Parts, HeaderMap};
 use std::{cell::RefCell, rc::Rc};
 
 /// The context provided by dioxus fullstack for server-side rendering.
@@ -13,6 +13,7 @@ use std::{cell::RefCell, rc::Rc};
 pub struct FullstackContext {
     current_status: Signal<StreamingStatus>,
     request_headers: Rc<RefCell<http::request::Parts>>,
+    response_headers: Rc<RefCell<Option<HeaderMap>>>,
     route_http_status: Signal<HttpError>,
 }
 
@@ -34,6 +35,7 @@ impl FullstackContext {
                 status: http::StatusCode::OK,
                 message: None,
             }),
+            response_headers: Rc::new(RefCell::new(Some(HeaderMap::new()))),
         }
     }
 
@@ -99,6 +101,23 @@ impl FullstackContext {
         self.route_http_status.set(status);
     }
 
+    /// Add a header to the response. This will be sent to the client when the response is committed.
+    pub fn add_response_header(
+        &self,
+        key: impl Into<http::header::HeaderName>,
+        value: impl Into<http::header::HeaderValue>,
+    ) {
+        if let Some(headers) = self.response_headers.borrow_mut().as_mut() {
+            headers.insert(key.into(), value.into());
+        }
+    }
+
+    /// Take the response headers out of the context. This will leave the context without any headers,
+    /// so it should only be called once when the response is being committed.
+    pub fn take_response_headers(&self) -> Option<HeaderMap> {
+        self.response_headers.borrow_mut().take()
+    }
+
     /// Set the current HTTP status for the route. This will be used when committing the response
     /// to the client.
     pub fn commit_http_status(status: StatusCode, message: Option<String>) {

packages/fullstack-server/src/server.rs

@@ -304,15 +304,23 @@ impl RenderHandleState {
             .await;
 
         match response {
-            Ok((status, freshness, rx)) => {
+            Ok((status, headers, freshness, rx)) => {
                 let mut response = Response::builder()
                     .status(status.status)
                     .header(CONTENT_TYPE, "text/html; charset=utf-8")
                     .body(Body::from_stream(rx))
                     .unwrap();
 
+                // Write our freshness header
                 freshness.write(response.headers_mut());
 
+                // write the other headers set by the user
+                for (key, value) in headers.into_iter() {
+                    if let Some(key) = key {
+                        response.headers_mut().insert(key, value);
+                    }
+                }
+
                 response
             }
 

packages/fullstack-server/src/ssr.rs

@@ -17,7 +17,7 @@ use dioxus_router::ParseRouteError;
 use dioxus_ssr::Renderer;
 use futures_channel::mpsc::Sender;
 use futures_util::{Stream, StreamExt};
-use http::{request::Parts, StatusCode};
+use http::{request::Parts, HeaderMap, StatusCode};
 use std::{
     collections::HashMap,
     fmt::Write,
@@ -103,6 +103,7 @@ impl SsrRendererPool {
     ) -> Result<
         (
             HttpError,
+            HeaderMap,
             RenderFreshness,
             impl Stream<Item = Result<String, IncrementalRendererError>>,
         ),
@@ -154,6 +155,7 @@ impl SsrRendererPool {
                     status: StatusCode::OK,
                     message: None,
                 },
+                HeaderMap::new(),
                 freshness,
                 ReceiverWithDrop {
                     receiver: rx,
@@ -292,9 +294,12 @@ impl SsrRendererPool {
 
             // Check the FullstackContext in case the user set the statuscode manually or via a layout.
             let http_status = streaming_context.current_http_status();
+            let headers = streaming_context
+                .take_response_headers()
+                .unwrap_or_default();
 
             // Now that we handled any errors from rendering, we can send the initial ok result
-            _ = initial_result_tx.send(Ok(http_status));
+            _ = initial_result_tx.send(Ok((http_status, headers)));
 
             // Wait long enough to assemble the `<head>` of the document before starting to stream
             let mut pre_body = String::new();
@@ -418,12 +423,13 @@ impl SsrRendererPool {
         let join_handle = Self::spawn_platform(create_render_future);
 
         // Wait for the initial result which determines the status code
-        let status = initial_result_rx
+        let (status, headers) = initial_result_rx
             .await
             .map_err(|err| SSRError::Incremental(IncrementalRendererError::Other(err.into())))??;
 
         Ok((
             status,
+            headers,
             RenderFreshness::now(None),
             ReceiverWithDrop {
                 receiver: rx,