From 9c20f0c5998ac44b253ce459398aa03e25cd803b Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Tue, 26 Jul 2022 01:28:33 +0200 Subject: [PATCH] Support setting request method, headers, and body (#32) * support setting headers * support setting body * support setting method --- src/bench.rs | 14 ++++++++ src/http/mod.rs | 14 ++++++-- src/http/user_input.rs | 32 ++++++++++++++++--- src/main.rs | 72 +++++++++++++++++++++++++++++++++++++++++- 4 files changed, 123 insertions(+), 9 deletions(-) diff --git a/src/bench.rs b/src/bench.rs index 21af0ed..c153e5b 100644 --- a/src/bench.rs +++ b/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; diff --git a/src/http/mod.rs b/src/http/mod.rs index a12ff20..41da9be 100644 --- a/src/http/mod.rs +++ b/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> { 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,6 +100,8 @@ 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(); @@ -102,7 +109,8 @@ async fn benchmark( // 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(); diff --git a/src/http/user_input.rs b/src/http/user_input.rs index 97f6b8f..748740f 100644 --- a/src/http/user_input.rs +++ b/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 { - 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 { + spawn_blocking(move || { + Self::blocking_new(protocol, string, method, headers, body) + }) + .await + .unwrap() } - fn blocking_new(protocol: BenchType, string: String) -> Result { + fn blocking_new( + protocol: BenchType, + string: String, + method: Method, + headers: HeaderMap, + body: Bytes, + ) -> Result { let uri = Uri::try_from(string)?; let scheme = uri .scheme() @@ -91,6 +110,9 @@ impl UserInput { host, host_header, uri, + method, + headers, + body, }) } } diff --git a/src/main.rs b/src/main.rs index 95ed6bc..3b94514 100644 --- a/src/main.rs +++ b/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::() .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::>>() { + 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 { 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")