Allow set custom headers from CLI

2 files changed, 33 insertions, 0 deletions

diff --git a/src/args.rs b/src/args.rs
index f736941..53f1638 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -117,6 +117,10 @@ struct CliArgs {
     /// Shown instead of host in page title and heading
     #[structopt(short = "t", long = "title")]
     title: Option<String>,
+
+    /// Custom header from user
+    #[structopt(long = "header")]
+    header: Option<String>,
 }
 
 /// Checks wether an interface is valid, i.e. it can be parsed into an IP address
@@ -225,6 +229,7 @@ pub fn parse_args() -> crate::MiniserveConfig {
         zip_enabled: args.enable_zip,
         dirs_first: args.dirs_first,
         title: args.title,
+        header: args.header,
     }
 }
 
diff --git a/src/main.rs b/src/main.rs
index c55e77f..149381f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -88,6 +88,9 @@ pub struct MiniserveConfig {
 
     /// Shown instead of host in page title and heading
     pub title: Option<String>,
+
+    /// If specified, header will be added
+    pub header: Option<String>,
 }
 
 fn main() {
@@ -248,6 +251,7 @@ async fn run() -> Result<(), ContextualError> {
 
     let srv = actix_web::HttpServer::new(move || {
         App::new()
+            .wrap(configure_header(&inside_config.clone()))
             .app_data(inside_config.clone())
             .wrap(middleware::Condition::new(
                 !inside_config.auth.is_empty(),
@@ -279,6 +283,30 @@ async fn run() -> Result<(), ContextualError> {
         .map_err(|e| ContextualError::IoError("".to_owned(), e))
 }
 
+fn configure_header(conf: &MiniserveConfig) -> middleware::DefaultHeaders {
+    let mut headers = [httparse::EMPTY_HEADER; 16];
+
+    match conf.clone().header {
+        Some(mut header) => {
+            // parse_headers need header newline ends
+            header.push('\n');
+            httparse::parse_headers(header.as_bytes(), &mut headers).expect("Bad header");
+
+            let mut header_middleware = middleware::DefaultHeaders::new();
+
+            for h in headers.iter() {
+                if h.name != httparse::EMPTY_HEADER.name {
+                    println!("h={:?}", h);
+                    header_middleware = header_middleware.header(h.name, h.value);
+                }
+            }
+
+            header_middleware
+        }
+        None => middleware::DefaultHeaders::new(),
+    }
+}
+
 /// Configures the Actix application
 fn configure_app(app: &mut web::ServiceConfig, conf: &MiniserveConfig) {
     let random_route = conf.random_route.clone().unwrap_or_default();