diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index ee354588..3a744392 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -10,6 +10,7 @@ jobs: matrix: os: [ubuntu-latest, windows-latest, macos-latest] rust: [stable, beta, nightly] + additional_args: ["", "--no-default-features --features rustls"] steps: - uses: actions/checkout@v2 - uses: actions-rs/toolchain@v1 @@ -20,6 +21,7 @@ jobs: - uses: actions-rs/cargo@v1 with: command: check + args: ${{ matrix.additional_args }} test: name: Test Suite @@ -28,6 +30,7 @@ jobs: matrix: os: [ubuntu-latest, windows-latest, macos-latest] rust: [stable, beta, nightly] + additional_args: ["", "--no-default-features --features rustls"] steps: - uses: actions/checkout@v2 - uses: actions-rs/toolchain@v1 @@ -38,6 +41,7 @@ jobs: - uses: actions-rs/cargo@v1 with: command: test + args: ${{ matrix.additional_args }} fmt: name: Rustfmt diff --git a/CHANGELOG.md b/CHANGELOG.md index 1217e8bd..cf2562e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,7 @@ # Unreleased +- Support randomly generated URL using rand_regex crate + # 0.5.9 (2023-06-12) - Fix -H Header parser diff --git a/Cargo.lock b/Cargo.lock index 0fdd2cc7..cc494eed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1015,7 +1015,9 @@ dependencies = [ "libc", "native-tls", "rand", + "rand_regex", "ratatui", + "regex-syntax", "rlimit", "rustls", "rustls-native-certs", @@ -1234,6 +1236,16 @@ dependencies = [ "getrandom", ] +[[package]] +name = "rand_regex" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a9fe2d7d9eeaf3279d1780452a5bbd26b31b27938787ef1c3e930d1e9cfbd" +dependencies = [ + "rand", + "regex-syntax", +] + [[package]] name = "ratatui" version = "0.21.0" @@ -1262,6 +1274,12 @@ version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +[[package]] +name = "regex-syntax" +version = "0.6.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" + [[package]] name = "resolv-conf" version = "0.7.0" diff --git a/Cargo.toml b/Cargo.toml index f3932808..543b311e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,6 +50,8 @@ base64 = "0.21.0" lazy_static = "1.4.0" rand = "0.8" trust-dns-resolver = "0.22.0" +rand_regex = "0.15.1" +regex-syntax = "0.6.22" [target.'cfg(unix)'.dependencies] rlimit = "0.9.0" diff --git a/README.md b/README.md index d62229b6..81a0fb0f 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,8 @@ Options: -z Duration of application to send requests. If duration is specified, n is ignored. Examples: -z 10s -z 3m. -q Rate limit for all, in queries per second (QPS) + --rand-regex-url Generate URL by rand_regex crate but dot is disabled for each query e.g. http://127.0.0.1/[a-z][a-z][0-9]. See https://docs.rs/rand_regex/latest/rand_regex/struct.Regex.html for details of syntax. + --max-repeat A parameter for the '--rand-regex-url'. The max_repeat parameter gives the maximum extra repeat counts the x*, x+ and x{n,} operators will become. [default: 4] --latency-correction Correct latency to avoid coordinated omission problem. It's ignored if -q is not set. --no-tui No realtime tui -j, --json Print results as JSON @@ -107,6 +109,20 @@ oha <-z or -n> -c -q --la You can avoid `Coordinated Omission Problem` by using `--latency-correction`. +## Dynamic url feature + +You can use `--rand-regex-url` option to generate random url for each connection. + +```sh +oha --rand-regex-url http://127.0.0.1/[a-z][a-z][0-9] +``` + +Each Urls are generated by [rand_regex](https://github.com/kennytm/rand_regex) crate but regex's dot is disabled since it's not useful for this purpose and it's very incovenient if url's dots are interpreted as regex's dot. + +Optionaly you can set `--max-repeat` option to limit max repeat count for each regex. e.g http://127.0.0.1/[a-z]* with `--max-repeat 4` will generate url like http://127.0.0.1/[a-z]{0,4} + +If keep-alive is enabled, path and query are generated for each request but shceme, host and port are fixed during connection. + # Contribution Feel free to help us! diff --git a/src/client.rs b/src/client.rs index 30207f65..99f7bd9a 100644 --- a/src/client.rs +++ b/src/client.rs @@ -4,6 +4,7 @@ use rand::prelude::*; use std::sync::Arc; use thiserror::Error; +use crate::url_generator::{UrlGenerator, UrlGeneratorError}; use crate::ConnectToEntry; #[derive(Debug, Clone)] @@ -39,8 +40,6 @@ impl RequestResult { #[allow(clippy::upper_case_acronyms)] struct DNS { - // To pick a random address from DNS. - rng: rand::rngs::StdRng, connect_to: Arc>, resolver: Arc< trust_dns_resolver::AsyncResolver< @@ -54,7 +53,11 @@ struct DNS { impl DNS { /// Perform a DNS lookup for a given url and returns (ip_addr, port) - async fn lookup(&mut self, url: &http::Uri) -> Result<(std::net::IpAddr, u16), ClientError> { + async fn lookup( + &self, + url: &http::Uri, + rng: &mut R, + ) -> Result<(std::net::IpAddr, u16), ClientError> { let host = url.host().ok_or(ClientError::HostNotFound)?; let port = get_http_port(url).ok_or(ClientError::PortNotFound)?; @@ -87,9 +90,7 @@ impl DNS { .iter() .collect::>(); - let addr = *addrs - .choose(&mut self.rng) - .ok_or(ClientError::DNSNoRecord)?; + let addr = *addrs.choose(rng).ok_or(ClientError::DNSNoRecord)?; Ok((addr, port)) } @@ -97,7 +98,7 @@ impl DNS { pub struct ClientBuilder { pub http_version: http::Version, - pub url: http::Uri, + pub url_generator: UrlGenerator, pub method: http::Method, pub headers: http::header::HeaderMap, pub body: Option<&'static [u8]>, @@ -122,14 +123,13 @@ pub struct ClientBuilder { impl ClientBuilder { pub fn build(&self) -> Client { Client { - url: self.url.clone(), + url_generator: self.url_generator.clone(), method: self.method.clone(), headers: self.headers.clone(), body: self.body, dns: DNS { resolver: self.resolver.clone(), connect_to: self.connect_to.clone(), - rng: rand::rngs::StdRng::from_entropy(), }, client: None, timeout: self.timeout, @@ -191,11 +191,13 @@ pub enum ClientError { InvalidUri(#[from] http::uri::InvalidUri), #[error("timeout")] Timeout, + #[error(transparent)] + UrlGeneratorError(#[from] UrlGeneratorError), } pub struct Client { http_version: http::Version, - url: http::Uri, + url_generator: UrlGenerator, method: http::Method, headers: http::header::HeaderMap, body: Option<&'static [u8]>, @@ -212,11 +214,12 @@ pub struct Client { impl Client { #[cfg(unix)] async fn client( - &mut self, + &self, addr: (std::net::IpAddr, u16), + url: &http::Uri, ) -> Result, ClientError> { - if self.url.scheme() == Some(&http::uri::Scheme::HTTPS) { - self.tls_client(addr).await + if url.scheme() == Some(&http::uri::Scheme::HTTPS) { + self.tls_client(addr, url).await } else if let Some(socket_path) = &self.unix_socket { let stream = tokio::net::UnixStream::connect(socket_path).await?; let (send, conn) = hyper::client::conn::handshake(stream).await?; @@ -234,11 +237,12 @@ impl Client { #[cfg(not(unix))] async fn client( - &mut self, + &self, addr: (std::net::IpAddr, u16), + url: &http::Uri, ) -> Result, ClientError> { - if self.url.scheme() == Some(&http::uri::Scheme::HTTPS) { - self.tls_client(addr).await + if url.scheme() == Some(&http::uri::Scheme::HTTPS) { + self.tls_client(addr, url).await } else { let stream = tokio::net::TcpStream::connect(addr).await?; stream.set_nodelay(true)?; @@ -251,8 +255,9 @@ impl Client { #[cfg(all(feature = "native-tls", not(feature = "rustls")))] async fn tls_client( - &mut self, + &self, addr: (std::net::IpAddr, u16), + url: &http::Uri, ) -> Result, ClientError> { let stream = tokio::net::TcpStream::connect(addr).await?; stream.set_nodelay(true)?; @@ -267,7 +272,7 @@ impl Client { }; let connector = tokio_native_tls::TlsConnector::from(connector); let stream = connector - .connect(self.url.host().ok_or(ClientError::HostNotFound)?, stream) + .connect(url.host().ok_or(ClientError::HostNotFound)?, stream) .await?; let (send, conn) = hyper::client::conn::handshake(stream).await?; @@ -277,8 +282,9 @@ impl Client { #[cfg(feature = "rustls")] async fn tls_client( - &mut self, + &self, addr: (std::net::IpAddr, u16), + url: &http::Uri, ) -> Result, ClientError> { let stream = tokio::net::TcpStream::connect(addr).await?; stream.set_nodelay(true)?; @@ -297,8 +303,7 @@ impl Client { .set_certificate_verifier(Arc::new(AcceptAnyServerCert)); } let connector = tokio_rustls::TlsConnector::from(Arc::new(config)); - let domain = - rustls::ServerName::try_from(self.url.host().ok_or(ClientError::HostNotFound)?)?; + let domain = rustls::ServerName::try_from(url.host().ok_or(ClientError::HostNotFound)?)?; let stream = connector.connect(domain, stream).await?; let (send, conn) = hyper::client::conn::handshake(stream).await?; @@ -328,7 +333,7 @@ impl Client { } } - pub async fn work(&mut self) -> Result { + pub async fn work(&mut self, rng: &mut R) -> Result { let timeout = if let Some(timeout) = self.timeout { tokio::time::sleep(timeout).boxed() } else { @@ -336,15 +341,16 @@ impl Client { }; let do_req = async { + let url = self.url_generator.generate(rng)?; let mut start = std::time::Instant::now(); let mut connection_time: Option = None; let mut send_request = if let Some(send_request) = self.client.take() { send_request } else { - let addr = self.dns.lookup(&self.url).await?; + let addr = self.dns.lookup(&url, rng).await?; let dns_lookup = std::time::Instant::now(); - let send_request = self.client(addr).await?; + let send_request = self.client(addr, &url).await?; let dialup = std::time::Instant::now(); connection_time = Some(ConnectionTime { dns_lookup, dialup }); @@ -355,13 +361,13 @@ impl Client { .is_err() { start = std::time::Instant::now(); - let addr = self.dns.lookup(&self.url).await?; + let addr = self.dns.lookup(&url, rng).await?; let dns_lookup = std::time::Instant::now(); - send_request = self.client(addr).await?; + send_request = self.client(addr, &url).await?; let dialup = std::time::Instant::now(); connection_time = Some(ConnectionTime { dns_lookup, dialup }); } - let request = self.request(&self.url)?; + let request = self.request(&url)?; match send_request.send_request(request).await { Ok(res) => { let (parts, mut stream) = res.into_parts(); @@ -375,12 +381,7 @@ impl Client { if self.redirect_limit != 0 { if let Some(location) = parts.headers.get("Location") { let (send_request_redirect, new_status, len) = self - .redirect( - send_request, - &self.url.clone(), - location, - self.redirect_limit, - ) + .redirect(send_request, &url, location, self.redirect_limit, rng) .await?; send_request = send_request_redirect; @@ -424,12 +425,13 @@ impl Client { } #[allow(clippy::type_complexity)] - fn redirect<'a>( - &'a mut self, + fn redirect<'a, R: Rng + Send>( + &'a self, send_request: hyper::client::conn::SendRequest, base_url: &'a http::Uri, location: &'a http::header::HeaderValue, limit: usize, + rng: &'a mut R, ) -> futures::future::BoxFuture< 'a, Result< @@ -460,20 +462,20 @@ impl Client { // reuse connection (send_request, None) } else { - let addr = self.dns.lookup(&url).await?; - (self.client(addr).await?, Some(send_request)) + let addr = self.dns.lookup(&url, rng).await?; + (self.client(addr, &url).await?, Some(send_request)) }; while futures::future::poll_fn(|ctx| send_request.poll_ready(ctx)) .await .is_err() { - let addr = self.dns.lookup(&url).await?; - send_request = self.client(addr).await?; + let addr = self.dns.lookup(&url, rng).await?; + send_request = self.client(addr, &url).await?; } let mut request = self.request(&url)?; - if url.authority() != self.url.authority() { + if url.authority() != base_url.authority() { request.headers_mut().insert( http::header::HOST, http::HeaderValue::from_str( @@ -494,7 +496,7 @@ impl Client { if let Some(location) = parts.headers.get("Location") { let (send_request_redirect, new_status, len) = self - .redirect(send_request, &url, location, limit - 1) + .redirect(send_request, &url, location, limit - 1, rng) .await?; send_request = send_request_redirect; status = new_status; @@ -584,11 +586,12 @@ pub async fn work( let futures = (0..n_workers) .map(|_| { let mut w = client_builder.build(); + let mut rng = StdRng::from_entropy(); let report_tx = report_tx.clone(); let counter = counter.clone(); tokio::spawn(async move { while counter.fetch_add(1, Ordering::Relaxed) < n_tasks { - let res = w.work().await; + let res = w.work(&mut rng).await; let is_cancel = is_too_many_open_files(&res); report_tx.send_async(res).await.unwrap(); if is_cancel { @@ -629,11 +632,12 @@ pub async fn work_with_qps( let futures = (0..n_workers) .map(|_| { let mut w = client_builder.build(); + let mut rng = StdRng::from_entropy(); let report_tx = report_tx.clone(); let rx = rx.clone(); tokio::spawn(async move { while let Ok(()) = rx.recv_async().await { - let res = w.work().await; + let res = w.work(&mut rng).await; let is_cancel = is_too_many_open_files(&res); report_tx.send_async(res).await.unwrap(); if is_cancel { @@ -674,11 +678,12 @@ pub async fn work_with_qps_latency_correction( let futures = (0..n_workers) .map(|_| { let mut w = client_builder.build(); + let mut rng = StdRng::from_entropy(); let report_tx = report_tx.clone(); let rx = rx.clone(); tokio::spawn(async move { while let Ok(start) = rx.recv_async().await { - let mut res = w.work().await; + let mut res = w.work(&mut rng).await; if let Ok(request_result) = &mut res { request_result.start_latency_correction = Some(start); @@ -710,9 +715,10 @@ pub async fn work_until( .map(|_| { let report_tx = report_tx.clone(); let mut w = client_builder.build(); + let mut rng = StdRng::from_entropy(); tokio::spawn(async move { loop { - let res = w.work().await; + let res = w.work(&mut rng).await; let is_cancel = is_too_many_open_files(&res); report_tx.send_async(res).await.unwrap(); if is_cancel { @@ -759,11 +765,12 @@ pub async fn work_until_with_qps( let futures = (0..n_workers) .map(|_| { let mut w = client_builder.build(); + let mut rng = StdRng::from_entropy(); let report_tx = report_tx.clone(); let rx = rx.clone(); tokio::spawn(async move { while let Ok(()) = rx.recv_async().await { - let res = w.work().await; + let res = w.work(&mut rng).await; let is_cancel = is_too_many_open_files(&res); report_tx.send_async(res).await.unwrap(); if is_cancel { @@ -811,11 +818,12 @@ pub async fn work_until_with_qps_latency_correction( let futures = (0..n_workers) .map(|_| { let mut w = client_builder.build(); + let mut rng = StdRng::from_entropy(); let report_tx = report_tx.clone(); let rx = rx.clone(); tokio::spawn(async move { while let Ok(start) = rx.recv_async().await { - let mut res = w.work().await; + let mut res = w.work(&mut rng).await; if let Ok(request_result) = &mut res { request_result.start_latency_correction = Some(start); diff --git a/src/main.rs b/src/main.rs index 73a908ef..71e3c166 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,15 +3,20 @@ use clap::Parser; use crossterm::tty::IsTty; use futures::prelude::*; use http::header::{HeaderName, HeaderValue}; +use http::Uri; use printer::PrintMode; +use rand::prelude::*; +use rand_regex::Regex; use std::sync::Arc; use std::{io::Read, str::FromStr}; +use url_generator::UrlGenerator; mod client; mod histogram; mod monitor; mod printer; mod timescale; +mod url_generator; #[cfg(all(target_env = "musl", target_pointer_width = "64"))] #[global_allocator] @@ -23,7 +28,7 @@ use client::{ClientError, RequestResult}; #[clap(author, about, version, override_usage = "oha [FLAGS] [OPTIONS] ")] struct Opts { #[clap(help = "Target URL.")] - url: http::Uri, + url: String, #[structopt( help = "Number of requests to run.", short = 'n', @@ -44,6 +49,18 @@ Examples: -z 10s -z 3m.", duration: Option, #[clap(help = "Rate limit for all, in queries per second (QPS)", short = 'q')] query_per_second: Option, + #[clap( + help = "Generate URL by rand_regex crate but dot is disabled for each query e.g. http://127.0.0.1/[a-z][a-z][0-9]. See https://docs.rs/rand_regex/latest/rand_regex/struct.Regex.html for details of syntax.", + default_value = "false", + long + )] + rand_regex_url: bool, + #[clap( + help = "A parameter for the '--rand-regex-url'. The max_repeat parameter gives the maximum extra repeat counts the x*, x+ and x{n,} operators will become.", + default_value = "4", + long + )] + max_repeat: u32, #[clap( help = "Correct latency to avoid coordinated omission problem. It's ignored if -q is not set.", long = "latency-correction" @@ -179,6 +196,26 @@ async fn main() -> anyhow::Result<()> { http::Version::HTTP_11 }; + let url_generator = if opts.rand_regex_url { + // Almost URL has dot in domain, so disable dot in regex for convenience. + let dot_disabled: String = opts + .url + .chars() + .map(|c| { + if c == '.' { + regex_syntax::escape(".") + } else { + c.to_string() + } + }) + .collect(); + UrlGenerator::new_dynamic(Regex::compile(&dot_disabled, opts.max_repeat)?) + } else { + UrlGenerator::new_static(Uri::from_str(&opts.url)?) + }; + + let url = url_generator.generate(&mut thread_rng())?; + let headers = { let mut headers: http::header::HeaderMap = Default::default(); @@ -199,7 +236,7 @@ async fn main() -> anyhow::Result<()> { headers.insert( http::header::HOST, http::header::HeaderValue::from_str( - opts.url.authority().context("get authority")?.as_str(), + url.authority().context("get authority")?.as_str(), )?, ); @@ -366,7 +403,7 @@ async fn main() -> anyhow::Result<()> { // client_builder builds client for each workers let client_builder = client::ClientBuilder { http_version, - url: opts.url, + url_generator, method: opts.method, headers, body, diff --git a/src/url_generator.rs b/src/url_generator.rs new file mode 100644 index 00000000..0eb38cfe --- /dev/null +++ b/src/url_generator.rs @@ -0,0 +1,42 @@ +use std::{borrow::Cow, str::FromStr, string::FromUtf8Error}; + +use http::{uri::InvalidUri, Uri}; +use rand::prelude::*; +use rand_regex::Regex; +use thiserror::Error; + +#[derive(Clone, Debug)] +pub enum UrlGenerator { + Static(Uri), + Dynamic(Regex), +} + +#[derive(Error, Debug)] +pub enum UrlGeneratorError { + #[error("{0}, generated url: {1}")] + InvalidUri(InvalidUri, String), + #[error(transparent)] + FromUtf8Error(#[from] FromUtf8Error), +} + +impl UrlGenerator { + pub fn new_static(url: Uri) -> Self { + Self::Static(url) + } + + pub fn new_dynamic(regex: Regex) -> Self { + Self::Dynamic(regex) + } + + pub fn generate(&self, rng: &mut R) -> Result, UrlGeneratorError> { + match self { + Self::Static(url) => Ok(Cow::Borrowed(url)), + Self::Dynamic(regex) => { + let generated = Distribution::>::sample(regex, rng)?; + Ok(Cow::Owned(Uri::from_str(generated.as_str()).map_err( + |e| UrlGeneratorError::InvalidUri(e, generated), + )?)) + } + } + } +} diff --git a/tests/tests.rs b/tests/tests.rs index 29d00005..addaceca 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -101,6 +101,33 @@ async fn get_query(p: &'static str) -> String { rx.try_recv().unwrap() } +async fn get_path_rand_regex(p: &'static str) -> String { + let (tx, rx) = flume::unbounded(); + let report_headers = warp::path!(String).map(move |path: String| { + tx.send(path).unwrap(); + "Hello World" + }); + + let _guard = PORT_LOCK.lock().await; + let port = get_port::tcp::TcpPort::any("127.0.0.1").unwrap(); + tokio::spawn(warp::serve(report_headers).run(([127, 0, 0, 1], port))); + // It's not guaranteed that the port is used here. + // So we can't drop guard here. + + tokio::task::spawn_blocking(move || { + Command::cargo_bin("oha") + .unwrap() + .args(["-n", "1", "--no-tui", "--rand-regex-url"]) + .arg(format!(r"http://127.0.0.1:{port}/{p}")) + .assert() + .success(); + }) + .await + .unwrap(); + + rx.try_recv().unwrap() +} + async fn redirect(n: usize, is_relative: bool, limit: usize) -> bool { let (tx, rx) = flume::unbounded(); let _guard = PORT_LOCK.lock().await; @@ -387,6 +414,16 @@ async fn test_query() { ); } +#[tokio::test] +async fn test_query_rand_regex() { + let query = get_path_rand_regex("[a-z][0-9][a-z]").await; + let chars = query.chars().collect::>(); + assert_eq!(chars.len(), 3); + assert!(chars[0].is_ascii_lowercase()); + assert!(chars[1].is_ascii_digit()); + assert!(chars[2].is_ascii_lowercase()); +} + #[tokio::test] async fn test_connect_to() { assert_eq!(