Re-do the way setting multiple headers works

2 files changed, 12 insertions, 17 deletions

diff --git a/src/args.rs b/src/args.rs
index 7710cce..8fdb067 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -120,9 +120,9 @@ struct CliArgs {
     #[structopt(short = "t", long = "title")]
     title: Option<String>,
 
-    /// Custom header from user
-    #[structopt(long = "header", parse(try_from_str = parse_header))]
-    header: Option<HeaderMap>,
+    /// Set custom header for responses
+    #[structopt(long = "header", parse(try_from_str = parse_header), number_of_values = 1)]
+    header: Vec<HeaderMap>,
 }
 
 /// Checks wether an interface is valid, i.e. it can be parsed into an IP address
@@ -178,15 +178,13 @@ 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 headers = [httparse::EMPTY_HEADER; 1];
     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 let Some(h) = headers.first() {
         if h.name != httparse::EMPTY_HEADER.name {
             header_map.insert(
                 HeaderName::from_bytes(&Bytes::copy_from_slice(h.name.as_bytes())).unwrap(),
diff --git a/src/main.rs b/src/main.rs
index 44298d7..a407136 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -91,7 +91,7 @@ pub struct MiniserveConfig {
     pub title: Option<String>,
 
     /// If specified, header will be added
-    pub header: Option<HeaderMap>,
+    pub header: Vec<HeaderMap>,
 }
 
 fn main() {
@@ -287,18 +287,15 @@ async fn run() -> Result<(), ContextualError> {
 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);
-                }
+    let mut default_headers = middleware::DefaultHeaders::new();
+    for header in headers {
+        for (header_name, header_value) in header.into_iter() {
+            if let Some(header_name) = header_name {
+                default_headers = default_headers.header(&header_name, header_value);
             }
-            default_headers
         }
-        _ => middleware::DefaultHeaders::new(),
     }
+    default_headers
 }
 
 /// Configures the Actix application