Add config option to specify additional headers for serve

CHANGELOG.md

@@ -40,6 +40,7 @@ Subheadings to categorize changes are `added, changed, deprecated, removed, fixe
 - It is now possible to disable the hashes in output file names with the new `--filehash` flag (for example `cargo build --filehash false`). Alternatively the `build.filehash` setting in `Trunk.toml` or the env var `CARGO_BUILD_FILEHASH` can be used.
 - Flags for enabling reference types & weak references in `wasm-bindgen`.
 - Added the `data-typescript` attribute to Rust assets. When present, `wasm-bindgen` will emit TS files for the WASM module.
+- Added `headers` config option for `trunk serve`.
 
 ### changed
 - Bump notify to 5.0.0-pre.13, which fixes [notify-rs/notify#356](https://github.com/notify-rs/notify/issues/356)

Cargo.toml

@@ -44,7 +44,7 @@ tokio = { version = "1", default-features = false, features = ["full"] }
 tokio-stream = { version = "0.1", default-features = false, features = ["fs", "sync"] }
 tokio-tungstenite = "0.17"
 toml = "0.5"
-tower-http = { version = "0.3", features = ["fs", "trace"] }
+tower-http = { version = "0.3", features = ["fs", "trace", "set-header"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 which = "4"

Trunk.toml

@@ -27,6 +27,8 @@ port = 8080
 open = false
 # Disable auto-reload of the web app.
 no_autoreload = false
+# Additional headers set for responses.
+#headers = { "test-header" = "header value", "test-header2" = "header value 2" }
 
 [clean]
 # The output dir for all final assets.

src/config/models.rs

@@ -127,6 +127,10 @@ pub struct ConfigOptsServe {
     #[clap(long = "no-autoreload")]
     #[serde(default)]
     pub no_autoreload: bool,
+    /// Additional headers to send in responses [default: none]
+    #[clap(skip)]
+    #[serde(default)]
+    pub headers: HashMap<String, String>,
 }
 
 /// Config options for the serve system.
@@ -336,6 +340,7 @@ impl ConfigOpts {
             proxy_insecure: cli.proxy_insecure,
             proxy_ws: cli.proxy_ws,
             no_autoreload: cli.no_autoreload,
+            headers: cli.headers,
         };
         let cfg = ConfigOpts {
             build: None,
@@ -510,6 +515,7 @@ impl ConfigOpts {
                 if l.open {
                     g.open = true;
                 }
+                g.headers.extend(l.headers);
                 Some(g)
             }
         };

src/config/rt.rs

@@ -218,6 +218,8 @@ pub struct RtcServe {
     pub proxies: Option<Vec<ConfigOptsProxy>>,
     /// Whether to disable auto-reload of the web page when a build completes.
     pub no_autoreload: bool,
+    /// Additional headers to include in responses.
+    pub headers: HashMap<String, String>,
 }
 
 impl RtcServe {
@@ -247,6 +249,7 @@ impl RtcServe {
             proxy_ws: opts.proxy_ws,
             proxies,
             no_autoreload: opts.no_autoreload,
+            headers: opts.headers,
         })
     }
 }

src/serve.rs

@@ -1,17 +1,20 @@
+use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
 use axum::body::{self, Body};
 use axum::extract::ws::{WebSocket, WebSocketUpgrade};
 use axum::extract::Extension;
-use axum::http::StatusCode;
+use axum::http::header::HeaderName;
+use axum::http::{HeaderValue, StatusCode};
 use axum::response::Response;
 use axum::routing::{get, get_service, Router};
 use axum::Server;
 use tokio::sync::broadcast;
 use tokio::task::JoinHandle;
 use tower_http::services::{ServeDir, ServeFile};
+use tower_http::set_header::SetResponseHeaderLayer;
 use tower_http::trace::TraceLayer;
 
 use crate::common::SERVER;
@@ -117,7 +120,7 @@ impl ServeSystem {
             &cfg,
             build_done_chan,
         ));
-        let router = router(state, cfg.clone());
+        let router = router(state, cfg.clone())?;
         let addr = (cfg.address, cfg.port).into();
         let server = Server::bind(&addr)
             .serve(router.into_make_service())
@@ -147,6 +150,8 @@ pub struct State {
     pub build_done_chan: broadcast::Sender<()>,
     /// Whether to disable autoreload
     pub no_autoreload: bool,
+    /// Additional headers to add to responses.
+    pub headers: HashMap<String, String>,
 }
 
 impl State {
@@ -166,13 +171,14 @@ impl State {
             public_url,
             build_done_chan,
             no_autoreload: cfg.no_autoreload,
+            headers: cfg.headers.clone(),
         }
     }
 }
 
 /// Build the Trunk router, this includes that static file server, the WebSocket server,
 /// (for autoreload & HMR in the future), as well as any user-defined proxies.
-fn router(state: Arc<State>, cfg: Arc<RtcServe>) -> Router {
+fn router(state: Arc<State>, cfg: Arc<RtcServe>) -> Result<Router> {
     // Build static file server, middleware, error handler & WS route for reloads.
     let public_route = if state.public_url == "/" {
         &state.public_url
@@ -183,19 +189,28 @@ fn router(state: Arc<State>, cfg: Arc<RtcServe>) -> Router {
             .unwrap_or(&state.public_url)
     };
 
+    let mut serve_dir = get_service(
+        ServeDir::new(&state.dist_dir).fallback(ServeFile::new(&state.dist_dir.join(INDEX_HTML))),
+    );
+    for (key, value) in &state.headers {
+        let name = HeaderName::from_bytes(key.as_bytes())
+            .with_context(|| format!("invalid header {:?}", key))?;
+        let value: HeaderValue = value
+            .parse()
+            .with_context(|| format!("invalid header value {:?} for header {}", value, name))?;
+        serve_dir = serve_dir.layer(SetResponseHeaderLayer::overriding(name, value))
+    }
+
     let mut router = Router::new()
         .fallback(
             Router::new().nest(
                 public_route,
-                get_service(
-                    ServeDir::new(&state.dist_dir)
-                        .fallback(ServeFile::new(&state.dist_dir.join(INDEX_HTML))),
-                )
-                .handle_error(|error| async move {
-                    tracing::error!(?error, "failed serving static file");
-                    StatusCode::INTERNAL_SERVER_ERROR
-                })
-                .layer(TraceLayer::new_for_http()),
+                get_service(serve_dir)
+                    .handle_error(|error| async move {
+                        tracing::error!(?error, "failed serving static file");
+                        StatusCode::INTERNAL_SERVER_ERROR
+                    })
+                    .layer(TraceLayer::new_for_http()),
             ),
         )
         .route(
@@ -268,7 +283,7 @@ fn router(state: Arc<State>, cfg: Arc<RtcServe>) -> Router {
         }
     }
 
-    router
+    Ok(router)
 }
 
 async fn handle_ws(mut ws: WebSocket, state: Arc<State>) {