Merge pull request #452 from deantvv/custom-header

Allow set custom headers from CLI

5 files changed, 100 insertions, 0 deletions

diff --git a/Cargo.lock b/Cargo.lock
index de0952f..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,8 @@ dependencies = [
  "futures",
  "grass",
  "hex",
+ "http",
+ "httparse",
  "libflate",
  "log",
  "maud",
diff --git a/Cargo.toml b/Cargo.toml
index c0af4f3..4d7201a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -45,6 +45,9 @@ actix-files = "0.5"
 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 f736941..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;
@@ -117,6 +119,10 @@ struct CliArgs {
     /// Shown instead of host in page title and heading
     #[structopt(short = "t", long = "title")]
     title: Option<String>,
+
+    /// Custom header from user
+    #[structopt(long = "header", parse(try_from_str = parse_header))]
+    header: Option<HeaderMap>,
 }
 
 /// Checks wether an interface is valid, i.e. it can be parsed into an IP address
@@ -170,6 +176,28 @@ 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 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 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
 pub fn parse_args() -> crate::MiniserveConfig {
     let args = CliArgs::from_args();
@@ -225,6 +253,7 @@ pub fn parse_args() -> crate::MiniserveConfig {
         zip_enabled: args.enable_zip,
         dirs_first: args.dirs_first,
         title: args.title,
+        header: args.header,
     }
 }
 
diff --git a/src/main.rs b/src/main.rs
index c55e77f..44298d7 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 http::header::HeaderMap;
 use std::io::{self, Write};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::thread;
@@ -88,6 +89,9 @@ pub struct MiniserveConfig {
 
     /// Shown instead of host in page title and heading
     pub title: Option<String>,
+
+    /// If specified, header will be added
+    pub header: Option<HeaderMap>,
 }
 
 fn main() {
@@ -248,6 +252,7 @@ async fn run() -> Result<(), ContextualError> {
 
     let srv = actix_web::HttpServer::new(move || {
         App::new()
+            .wrap(configure_header(&inside_config.clone()))
             .app_data(inside_config.clone())
             .wrap(middleware::Condition::new(
                 !inside_config.auth.is_empty(),
@@ -279,6 +284,23 @@ async fn run() -> Result<(), ContextualError> {
         .map_err(|e| ContextualError::IoError("".to_owned(), e))
 }
 
+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);
+                }
+            }
+            default_headers
+        }
+        _ => middleware::DefaultHeaders::new(),
+    }
+}
+
 /// Configures the Actix application
 fn configure_app(app: &mut web::ServiceConfig, conf: &MiniserveConfig) {
     let random_route = conf.random_route.clone().unwrap_or_default();
diff --git a/tests/header.rs b/tests/header.rs
new file mode 100644
index 0000000..187730f
--- /dev/null
+++ b/tests/header.rs
@@ -0,0 +1,43 @@
+mod fixtures;
+
+use assert_cmd::prelude::*;
+use assert_fs::fixture::TempDir;
+use fixtures::{port, tmpdir, Error};
+use rstest::rstest;
+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())
+)]
+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.clone())
+        .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);
+        }
+    }
+
+    child.kill()?;
+
+    Ok(())
+}