Merge pull request #452 from deantvv/custom-header

Allow set custom headers from CLI

2 files changed, 51 insertions, 0 deletions

diff --git a/src/args.rs b/src/args.rs
index f736941..7710cce 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -1,3 +1,5 @@
+use bytes::Bytes;
+use http::header::{HeaderMap, HeaderName, HeaderValue};
 use port_check::free_local_port;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 use std::path::PathBuf;
@@ -117,6 +119,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", parse(try_from_str = parse_header))]
+    header: Option<HeaderMap>,
 }
 
 /// Checks wether an interface is valid, i.e. it can be parsed into an IP address
@@ -170,6 +176,28 @@ fn parse_auth(src: &str) -> Result<auth::RequiredAuth, ContextualError> {
     })
 }
 
+/// Custom header parser (allow multiple headers input)
+pub fn parse_header(src: &str) -> Result<HeaderMap, httparse::Error> {
+    // Max customized header is limitted to 16
+    let mut headers = [httparse::EMPTY_HEADER; 16];
+    let mut header = src.to_string();
+    header.push('\n');
+    httparse::parse_headers(header.as_bytes(), &mut headers)?;
+
+    let mut header_map = HeaderMap::new();
+
+    for h in headers.iter() {
+        if h.name != httparse::EMPTY_HEADER.name {
+            header_map.insert(
+                HeaderName::from_bytes(&Bytes::copy_from_slice(h.name.as_bytes())).unwrap(),
+                HeaderValue::from_bytes(h.value).unwrap(),
+            );
+        }
+    }
+
+    Ok(header_map)
+}
+
 /// Parses the command line arguments
 pub fn parse_args() -> crate::MiniserveConfig {
     let args = CliArgs::from_args();
@@ -225,6 +253,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..44298d7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,6 +5,7 @@ use actix_web::{
 };
 use actix_web::{middleware, App, HttpRequest, HttpResponse};
 use actix_web_httpauth::middleware::HttpAuthentication;
+use http::header::HeaderMap;
 use std::io::{self, Write};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::thread;
@@ -88,6 +89,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<HeaderMap>,
 }
 
 fn main() {
@@ -248,6 +252,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 +284,23 @@ async fn run() -> Result<(), ContextualError> {
         .map_err(|e| ContextualError::IoError("".to_owned(), e))
 }
 
+fn configure_header(conf: &MiniserveConfig) -> middleware::DefaultHeaders {
+    let headers = conf.clone().header;
+
+    match headers {
+        Some(headers) => {
+            let mut default_headers = middleware::DefaultHeaders::new();
+            for (header_name, header_value) in headers.into_iter() {
+                if let Some(header_name) = header_name {
+                    default_headers = default_headers.header(&header_name, header_value);
+                }
+            }
+            default_headers
+        }
+        _ => 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();