Multiple headers support for custom headers

5 files changed, 55 insertions, 44 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 98a000e..55ba636 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1590,6 +1590,7 @@ dependencies = [
  "alphanumeric-sort",
  "assert_cmd",
  "assert_fs",
+ "bytes 1.0.1",
  "bytesize",
  "chrono",
  "chrono-humanize",
@@ -1597,6 +1598,7 @@ dependencies = [
  "futures",
  "grass",
  "hex",
+ "http",
  "httparse",
  "libflate",
  "log",
diff --git a/Cargo.toml b/Cargo.toml
index cc4d5fc..4d7201a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -46,6 +46,8 @@ actix-multipart = "0.3.0"
 actix-web-httpauth = "0.5.0"
 mime = "0.3"
 httparse = "1"
+http = "0.2.3"
+bytes = "1"
 
 [dev-dependencies]
 assert_cmd = "1.0"
diff --git a/src/args.rs b/src/args.rs
index ecfe151..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;
@@ -120,7 +122,7 @@ struct CliArgs {
 
     /// Custom header from user
     #[structopt(long = "header", parse(try_from_str = parse_header))]
-    header: Option<Header>,
+    header: Option<HeaderMap>,
 }
 
 /// Checks wether an interface is valid, i.e. it can be parsed into an IP address
@@ -174,37 +176,26 @@ fn parse_auth(src: &str) -> Result<auth::RequiredAuth, ContextualError> {
     })
 }
 
-/// A own header modified from [httparse](https://docs.rs/httparse/1.3.5/src/httparse/lib.rs.html#415-425)
-#[derive(Clone, Eq, PartialEq, Debug)]
-pub struct Header {
-    /// The name portion of a header.
-    ///
-    /// A header name must be valid ASCII-US, so it's safe to store as a `String`.
-    pub name: String,
-    /// The value portion of a header.
-    ///
-    /// While headers **should** be ASCII-US, the specification allows for
-    /// values that may not be, and so the value is stored as bytes.
-    pub value: Vec<u8>,
-}
-
-impl Header {
-    fn new(header: &httparse::Header) -> Self {
-        Header {
-            name: header.name.to_string(),
-            value: header.value.to_vec(),
-        }
-    }
-}
-
-fn parse_header(src: &str) -> Result<Header, httparse::Error> {
-    let mut headers = [httparse::EMPTY_HEADER; 1];
+/// 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 header = Header::new(&headers[0]);
-    Ok(header)
+    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
diff --git a/src/main.rs b/src/main.rs
index 0dd692e..44298d7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,7 +5,7 @@ use actix_web::{
 };
 use actix_web::{middleware, App, HttpRequest, HttpResponse};
 use actix_web_httpauth::middleware::HttpAuthentication;
-use args::Header;
+use http::header::HeaderMap;
 use std::io::{self, Write};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::thread;
@@ -91,7 +91,7 @@ pub struct MiniserveConfig {
     pub title: Option<String>,
 
     /// If specified, header will be added
-    pub header: Option<Header>,
+    pub header: Option<HeaderMap>,
 }
 
 fn main() {
@@ -285,11 +285,17 @@ async fn run() -> Result<(), ContextualError> {
 }
 
 fn configure_header(conf: &MiniserveConfig) -> middleware::DefaultHeaders {
-    let header = conf.clone().header;
-
-    match header {
-        Some(header) if header.name != httparse::EMPTY_HEADER.name => {
-            middleware::DefaultHeaders::new().header(&header.name, header.value)
+    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(),
     }
diff --git a/tests/header.rs b/tests/header.rs
index a0734e5..187730f 100644
--- a/tests/header.rs
+++ b/tests/header.rs
@@ -8,18 +8,17 @@ use std::process::{Command, Stdio};
 use std::thread::sleep;
 use std::time::Duration;
 
-#[rstest]
-fn custom_header_set(tmpdir: TempDir, port: u16) -> Result<(), Error> {
-    let header_name = "x-info";
-    let header_value = "123";
-    let header_str = format!("{}: {}", header_name, header_value);
-
-    let _ = Command::cargo_bin("miniserve")?
+#[rstest(header,
+    case("x-info: 123".to_string()),
+    case("x-info1: 123\r\nx-info2: 345".to_string())
+)]
+fn custom_header_set(tmpdir: TempDir, port: u16, header: String) -> Result<(), Error> {
+    let mut child = Command::cargo_bin("miniserve")?
         .arg(tmpdir.path())
         .arg("-p")
         .arg(port.to_string())
         .arg("--header")
-        .arg(header_str)
+        .arg(header.clone())
         .stdout(Stdio::null())
         .spawn()?;
 
@@ -27,7 +26,18 @@ fn custom_header_set(tmpdir: TempDir, port: u16) -> Result<(), Error> {
 
     let resp = reqwest::blocking::get(format!("http://localhost:{}", port).as_str())?;
 
-    assert_eq!(resp.headers().get(header_name).unwrap(), header_value);
+    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);
+        }
+    }
+
+    child.kill()?;
 
     Ok(())
 }