Re-do the way setting multiple headers works

4 files changed, 25 insertions, 34 deletions

diff --git a/README.md b/README.md
index 2e1e14f..5418b8a 100644
--- a/README.md
+++ b/README.md
@@ -118,6 +118,8 @@ Sometimes this is just a more practical and quick way than doing things properly
         -d, --color-scheme-dark <color-scheme-dark>
                 Default color scheme [default: archlinux]  [possible values: squirrel, archlinux,
                 zenburn, monokai]
+            --header <header>...
+                Set custom header for responses
             --index <index_file>
                 The name of a directory index file to serve, like "index.html"
 
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
diff --git a/tests/header.rs b/tests/header.rs
index 187730f..774e5ef 100644
--- a/tests/header.rs
+++ b/tests/header.rs
@@ -8,33 +8,27 @@ use std::process::{Command, Stdio};
 use std::thread::sleep;
 use std::time::Duration;
 
-#[rstest(header,
-    case("x-info: 123".to_string()),
-    case("x-info1: 123\r\nx-info2: 345".to_string())
+#[rstest(headers,
+    case(vec!["x-info: 123".to_string()]),
+    case(vec!["x-info1: 123".to_string(), "x-info2: 345".to_string()])
 )]
-fn custom_header_set(tmpdir: TempDir, port: u16, header: String) -> Result<(), Error> {
+fn custom_header_set(tmpdir: TempDir, port: u16, headers: Vec<String>) -> Result<(), Error> {
     let mut child = Command::cargo_bin("miniserve")?
         .arg(tmpdir.path())
         .arg("-p")
         .arg(port.to_string())
-        .arg("--header")
-        .arg(header.clone())
-        .stdout(Stdio::null())
-        .spawn()?;
+        .args(headers.iter().flat_map(|h| vec!["--header", h]))
+        .stdout(Stdio::null()).spawn()?;
 
     sleep(Duration::from_secs(1));
 
     let resp = reqwest::blocking::get(format!("http://localhost:{}", port).as_str())?;
 
-    let mut headers = [httparse::EMPTY_HEADER; 4];
-    let mut header = header.clone();
-    header.push('\n');
-    httparse::parse_headers(header.as_bytes(), &mut headers)?;
-
-    for h in headers.iter() {
-        if h.name != httparse::EMPTY_HEADER.name {
-            assert_eq!(resp.headers().get(h.name).unwrap(), h.value);
-        }
+    for header in headers {
+        let mut header_split = header.splitn(2, ':');
+        let header_name = header_split.next().unwrap();
+        let header_value = header_split.next().unwrap().trim();
+        assert_eq!(resp.headers().get(header_name).unwrap(), header_value);
     }
 
     child.kill()?;