set server header

2 files changed, 18 insertions, 15 deletions

diff --git a/src/main.rs b/src/main.rs
index c325e61..980dd36 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,7 +15,7 @@ use error::ServiceError;
 use http_body_util::{combinators::BoxBody, BodyExt};
 use hyper::{
     body::Incoming,
-    header::{CONTENT_TYPE, HOST},
+    header::{CONTENT_TYPE, HOST, SERVER},
     http::HeaderValue,
     server::conn::http1,
     service::service_fn,
@@ -161,10 +161,24 @@ async fn service(
         ))
         .ok_or(ServiceError::NoHost)?;
 
-    match route {
+    let mut resp = match route {
         HostConfig::Backend { backend } => proxy_request(req, addr, backend).await,
         HostConfig::Files { files } => serve_files(req, files).await,
-    }
+    }?;
+
+    let server_header = resp.headers().get(SERVER).cloned();
+    eprintln!(" {server_header:?}");
+    resp.headers_mut().insert(
+        SERVER,
+        HeaderValue::from_str(&if let Some(o) = server_header {
+            format!("{} via gnix", o.to_str().ok().unwrap_or("invalid"))
+        } else {
+            format!("gnix")
+        })
+        .unwrap(),
+    );
+
+    return Ok(resp);
 }
 
 pub fn remove_port(s: &str) -> &str {
diff --git a/src/proxy.rs b/src/proxy.rs
index aff1ae5..036bbc1 100644
--- a/src/proxy.rs
+++ b/src/proxy.rs
@@ -2,7 +2,7 @@ use crate::ServiceError;
 use http_body_util::{combinators::BoxBody, BodyExt};
 use hyper::{
     body::Incoming,
-    header::{SERVER, UPGRADE},
+    header::UPGRADE,
     http::{
         uri::{PathAndQuery, Scheme},
         HeaderValue,
@@ -67,17 +67,6 @@ pub async fn proxy_request(
             .map_err(ServiceError::Hyper)?
     };
 
-    let server_header = resp.headers().get(SERVER).cloned();
-    resp.headers_mut().insert(
-        SERVER,
-        HeaderValue::from_str(&if let Some(o) = server_header {
-            format!("{} via gnix", o.to_str().unwrap())
-        } else {
-            format!("gnix")
-        })
-        .unwrap(),
-    );
-
     if do_upgrade {
         let on_upgrade_upstream = resp.extensions_mut().remove::<OnUpgrade>();
         tokio::task::spawn(async move {