From 5f9fac533a502f1ea1b2c7f7fe44c7b114f57016 Mon Sep 17 00:00:00 2001 From: Thomas de Grenier de Latour Date: Tue, 11 Feb 2025 00:41:24 +0100 Subject: [PATCH 1/3] Randomize --connect-to when used several times for the same host:port --- README.md | 1 + src/client.rs | 7 +++++-- src/main.rs | 3 ++- 3 files changed, 8 insertions(+), 3 deletions(-) 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, From 328c7d489e441ae599bd24aa95861e9afd65244d Mon Sep 17 00:00:00 2001 From: Thomas de Grenier de Latour Date: Tue, 11 Feb 2025 00:42:01 +0100 Subject: [PATCH 2/3] Add a test for the randomized --connect-to use-case --- tests/tests.rs | 63 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/tests/tests.rs b/tests/tests.rs index 33dcff62..6c63eec0 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(); + "Success1" + }), + ); + + 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 let Ok(_) = rx1.try_recv() { + count1 += 1; + } else if let Ok(_) = rx2.try_recv() { + 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!( From 1bb9e1b7f792f9f2c694d935f91879eb7e3be0fc Mon Sep 17 00:00:00 2001 From: Thomas de Grenier de Latour Date: Tue, 11 Feb 2025 09:22:15 +0100 Subject: [PATCH 3/3] Fix clippy warning --- tests/tests.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/tests.rs b/tests/tests.rs index 6c63eec0..eb443985 100644 --- a/tests/tests.rs +++ b/tests/tests.rs @@ -319,7 +319,7 @@ async fn distribution_on_two_matching_connect_to(host: &'static str) -> (i32, i3 "/", get(move || async move { tx2.send(()).unwrap(); - "Success1" + "Success2" }), ); @@ -347,9 +347,9 @@ async fn distribution_on_two_matching_connect_to(host: &'static str) -> (i32, i3 let mut count1 = 0; let mut count2 = 0; loop { - if let Ok(_) = rx1.try_recv() { + if rx1.try_recv().is_ok() { count1 += 1; - } else if let Ok(_) = rx2.try_recv() { + } else if rx2.try_recv().is_ok() { count2 += 1; } else { break;