From 7bbb7c9d9b5efd9ccc4526df2f0f98f5e943938d Mon Sep 17 00:00:00 2001 From: Siavash Sameni Date: Tue, 31 Mar 2026 18:28:48 +0400 Subject: [PATCH] Add dual-stack IPv4+IPv6 listening Server now binds on both IPv4 (0.0.0.0) and IPv6 (::) by default. Uses tokio::select! to accept from whichever listener has a connection. New flags: --listen IPv4 listen address (default: 0.0.0.0, "none" to disable) --listen6 IPv6 listen address (default: ::, "none" to disable) Examples: btest -s # listen on both v4 and v6 btest -s --listen6 none # IPv4 only btest -s --listen none # IPv6 only btest -s --listen 192.168.1.1 # specific IPv4 address Co-Authored-By: Claude Opus 4.6 (1M context) --- src/main.rs | 12 +++++++- src/server.rs | 60 +++++++++++++++++++++++++++++++++++---- tests/ecsrp5_test.rs | 10 +++++-- tests/integration_test.rs | 2 +- 4 files changed, 73 insertions(+), 11 deletions(-) diff --git a/src/main.rs b/src/main.rs index 6b64cf9..2633640 100644 --- a/src/main.rs +++ b/src/main.rs @@ -50,6 +50,14 @@ struct Cli { #[arg(short = 'P', long = "port", default_value_t = BTEST_PORT)] port: u16, + /// Listen address for IPv4 (default: 0.0.0.0, use "none" to disable) + #[arg(long = "listen", default_value = "0.0.0.0")] + listen_addr: String, + + /// Listen address for IPv6 (default: ::, use "none" to disable) + #[arg(long = "listen6", default_value = "::")] + listen6_addr: String, + /// Authentication username #[arg(short = 'a', long = "authuser")] auth_user: Option, @@ -101,8 +109,10 @@ async fn main() -> anyhow::Result<()> { if cli.server { // Server mode + let v4 = if cli.listen_addr.eq_ignore_ascii_case("none") { None } else { Some(cli.listen_addr) }; + let v6 = if cli.listen6_addr.eq_ignore_ascii_case("none") { None } else { Some(cli.listen6_addr) }; tracing::info!("Starting btest server on port {}", cli.port); - server::run_server(cli.port, cli.auth_user, cli.auth_pass, cli.ecsrp5).await?; + server::run_server(cli.port, cli.auth_user, cli.auth_pass, cli.ecsrp5, v4, v6).await?; } else if let Some(host) = cli.client { // Client mode - must specify at least one direction if !cli.transmit && !cli.receive { diff --git a/src/server.rs b/src/server.rs index 851fefa..be38651 100644 --- a/src/server.rs +++ b/src/server.rs @@ -27,10 +27,9 @@ pub async fn run_server( auth_user: Option, auth_pass: Option, use_ecsrp5: bool, + listen_v4: Option, + listen_v6: Option, ) -> Result<()> { - let addr = format!("0.0.0.0:{}", port); - let listener = TcpListener::bind(&addr).await?; - // Pre-derive EC-SRP5 credentials if enabled let ecsrp5_creds = if use_ecsrp5 { match (auth_user.as_deref(), auth_pass.as_deref()) { @@ -47,13 +46,62 @@ pub async fn run_server( None }; - tracing::info!("btest server listening on {}", addr); - let udp_port_offset = Arc::new(std::sync::atomic::AtomicU16::new(0)); let sessions: SessionMap = Arc::new(Mutex::new(HashMap::new())); + // Bind IPv4 listener + let v4_listener = if let Some(ref addr) = listen_v4 { + let bind_addr = format!("{}:{}", addr, port); + match TcpListener::bind(&bind_addr).await { + Ok(l) => { + tracing::info!("Listening on {} (IPv4)", bind_addr); + Some(l) + } + Err(e) => { + tracing::error!("Failed to bind {}: {}", bind_addr, e); + None + } + } + } else { + None + }; + + // Bind IPv6 listener + let v6_listener = if let Some(ref addr) = listen_v6 { + let bind_addr = format!("[{}]:{}", addr, port); + match TcpListener::bind(&bind_addr).await { + Ok(l) => { + tracing::info!("Listening on {} (IPv6)", bind_addr); + Some(l) + } + Err(e) => { + tracing::error!("Failed to bind {}: {}", bind_addr, e); + None + } + } + } else { + None + }; + + if v4_listener.is_none() && v6_listener.is_none() { + return Err(crate::protocol::BtestError::Protocol( + "No listeners bound. Check --listen and --listen6 addresses.".into(), + )); + } + loop { - let (stream, peer) = listener.accept().await?; + // Accept from whichever listener has a connection ready + let (stream, peer) = match (&v4_listener, &v6_listener) { + (Some(v4), Some(v6)) => { + tokio::select! { + r = v4.accept() => r?, + r = v6.accept() => r?, + } + } + (Some(v4), None) => v4.accept().await?, + (None, Some(v6)) => v6.accept().await?, + (None, None) => unreachable!(), + }; tracing::info!("New connection from {}", peer); let auth_user = auth_user.clone(); diff --git a/tests/ecsrp5_test.rs b/tests/ecsrp5_test.rs index 98246c0..21efa03 100644 --- a/tests/ecsrp5_test.rs +++ b/tests/ecsrp5_test.rs @@ -10,7 +10,9 @@ async fn start_ecsrp5_server(port: u16) { port, Some("testuser".into()), Some("testpass".into()), - true, // ecsrp5 + true, + Some("127.0.0.1".into()), + None, ) .await; }); @@ -23,7 +25,9 @@ async fn start_md5_server(port: u16) { port, Some("testuser".into()), Some("testpass".into()), - false, // md5 + false, + Some("127.0.0.1".into()), + None, ) .await; }); @@ -32,7 +36,7 @@ async fn start_md5_server(port: u16) { async fn start_noauth_server(port: u16) { tokio::spawn(async move { - let _ = btest_rs::server::run_server(port, None, None, false).await; + let _ = btest_rs::server::run_server(port, None, None, false, Some("127.0.0.1".into()), None).await; }); tokio::time::sleep(Duration::from_millis(200)).await; } diff --git a/tests/integration_test.rs b/tests/integration_test.rs index a04599c..5c397f2 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -8,7 +8,7 @@ async fn start_test_server(port: u16, auth_user: Option<&str>, auth_pass: Option let user = auth_user.map(String::from); let pass = auth_pass.map(String::from); tokio::spawn(async move { - let _ = btest_rs::server::run_server(port, user, pass, false).await; + let _ = btest_rs::server::run_server(port, user, pass, false, Some("127.0.0.1".into()), None).await; }); tokio::time::sleep(Duration::from_millis(100)).await; }