Multiple headers support for custom headers

2 files changed, 32 insertions, 35 deletions

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(),
     }