nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: hatoo/oha
Pull request #182 (2022-08-21 06:39:59)

Support Ipv6 host
#181 


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1694b28..8af8266 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 # Unreleased
 
+- Support Ipv6 host #181
 - Print min, max, average and pXX for Requests per second in JSON output like bombardier #177
 - Add JSON Output #169
 - Fix QPS control to send with correct rate for first 1 sec
diff --git a/src/client.rs b/src/client.rs
index b8b16db..faece49 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -67,6 +67,14 @@ impl DNS {
             (host, port)
         };
 
+        let host = if host.starts_with('[') && host.ends_with(']') {
+            // host is [ipv6] format
+            // remove first [ and last ]
+            &host[1..host.len() - 1]
+        } else {
+            host
+        };
+
         // Perform actual DNS lookup, either on the original (host, port), or
         // on the (host, port) specified with `--connect-to`.
         let addrs = self
diff --git a/src/main.rs b/src/main.rs
index 09c5926..fe2c054 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -140,7 +140,7 @@ impl FromStr for ConnectToEntry {
         let (target_host, target_port) = if s.starts_with('[') {
             // parse bracketed IPv6 address
             let end = s.find("]:").ok_or(expected_syntax)?;
-            let target_host = &s[1..end];
+            let target_host = &s[..end + 1];
             let target_port = &s[end + 2..]; // shouldn't panic because we looked for `]:`, not just `]`
             (target_host, target_port)
         } else {


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 58c6a6d..dd9f350 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -386,3 +386,37 @@ async fn test_redirect() {
         assert!(!redirect(n, false, 10).await);
     }
 }
+
+#[tokio::test]
+async fn test_ipv6() {
+    let (tx, rx) = flume::unbounded();
+    let report_host =
+        warp::get()
+            .and(warp::filters::header::header("host"))
+            .map(move |host: String| {
+                tx.send(host).unwrap();
+                "Hello World"
+            });
+
+    let _guard = PORT_LOCK.lock().await;
+    // sic. the `get_port` crate doesn't support IpV6 addresses, so we check
+    // with 127.0.0.1 even though we bind on ::1 later.
+    let port = get_port::tcp::TcpPort::any("127.0.0.1").unwrap();
+    let addr = Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1);
+    tokio::spawn(warp::serve(report_host).run((addr, 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"])
+            .arg(format!("http://[::1]:{}/", port))
+            .assert()
+            .success();
+    })
+    .await
+    .unwrap();
+
+    rx.try_recv().unwrap();
+}