Support setting request method, headers, and body (lnx-search#32)

* support setting headers

* support setting body

* support setting method

src/bench.rs

@@ -1,9 +1,11 @@
 use std::fmt::Display;
 use std::time::Duration;
 
+use ::http::{HeaderMap, Method};
 use anyhow::{anyhow, Result};
 use colored::*;
 use futures_util::StreamExt;
+use hyper::body::Bytes;
 
 use crate::results::WorkerResult;
 use crate::utils::div_mod;
@@ -36,6 +38,15 @@ pub struct BenchmarkSettings {
 
     /// The number of rounds to repeat.
     pub rounds: usize,
+
+    /// The request method.
+    pub method: Method,
+
+    /// Additional request headers.
+    pub headers: HeaderMap,
+
+    /// Request body.
+    pub body: Bytes,
 }
 
 /// Builds the runtime with the given settings and blocks on the main future.
@@ -79,6 +90,9 @@ async fn run(settings: BenchmarkSettings) -> Result<()> {
         settings.connections,
         settings.host.trim().to_string(),
         settings.bench_type,
+        settings.method,
+        settings.headers,
+        settings.body,
         predict_size as usize,
     )
     .await;

src/http/mod.rs

@@ -6,7 +6,8 @@ use anyhow::anyhow;
 use futures_util::stream::FuturesUnordered;
 use futures_util::TryFutureExt;
 use http::header::{self, HeaderMap};
-use http::Request;
+use http::{Method, Request};
+use hyper::body::Bytes;
 use hyper::client::conn::{self, SendRequest};
 use hyper::Body;
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -51,10 +52,14 @@ pub async fn start_tasks(
     connections: usize,
     uri_string: String,
     bench_type: BenchType,
+    method: Method,
+    headers: HeaderMap,
+    body: Bytes,
     _predicted_size: usize,
 ) -> anyhow::Result<FuturesUnordered<Handle>> {
     let deadline = Instant::now() + time_for;
-    let user_input = UserInput::new(bench_type, uri_string).await?;
+    let user_input =
+        UserInput::new(bench_type, uri_string, method, headers, body).await?;
 
     let handles = FuturesUnordered::new();
 
@@ -95,14 +100,17 @@ async fn benchmark(
         request_headers.insert(header::HOST, user_input.host_header);
     }
 
+    request_headers.extend(user_input.headers);
+
     let mut request_times = Vec::new();
     let mut error_map = HashMap::new();
 
     // Benchmark loop.
     // Futures must not be awaited without timeout.
     loop {
         // Create request from **parsed** data.
-        let mut request = Request::new(Body::empty());
+        let mut request = Request::new(Body::from(user_input.body.clone()));
+        *request.method_mut() = user_input.method.clone();
         *request.uri_mut() = user_input.uri.clone();
         *request.headers_mut() = request_headers.clone();
 

src/http/user_input.rs

@@ -4,6 +4,8 @@ use std::net::{SocketAddr, ToSocketAddrs};
 use anyhow::{anyhow, Result};
 use http::header::HeaderValue;
 use http::uri::Uri;
+use http::{HeaderMap, Method};
+use hyper::body::Bytes;
 use tokio::task::spawn_blocking;
 use tokio_native_tls::TlsConnector;
 
@@ -31,16 +33,33 @@ pub(crate) struct UserInput {
     pub(crate) host: String,
     pub(crate) host_header: HeaderValue,
     pub(crate) uri: Uri,
+    pub(crate) method: Method,
+    pub(crate) headers: HeaderMap,
+    pub(crate) body: Bytes,
 }
 
 impl UserInput {
-    pub(crate) async fn new(protocol: BenchType, string: String) -> Result<Self> {
-        spawn_blocking(move || Self::blocking_new(protocol, string))
-            .await
-            .unwrap()
+    pub(crate) async fn new(
+        protocol: BenchType,
+        string: String,
+        method: Method,
+        headers: HeaderMap,
+        body: Bytes,
+    ) -> Result<Self> {
+        spawn_blocking(move || {
+            Self::blocking_new(protocol, string, method, headers, body)
+        })
+        .await
+        .unwrap()
     }
 
-    fn blocking_new(protocol: BenchType, string: String) -> Result<Self> {
+    fn blocking_new(
+        protocol: BenchType,
+        string: String,
+        method: Method,
+        headers: HeaderMap,
+        body: Bytes,
+    ) -> Result<Self> {
         let uri = Uri::try_from(string)?;
         let scheme = uri
             .scheme()
@@ -91,6 +110,9 @@ impl UserInput {
             host,
             host_header,
             uri,
+            method,
+            headers,
+            body,
         })
     }
 }

src/main.rs

@@ -1,7 +1,12 @@
 extern crate clap;
 
-use anyhow::{Error, Result};
+use std::str::FromStr;
+
+use ::http::header::HeaderName;
+use ::http::{HeaderMap, HeaderValue, Method};
+use anyhow::{Context, Error, Result};
 use clap::{App, Arg, ArgMatches};
+use hyper::body::Bytes;
 use regex::Regex;
 use tokio::time::Duration;
 
@@ -77,6 +82,33 @@ fn main() {
         .parse::<usize>()
         .unwrap_or(1);
 
+    let method = match args
+        .value_of("method")
+        .map(|method| Method::from_str(&method.to_uppercase()))
+        .transpose()
+    {
+        Ok(method) => method.unwrap_or(Method::GET),
+        Err(e) => {
+            eprintln!("failed to parse method: {}", e);
+            return;
+        },
+    };
+
+    let headers = if let Some(headers) = args.values_of("header") {
+        match headers.map(parse_header).collect::<Result<HeaderMap<_>>>() {
+            Ok(headers) => headers,
+            Err(e) => {
+                eprintln!("failed to parse header: {}", e);
+                return;
+            },
+        }
+    } else {
+        HeaderMap::new()
+    };
+
+    let body: &str = args.value_of("body").unwrap_or_default();
+    let body = Bytes::copy_from_slice(body.as_bytes());
+
     let settings = bench::BenchmarkSettings {
         threads,
         connections: conns,
@@ -86,6 +118,9 @@ fn main() {
         display_percentile: pct,
         display_json: json,
         rounds,
+        method,
+        headers,
+        body,
     };
 
     bench::start_benchmark(settings);
@@ -137,6 +172,15 @@ fn parse_duration(duration: &str) -> Result<Duration> {
     Ok(dur)
 }
 
+fn parse_header(value: &str) -> Result<(HeaderName, HeaderValue)> {
+    let (key, value) = value
+        .split_once(": ")
+        .context("Header value missing colon (\": \")")?;
+    let key = HeaderName::from_str(key).context("Invalid header name")?;
+    let value = HeaderValue::from_str(value).context("Invalid header value")?;
+    Ok((key, value))
+}
+
 /// Contains Clap's app setup.
 fn parse_args() -> ArgMatches<'static> {
     App::new("ReWrk")
@@ -204,6 +248,32 @@ fn parse_args() -> ArgMatches<'static> {
                 .takes_value(true)
                 .required(false),
         )
+        .arg(
+            Arg::with_name("method")
+                .long("method")
+                .short("m")
+                .help("Set request method e.g. '-m get'")
+                .takes_value(true)
+                .required(false)
+                .multiple(true),
+        )
+        .arg(
+            Arg::with_name("header")
+                .long("header")
+                .short("H")
+                .help("Add header to request e.g. '-H \"content-type: text/plain\"'")
+                .takes_value(true)
+                .required(false)
+                .multiple(true),
+        )
+        .arg(
+            Arg::with_name("body")
+                .long("body")
+                .short("b")
+                .help("Add body to request e.g. '-b \"foo\"'")
+                .takes_value(true)
+                .required(false),
+        )
         //.arg(
         //    Arg::with_name("random")
         //        .long("rand")