Add config option to specify additional headers for serve

trunk-rs · Feb 17, 2022 · 138d7cb · 138d7cb
1 parent d46ba85
commit 138d7cb
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@ Subheadings to categorize changes are `added, changed, deprecated, removed, fixe
 ### added
 - Added the `--address` option for `trunk serve`.
 - Open autoreload websocket using wss when assets are served over a secure connection.
+- Added `headers` config option for `trunk server`.
 
 ### changed
 - Bump notify to 5.0.0-pre.13, which fixes [notify-rs/notify#356](https://github.com/notify-rs/notify/issues/356)

diff --git a/Trunk.toml b/Trunk.toml
@@ -25,6 +25,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" }
 
 [clean]
 # The output dir for all final assets.

diff --git a/src/config/models.rs b/src/config/models.rs
@@ -108,6 +108,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.
@@ -295,6 +299,7 @@ impl ConfigOpts {
             proxy_rewrite: cli.proxy_rewrite,
             proxy_ws: cli.proxy_ws,
             no_autoreload: cli.no_autoreload,
+            headers: cli.headers,
         };
         let cfg = ConfigOpts {
             build: None,
@@ -449,6 +454,7 @@ impl ConfigOpts {
                 if l.open {
                     g.open = true;
                 }
+                g.headers.extend(l.headers);
                 Some(g)
             }
         };

diff --git a/src/config/rt.rs b/src/config/rt.rs
@@ -157,6 +157,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 {
@@ -175,6 +177,7 @@ impl RtcServe {
             proxy_ws: opts.proxy_ws,
             proxies,
             no_autoreload: opts.no_autoreload,
+            headers: opts.headers,
         })
     }
 }

diff --git a/src/serve.rs b/src/serve.rs
@@ -19,6 +19,8 @@ use crate::common::SERVER;
 use crate::config::RtcServe;
 use crate::proxy::{ProxyHandlerHttp, ProxyHandlerWebSocket};
 use crate::watch::WatchSystem;
+use std::collections::HashMap;
+use axum::http::header::HeaderName;
 
 const INDEX_HTML: &str = "index.html";
 
@@ -123,6 +125,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 {
@@ -134,6 +138,7 @@ impl State {
             public_url,
             build_done_chan,
             no_autoreload: cfg.no_autoreload,
+            headers: cfg.headers.clone(),
         }
     }
 }
@@ -156,10 +161,18 @@ async fn serve_dist(req: Request<Body>) -> ServerResult<Response<Body>> {
         // If accept does not contain `*/*` or `text/html`, then return.
         (ResolveResult::NotFound, Some(Ok(accept_header))) if accept_header.contains("*/*") || accept_header.contains("text/html") => (),
         _ => {
-            return Ok(ResponseBuilder::new()
+            let mut response = ResponseBuilder::new()
                 .request(&req)
                 .build(res)
-                .context("error serving from dist dir")?);
+                .context("error serving from dist dir")?;
+            for (key, value) in &state.headers {
+                let name = HeaderName::from_bytes(key.as_bytes())
+                    .with_context(|| format!("invalid header {:?}", key))?;
+                let value = value.parse()
+                    .with_context(|| format!("invalid header value {:?}", value))?;
+                response.headers_mut().append(name, value);
+            }
+            return Ok(response);
         }
     };