diff --git a/README.md b/README.md index b4580369..1105d41c 100644 --- a/README.md +++ b/README.md @@ -192,6 +192,7 @@ Options: Accept invalid certs. --connect-to Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443' + Note: if used several times for the same host:port:target_host:target_port, a random choice is made --disable-color Disable the color scheme. --unix-socket diff --git a/src/client.rs b/src/client.rs index bba17310..4f5c42c2 100644 --- a/src/client.rs +++ b/src/client.rs @@ -79,11 +79,14 @@ impl Dns { .port_or_known_default() .ok_or(ClientError::PortNotFound)?; - // Try to find an override (passed via `--connect-to`) that applies to this (host, port) + // Try to find an override (passed via `--connect-to`) that applies to this (host, port), + // choosing one randomly if several match. let (host, port) = if let Some(entry) = self .connect_to .iter() - .find(|entry| entry.requested_port == port && entry.requested_host == host) + .filter(|entry| entry.requested_port == port && entry.requested_host == host) + .collect::>() + .choose(rng) { (entry.target_host.as_str(), entry.target_port) } else { diff --git a/src/main.rs b/src/main.rs index 8f0003d4..eb05b005 100644 --- a/src/main.rs +++ b/src/main.rs @@ -236,7 +236,8 @@ Note: If qps is specified, burst will be ignored", #[arg(help = "Accept invalid certs.", long = "insecure")] insecure: bool, #[arg( - help = "Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443'", + help = "Override DNS resolution and default port numbers with strings like 'example.org:443:localhost:8443' +Note: if used several times for the same host:port:target_host:target_port, a random choice is made", long = "connect-to" )] connect_to: Vec, diff --git a/tests/tests.rs b/tests/tests.rs index 33dcff62..eb443985 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -302,6 +302,62 @@ async fn burst_10_req_delay_2s_rate_4(iteration: u8, args: &[&str]) -> usize { count } +// Randomly spread 100 requests on two matching --connect-to targets, and return a count for each +async fn distribution_on_two_matching_connect_to(host: &'static str) -> (i32, i32) { + let (tx1, rx1) = flume::unbounded(); + let (tx2, rx2) = flume::unbounded(); + + let app1 = Router::new().route( + "/", + get(move || async move { + tx1.send(()).unwrap(); + "Success1" + }), + ); + + let app2 = Router::new().route( + "/", + get(move || async move { + tx2.send(()).unwrap(); + "Success2" + }), + ); + + let (listener1, port1) = bind_port().await; + tokio::spawn(async { axum::serve(listener1, app1).await }); + + let (listener2, port2) = bind_port().await; + tokio::spawn(async { axum::serve(listener2, app2).await }); + + tokio::task::spawn_blocking(move || { + Command::cargo_bin("oha") + .unwrap() + .args(["-n", "100", "--no-tui"]) + .arg(format!("http://{host}/")) + .arg("--connect-to") + .arg(format!("{host}:80:localhost:{port1}")) + .arg("--connect-to") + .arg(format!("{host}:80:localhost:{port2}")) + .assert() + .success(); + }) + .await + .unwrap(); + + let mut count1 = 0; + let mut count2 = 0; + loop { + if rx1.try_recv().is_ok() { + count1 += 1; + } else if rx2.try_recv().is_ok() { + count2 += 1; + } else { + break; + } + } + (count1, count2) +} + #[tokio::test] async fn test_enable_compression_default() { let req = get_req("/", &[]).await; @@ -596,6 +652,13 @@ async fn test_connect_to() { ) } +#[tokio::test] +async fn test_connect_to_randomness() { + let (count1, count2) = distribution_on_two_matching_connect_to("invalid.example.org").await; + assert!(count1 >= 10 && count2 >= 10); // should not be too flaky with 100 coin tosses + assert!(count1 + count2 == 100); +} + #[tokio::test] async fn test_connect_to_ipv6_target() { assert_eq!(