Move headers to Handler

Author: andrewhickman

src/config.rs

@@ -39,6 +39,8 @@ pub struct Config {
 pub struct Route {
     pub route: route::Route,
     pub rewrite_path: Option<String>,
+    #[serde(with = "http_serde::header_map", default)]
+    pub response_headers: http::HeaderMap,
     #[serde(flatten)]
     pub kind: RouteKind,
 }
@@ -74,8 +76,6 @@ pub struct JsonRoute {
 pub struct ProxyRoute {
     #[serde(rename = "url", with = "http_serde::uri")]
     pub uri: Uri,
-    #[serde(with = "http_serde::header_map", default)]
-    pub response_headers: http::HeaderMap,
 }
 
 impl Config {

src/handler/mod.rs

@@ -15,6 +15,7 @@ use crate::config;
 pub struct Handler {
     kind: HandlerKind,
     path_rewriter: Option<PathRewriter>,
+    response_headers: http::HeaderMap,
 }
 
 #[derive(Debug)]
@@ -37,6 +38,7 @@ impl Handler {
             rewrite_path,
             route,
             kind,
+            response_headers,
         } = route;
         let path_rewriter = rewrite_path.map(|replace| {
             let regex = route.to_regex();
@@ -53,6 +55,7 @@ impl Handler {
         Ok(Handler {
             path_rewriter,
             kind,
+            response_headers,
         })
     }
 
@@ -65,12 +68,18 @@ impl Handler {
             None => request.uri().path().to_owned(),
         };
 
-        match &self.kind {
+        let mut result = match &self.kind {
             HandlerKind::File(file) => file.handle(request).await,
             HandlerKind::Dir(dir) => dir.handle(request, &path).await,
             HandlerKind::Proxy(proxy) => proxy.handle(request, &path).await,
             HandlerKind::Json(json) => json.handle(request, &path).await,
+        };
+
+        if let Ok(response) = &mut result {
+            response.headers_mut().extend(self.response_headers.clone());
         }
+
+        result
     }
 }
 

src/handler/proxy.rs

@@ -52,10 +52,7 @@ impl ProxyHandler {
         log::debug!("Forwarding request to `{}`", request.uri());
 
         match self.client.request(request).await {
-            Ok(mut response) => {
-                response
-                    .headers_mut()
-                    .extend(self.config.response_headers.clone());
+            Ok(response) => {
                 Ok(response)
             }
             Err(err) => {