From 7b4402238a19f187321a68088e6542d8d5fa8572 Mon Sep 17 00:00:00 2001
From: Dean Li <deantvv@gmail.com>
Date: Mon, 22 Feb 2021 12:11:11 +0800
Subject: [PATCH] Move the parsing header logic to args.rs

---
 Cargo.toml  |  2 +-
 src/args.rs | 37 +++++++++++++++++++++++++++++++++++--
 src/main.rs | 26 +++++++-------------------
 3 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 537b215..cc4d5fc 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -45,7 +45,7 @@ actix-files = "0.5"
 actix-multipart = "0.3.0"
 actix-web-httpauth = "0.5.0"
 mime = "0.3"
-httparse = "1.3.5"
+httparse = "1"
 
 [dev-dependencies]
 assert_cmd = "1.0"
diff --git a/src/args.rs b/src/args.rs
index 53f1638..ecfe151 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -119,8 +119,8 @@ struct CliArgs {
     title: Option<String>,
 
     /// Custom header from user
-    #[structopt(long = "header")]
-    header: Option<String>,
+    #[structopt(long = "header", parse(try_from_str = parse_header))]
+    header: Option<Header>,
 }
 
 /// Checks wether an interface is valid, i.e. it can be parsed into an IP address
@@ -174,6 +174,39 @@ 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];
+    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)
+}
+
 /// Parses the command line arguments
 pub fn parse_args() -> crate::MiniserveConfig {
     let args = CliArgs::from_args();
diff --git a/src/main.rs b/src/main.rs
index 149381f..0dd692e 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 args::Header;
 use std::io::{self, Write};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::thread;
@@ -90,7 +91,7 @@ pub struct MiniserveConfig {
     pub title: Option<String>,
 
     /// If specified, header will be added
-    pub header: Option<String>,
+    pub header: Option<Header>,
 }
 
 fn main() {
@@ -284,26 +285,13 @@ async fn run() -> Result<(), ContextualError> {
 }
 
 fn configure_header(conf: &MiniserveConfig) -> middleware::DefaultHeaders {
-    let mut headers = [httparse::EMPTY_HEADER; 16];
+    let header = conf.clone().header;
 
-    match conf.clone().header {
-        Some(mut header) => {
-            // parse_headers need header newline ends
-            header.push('\n');
-            httparse::parse_headers(header.as_bytes(), &mut headers).expect("Bad header");
-
-            let mut header_middleware = middleware::DefaultHeaders::new();
-
-            for h in headers.iter() {
-                if h.name != httparse::EMPTY_HEADER.name {
-                    println!("h={:?}", h);
-                    header_middleware = header_middleware.header(h.name, h.value);
-                }
-            }
-
-            header_middleware
+    match header {
+        Some(header) if header.name != httparse::EMPTY_HEADER.name => {
+            middleware::DefaultHeaders::new().header(&header.name, header.value)
         }
-        None => middleware::DefaultHeaders::new(),
+        _ => middleware::DefaultHeaders::new(),
     }
 }