Initial commit: MikroTik btest server & client in Rust

Full reimplementation of the MikroTik Bandwidth Test protocol:
- Server mode: accepts connections from MikroTik devices on port 2000
- Client mode: connects to MikroTik btest servers
- TCP and UDP protocols with bidirectional support
- MD5 challenge-response authentication
- Dynamic speed adjustment (1.5x algorithm)
- Status exchange matching original C pselect() behavior
- Docker support with multi-stage build

Tested against MikroTik RouterOS achieving:
- 1.05 Gbps server RX (single connection)
- 530 Mbps client TCP download
- 840 Mbps client TCP upload
- 433 Mbps client UDP download

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Siavash Sameni
2026-03-31 11:56:34 +04:00
commit d9007dc169
18 changed files with 2801 additions and 0 deletions

168
src/auth.rs Normal file
View File

@@ -0,0 +1,168 @@
use md5::{Digest, Md5};
use rand::RngCore;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use crate::protocol::{self, BtestError, Result, AUTH_FAILED, AUTH_OK, AUTH_REQUIRED};
pub fn generate_challenge() -> [u8; 16] {
let mut nonce = [0u8; 16];
rand::thread_rng().fill_bytes(&mut nonce);
nonce
}
/// Compute the double-MD5 response: MD5(password + MD5(password + challenge))
pub fn compute_auth_hash(password: &str, challenge: &[u8; 16]) -> [u8; 16] {
// hash1 = MD5(password + challenge)
let mut hasher = Md5::new();
hasher.update(password.as_bytes());
hasher.update(challenge);
let hash1 = hasher.finalize();
// hash2 = MD5(password + hash1)
let mut hasher = Md5::new();
hasher.update(password.as_bytes());
hasher.update(&hash1);
hasher.finalize().into()
}
/// Server-side: send auth challenge and verify response.
/// Returns Ok(()) if auth succeeds or no auth is configured.
pub async fn server_authenticate<S: AsyncReadExt + AsyncWriteExt + Unpin>(
stream: &mut S,
username: Option<&str>,
password: Option<&str>,
) -> Result<()> {
match (username, password) {
(None, None) => {
// No auth required
stream.write_all(&AUTH_OK).await?;
stream.flush().await?;
Ok(())
}
(_, Some(pass)) => {
// Send auth challenge
stream.write_all(&AUTH_REQUIRED).await?;
let challenge = generate_challenge();
stream.write_all(&challenge).await?;
stream.flush().await?;
// Receive response: 16 bytes hash + 32 bytes username
let mut response = [0u8; 48];
stream.read_exact(&mut response).await?;
let received_hash = &response[0..16];
let received_user = &response[16..48];
// Extract username (null-terminated)
let user_end = received_user
.iter()
.position(|&b| b == 0)
.unwrap_or(32);
let received_username = std::str::from_utf8(&received_user[..user_end])
.unwrap_or("");
// Verify username if configured
if let Some(expected_user) = username {
if received_username != expected_user {
tracing::warn!("Auth failed: username mismatch (got '{}')", received_username);
stream.write_all(&AUTH_FAILED).await?;
stream.flush().await?;
return Err(BtestError::AuthFailed);
}
}
// Verify hash
let expected_hash = compute_auth_hash(pass, &challenge);
if received_hash != expected_hash {
tracing::warn!("Auth failed: hash mismatch for user '{}'", received_username);
stream.write_all(&AUTH_FAILED).await?;
stream.flush().await?;
return Err(BtestError::AuthFailed);
}
tracing::info!("Auth successful for user '{}'", received_username);
stream.write_all(&AUTH_OK).await?;
stream.flush().await?;
Ok(())
}
(Some(_), None) => {
// Username but no password - treat as no auth
stream.write_all(&AUTH_OK).await?;
stream.flush().await?;
Ok(())
}
}
}
/// Client-side: respond to auth challenge if required.
pub async fn client_authenticate<S: AsyncReadExt + AsyncWriteExt + Unpin>(
stream: &mut S,
resp: [u8; 4],
username: &str,
password: &str,
) -> Result<()> {
if resp == AUTH_OK {
return Ok(());
}
if resp == AUTH_REQUIRED {
// Read 16-byte challenge
let mut challenge = [0u8; 16];
stream.read_exact(&mut challenge).await?;
// Compute response
let hash = compute_auth_hash(password, &challenge);
// Build 48-byte response: 16 hash + 32 username
let mut response = [0u8; 48];
response[0..16].copy_from_slice(&hash);
let user_bytes = username.as_bytes();
let copy_len = user_bytes.len().min(32);
response[16..16 + copy_len].copy_from_slice(&user_bytes[..copy_len]);
stream.write_all(&response).await?;
stream.flush().await?;
// Read auth result
let result = protocol::recv_response(stream).await?;
if result == AUTH_OK {
tracing::info!("Authentication successful");
Ok(())
} else {
Err(BtestError::AuthFailed)
}
} else if resp == [0x03, 0x00, 0x00, 0x00] {
Err(BtestError::Protocol(
"EC-SRP5 authentication (RouterOS >= 6.43) is not supported".into(),
))
} else {
Err(BtestError::Protocol(format!(
"Unexpected auth response: {:02x?}",
resp
)))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_auth_hash_known_vector() {
// From the Perl reference: password "test", challenge as hex "ad32d6f94d28161625f2f390bb895637"
let challenge: [u8; 16] = [
0xad, 0x32, 0xd6, 0xf9, 0x4d, 0x28, 0x16, 0x16, 0x25, 0xf2, 0xf3, 0x90, 0xbb, 0x89,
0x56, 0x37,
];
let hash = compute_auth_hash("test", &challenge);
let hex: String = hash.iter().map(|b| format!("{:02x}", b)).collect();
assert_eq!(hex, "3c968565bc0314f281a6da1571cf7255");
}
#[test]
fn test_empty_password() {
let challenge = generate_challenge();
let hash = compute_auth_hash("", &challenge);
assert_eq!(hash.len(), 16);
}
}

146
src/bandwidth.rs Normal file
View File

@@ -0,0 +1,146 @@
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64};
use std::sync::Arc;
use std::time::Duration;
/// Shared state for bandwidth tracking between TX/RX threads and status reporter.
#[derive(Debug)]
pub struct BandwidthState {
pub tx_bytes: AtomicU64,
pub rx_bytes: AtomicU64,
pub tx_speed: AtomicU32,
pub tx_speed_changed: AtomicBool,
pub running: AtomicBool,
pub rx_packets: AtomicU64,
pub rx_lost_packets: AtomicU64,
pub last_udp_seq: AtomicU32,
}
impl BandwidthState {
pub fn new() -> Arc<Self> {
Arc::new(Self {
tx_bytes: AtomicU64::new(0),
rx_bytes: AtomicU64::new(0),
tx_speed: AtomicU32::new(0),
tx_speed_changed: AtomicBool::new(false),
running: AtomicBool::new(true),
rx_packets: AtomicU64::new(0),
rx_lost_packets: AtomicU64::new(0),
last_udp_seq: AtomicU32::new(0),
})
}
}
/// Calculate the sleep interval between packets to achieve target bandwidth.
/// Returns None if speed is unlimited (0).
pub fn calc_send_interval(tx_speed_bps: u32, tx_size: u16) -> Option<Duration> {
if tx_speed_bps == 0 {
return None;
}
let bits_per_packet = tx_size as u64 * 8;
let interval_ns = (1_000_000_000u64 * bits_per_packet) / tx_speed_bps as u64;
// Replicate MikroTik behavior: if interval > 500ms, clamp to 1 second
if interval_ns > 500_000_000 {
Some(Duration::from_secs(1))
} else {
Some(Duration::from_nanos(interval_ns.max(1)))
}
}
/// Format a bandwidth value in human-readable form.
pub fn format_bandwidth(bits_per_sec: f64) -> String {
if bits_per_sec >= 1_000_000_000.0 {
format!("{:.2} Gbps", bits_per_sec / 1_000_000_000.0)
} else if bits_per_sec >= 1_000_000.0 {
format!("{:.2} Mbps", bits_per_sec / 1_000_000.0)
} else if bits_per_sec >= 1_000.0 {
format!("{:.2} Kbps", bits_per_sec / 1_000.0)
} else {
format!("{:.0} bps", bits_per_sec)
}
}
/// Parse bandwidth string like "100M", "1G", "500K", "1000000"
pub fn parse_bandwidth(s: &str) -> std::result::Result<u32, anyhow::Error> {
let s = s.trim();
if s.is_empty() {
return Err(anyhow::anyhow!("Empty bandwidth string"));
}
let (num_str, multiplier) = match s.as_bytes().last() {
Some(b'G' | b'g') => (&s[..s.len() - 1], 1_000_000_000u64),
Some(b'M' | b'm') => (&s[..s.len() - 1], 1_000_000u64),
Some(b'K' | b'k') => (&s[..s.len() - 1], 1_000u64),
_ => (s, 1u64),
};
let num: f64 = num_str
.parse()
.map_err(|e| anyhow::anyhow!("Invalid bandwidth number '{}': {}", num_str, e))?;
let result = (num * multiplier as f64) as u64;
if result > u32::MAX as u64 {
Err(anyhow::anyhow!("Bandwidth {} exceeds maximum (4 Gbps)", s))
} else {
Ok(result as u32)
}
}
/// Print a status line for a reporting interval.
pub fn print_status(
interval_num: u32,
direction: &str,
bytes: u64,
elapsed: Duration,
lost_packets: Option<u64>,
) {
let secs = elapsed.as_secs_f64();
let bits = bytes as f64 * 8.0;
let bw = if secs > 0.0 { bits / secs } else { 0.0 };
let loss_str = match lost_packets {
Some(lost) if lost > 0 => format!(" lost: {}", lost),
_ => String::new(),
};
println!(
"[{:4}] {:>3} {} ({} bytes){}",
interval_num,
direction,
format_bandwidth(bw),
bytes,
loss_str,
);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_bandwidth() {
assert_eq!(parse_bandwidth("100M").unwrap(), 100_000_000);
assert_eq!(parse_bandwidth("1G").unwrap(), 1_000_000_000);
assert_eq!(parse_bandwidth("500K").unwrap(), 500_000);
assert_eq!(parse_bandwidth("1000000").unwrap(), 1_000_000);
assert_eq!(parse_bandwidth("1.5M").unwrap(), 1_500_000);
}
#[test]
fn test_calc_interval() {
// 100Mbps with 1500 byte packets
let interval = calc_send_interval(100_000_000, 1500).unwrap();
// Expected: (1e9 * 1500 * 8) / 100_000_000 = 120_000 ns = 120 us
assert_eq!(interval.as_nanos(), 120_000);
// Unlimited
assert!(calc_send_interval(0, 1500).is_none());
}
#[test]
fn test_format_bandwidth() {
assert_eq!(format_bandwidth(100_000_000.0), "100.00 Mbps");
assert_eq!(format_bandwidth(1_500_000_000.0), "1.50 Gbps");
assert_eq!(format_bandwidth(500_000.0), "500.00 Kbps");
}
}

440
src/client.rs Normal file
View File

@@ -0,0 +1,440 @@
use std::net::SocketAddr;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpStream, UdpSocket};
use crate::auth;
use crate::bandwidth::{self, BandwidthState};
use crate::protocol::*;
pub async fn run_client(
host: &str,
port: u16,
direction: u8,
use_udp: bool,
tx_speed: u32,
rx_speed: u32,
auth_user: Option<String>,
auth_pass: Option<String>,
nat_mode: bool,
) -> Result<()> {
let addr = format!("{}:{}", host, port);
tracing::info!("Connecting to {}...", addr);
let mut stream = TcpStream::connect(&addr).await?;
stream.set_nodelay(true)?;
recv_hello(&mut stream).await?;
tracing::info!("Connected to server");
let proto = if use_udp { CMD_PROTO_UDP } else { CMD_PROTO_TCP };
let mut cmd = Command::new(proto, direction);
cmd.local_tx_speed = tx_speed;
cmd.remote_tx_speed = rx_speed;
send_command(&mut stream, &cmd).await?;
let resp = recv_response(&mut stream).await?;
match (auth_user.as_deref(), auth_pass.as_deref()) {
(Some(user), Some(pass)) => {
auth::client_authenticate(&mut stream, resp, user, pass).await?;
}
_ => {
if resp == AUTH_REQUIRED {
return Err(BtestError::Protocol(
"Server requires authentication but no credentials provided".into(),
));
}
if resp == [0x03, 0x00, 0x00, 0x00] {
return Err(BtestError::Protocol(
"Server requires EC-SRP5 authentication (RouterOS >= 6.43) which is not yet supported. \
Try disabling authentication on the MikroTik btest server, or provide -a/-p credentials".into(),
));
}
if resp != AUTH_OK {
return Err(BtestError::Protocol(format!(
"Unexpected server response: {:02x?}",
resp
)));
}
}
}
tracing::info!(
"Starting {} {} test",
if use_udp { "UDP" } else { "TCP" },
match direction {
CMD_DIR_RX => "upload (client TX)",
CMD_DIR_TX => "download (client RX)",
CMD_DIR_BOTH => "bidirectional",
_ => "unknown",
},
);
if use_udp {
run_udp_test_client(&mut stream, host, &cmd, nat_mode).await
} else {
run_tcp_test_client(stream, cmd).await
}
}
// --- TCP Test Client ---
async fn run_tcp_test_client(stream: TcpStream, cmd: Command) -> Result<()> {
let state = BandwidthState::new();
let tx_size = cmd.tx_size as usize;
let client_should_tx = cmd.client_tx();
let client_should_rx = cmd.client_rx();
let tx_speed = cmd.local_tx_speed;
let (reader, writer) = stream.into_split();
// IMPORTANT: Do NOT drop unused halves - dropping OwnedWriteHalf sends TCP FIN,
// causing the peer to think we disconnected. Use Option to conditionally move.
let mut _writer_keepalive = None;
let mut _reader_keepalive = None;
let state_tx = state.clone();
let tx_handle = if client_should_tx {
Some(tokio::spawn(async move {
tcp_client_tx_loop(writer, tx_size, tx_speed, state_tx).await
}))
} else {
_writer_keepalive = Some(writer);
None
};
let state_rx = state.clone();
let rx_handle = if client_should_rx {
Some(tokio::spawn(async move {
tcp_client_rx_loop(reader, state_rx).await
}))
} else {
_reader_keepalive = Some(reader);
None
};
client_status_loop(&cmd, &state).await;
state.running.store(false, Ordering::SeqCst);
if let Some(h) = tx_handle { let _ = h.await; }
if let Some(h) = rx_handle { let _ = h.await; }
Ok(())
}
async fn tcp_client_tx_loop(
mut writer: tokio::net::tcp::OwnedWriteHalf,
tx_size: usize,
tx_speed: u32,
state: Arc<BandwidthState>,
) {
tokio::time::sleep(Duration::from_millis(100)).await;
let mut packet = vec![0u8; tx_size];
packet[0] = STATUS_MSG_TYPE;
let mut interval = bandwidth::calc_send_interval(tx_speed, tx_size as u16);
let mut next_send = Instant::now();
while state.running.load(Ordering::Relaxed) {
if writer.write_all(&packet).await.is_err() {
break;
}
state.tx_bytes.fetch_add(tx_size as u64, Ordering::Relaxed);
if state.tx_speed_changed.load(Ordering::Relaxed) {
state.tx_speed_changed.store(false, Ordering::Relaxed);
let new_speed = state.tx_speed.load(Ordering::Relaxed);
interval = bandwidth::calc_send_interval(new_speed, tx_size as u16);
next_send = Instant::now();
}
match interval {
Some(iv) => {
next_send += iv;
let now = Instant::now();
if next_send > now {
tokio::time::sleep(next_send - now).await;
}
}
None => {
tokio::task::yield_now().await;
}
}
}
}
async fn tcp_client_rx_loop(
mut reader: tokio::net::tcp::OwnedReadHalf,
state: Arc<BandwidthState>,
) {
let mut buf = vec![0u8; 65536];
while state.running.load(Ordering::Relaxed) {
match reader.read(&mut buf).await {
Ok(0) | Err(_) => break,
Ok(n) => {
state.rx_bytes.fetch_add(n as u64, Ordering::Relaxed);
}
}
}
}
// --- UDP Test Client ---
async fn run_udp_test_client(
stream: &mut TcpStream,
host: &str,
cmd: &Command,
nat_mode: bool,
) -> Result<()> {
let mut port_buf = [0u8; 2];
stream.read_exact(&mut port_buf).await?;
let server_udp_port = u16::from_be_bytes(port_buf);
let client_udp_port = server_udp_port + BTEST_PORT_CLIENT_OFFSET;
tracing::info!(
"UDP ports: server={}, client={}",
server_udp_port, client_udp_port,
);
let udp = UdpSocket::bind(format!("0.0.0.0:{}", client_udp_port)).await?;
let server_udp_addr: SocketAddr =
format!("{}:{}", host, server_udp_port).parse().unwrap();
udp.connect(server_udp_addr).await?;
if nat_mode {
tracing::info!("NAT mode: sending probe packet");
udp.send(&[]).await?;
}
let state = BandwidthState::new();
let tx_size = cmd.tx_size as usize;
let client_should_tx = cmd.client_tx();
let client_should_rx = cmd.client_rx();
let tx_speed = cmd.local_tx_speed;
let udp = Arc::new(udp);
let state_tx = state.clone();
let udp_tx = udp.clone();
let tx_handle = if client_should_tx {
Some(tokio::spawn(async move {
udp_client_tx_loop(&udp_tx, tx_size, tx_speed, state_tx).await
}))
} else {
None
};
let state_rx = state.clone();
let udp_rx = udp.clone();
let rx_handle = if client_should_rx {
Some(tokio::spawn(async move {
udp_client_rx_loop(&udp_rx, state_rx).await
}))
} else {
None
};
udp_client_status_loop(stream, cmd, &state).await;
state.running.store(false, Ordering::SeqCst);
if let Some(h) = tx_handle { let _ = h.await; }
if let Some(h) = rx_handle { let _ = h.await; }
Ok(())
}
async fn udp_client_tx_loop(
socket: &UdpSocket,
tx_size: usize,
initial_tx_speed: u32,
state: Arc<BandwidthState>,
) {
let mut seq: u32 = 0;
let mut packet = vec![0u8; tx_size];
let mut interval = bandwidth::calc_send_interval(initial_tx_speed, tx_size as u16);
let mut next_send = Instant::now();
let mut consecutive_errors: u32 = 0;
while state.running.load(Ordering::Relaxed) {
packet[0..4].copy_from_slice(&seq.to_be_bytes());
match socket.send(&packet).await {
Ok(n) => {
seq = seq.wrapping_add(1);
state.tx_bytes.fetch_add(n as u64, Ordering::Relaxed);
consecutive_errors = 0;
}
Err(_) => {
consecutive_errors += 1;
if consecutive_errors > 1000 {
tracing::warn!("UDP TX: too many consecutive send errors, stopping");
break;
}
tokio::time::sleep(Duration::from_micros(200)).await;
continue;
}
}
if state.tx_speed_changed.load(Ordering::Relaxed) {
state.tx_speed_changed.store(false, Ordering::Relaxed);
let new_speed = state.tx_speed.load(Ordering::Relaxed);
interval = bandwidth::calc_send_interval(new_speed, tx_size as u16);
next_send = Instant::now();
tracing::debug!("TX speed adjusted to {} bps ({:.2} Mbps)",
new_speed, new_speed as f64 / 1_000_000.0);
}
match interval {
Some(iv) => {
next_send += iv;
let now = Instant::now();
if next_send > now {
tokio::time::sleep(next_send - now).await;
}
}
None => {
if seq % 64 == 0 {
tokio::task::yield_now().await;
}
}
}
}
}
async fn udp_client_rx_loop(socket: &UdpSocket, state: Arc<BandwidthState>) {
let mut buf = vec![0u8; 65536];
let mut last_seq: Option<u32> = None;
while state.running.load(Ordering::Relaxed) {
match tokio::time::timeout(Duration::from_secs(5), socket.recv(&mut buf)).await {
Ok(Ok(n)) if n >= 4 => {
state.rx_bytes.fetch_add(n as u64, Ordering::Relaxed);
state.rx_packets.fetch_add(1, Ordering::Relaxed);
let seq = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
if let Some(last) = last_seq {
let expected = last.wrapping_add(1);
if seq > expected {
let lost = seq - expected;
state.rx_lost_packets.fetch_add(lost as u64, Ordering::Relaxed);
}
}
last_seq = Some(seq);
}
Ok(Ok(_)) => {}
Ok(Err(e)) => {
tracing::debug!("UDP recv error: {}", e);
tokio::time::sleep(Duration::from_millis(10)).await;
}
Err(_) => {
tracing::debug!("UDP RX timeout");
}
}
}
}
// --- Status Loops ---
async fn client_status_loop(cmd: &Command, state: &BandwidthState) {
let mut seq: u32 = 0;
let mut interval = tokio::time::interval(Duration::from_secs(1));
loop {
interval.tick().await;
if !state.running.load(Ordering::Relaxed) {
break;
}
seq += 1;
if cmd.client_tx() {
let tx = state.tx_bytes.swap(0, Ordering::Relaxed);
bandwidth::print_status(seq, "TX", tx, Duration::from_secs(1), None);
}
if cmd.client_rx() {
let rx = state.rx_bytes.swap(0, Ordering::Relaxed);
bandwidth::print_status(seq, "RX", rx, Duration::from_secs(1), None);
}
}
}
/// UDP status exchange - sequential like C pselect():
/// 1. Wait up to 1 second for server status
/// 2. Read and process if available
/// 3. ALWAYS send our status
async fn udp_client_status_loop(
stream: &mut TcpStream,
cmd: &Command,
state: &BandwidthState,
) {
let mut seq: u32 = 0;
let (mut reader, mut writer) = tokio::io::split(stream);
let mut status_buf = [0u8; STATUS_MSG_SIZE];
let mut next_status = Instant::now() + Duration::from_secs(1);
loop {
if !state.running.load(Ordering::Relaxed) {
break;
}
let now = Instant::now();
let wait_time = if next_status > now {
next_status - now
} else {
Duration::ZERO
};
match tokio::time::timeout(wait_time, reader.read_exact(&mut status_buf)).await {
Ok(Ok(_)) => {
let server_status = StatusMessage::deserialize(&status_buf);
if server_status.bytes_received > 0 && cmd.client_tx() {
let new_speed =
((server_status.bytes_received as u64 * 8 * 3) / 2) as u32;
state.tx_speed.store(new_speed, Ordering::Relaxed);
state.tx_speed_changed.store(true, Ordering::Relaxed);
tracing::debug!(
"Server received {} bytes → TX {:.2} Mbps",
server_status.bytes_received,
new_speed as f64 / 1_000_000.0,
);
}
if Instant::now() < next_status {
continue;
}
}
Ok(Err(_)) => {
state.running.store(false, Ordering::SeqCst);
break;
}
Err(_) => {}
}
// ALWAYS send our status every 1 second
seq += 1;
next_status = Instant::now() + Duration::from_secs(1);
let rx_bytes = state.rx_bytes.swap(0, Ordering::Relaxed);
let tx_bytes = state.tx_bytes.swap(0, Ordering::Relaxed);
let lost = state.rx_lost_packets.swap(0, Ordering::Relaxed);
let status = StatusMessage {
seq,
bytes_received: rx_bytes as u32,
};
if writer.write_all(&status.serialize()).await.is_err() {
state.running.store(false, Ordering::SeqCst);
break;
}
let _ = writer.flush().await;
if cmd.client_tx() {
bandwidth::print_status(seq, "TX", tx_bytes, Duration::from_secs(1), None);
}
if cmd.client_rx() {
bandwidth::print_status(seq, "RX", rx_bytes, Duration::from_secs(1), Some(lost));
}
}
}

5
src/lib.rs Normal file
View File

@@ -0,0 +1,5 @@
pub mod auth;
pub mod bandwidth;
pub mod client;
pub mod protocol;
pub mod server;

136
src/main.rs Normal file
View File

@@ -0,0 +1,136 @@
mod auth;
mod bandwidth;
mod client;
mod protocol;
mod server;
use clap::Parser;
use tracing_subscriber::EnvFilter;
use crate::protocol::*;
#[derive(Parser, Debug)]
#[command(
name = "btest",
about = "MikroTik Bandwidth Test (btest) - server and client",
version,
long_about = "Compatible bandwidth testing tool for MikroTik RouterOS devices.\n\
Supports TCP and UDP modes with optional authentication."
)]
struct Cli {
/// Run in server mode
#[arg(short = 's', long = "server", conflicts_with = "client")]
server: bool,
/// Run in client mode, connecting to the specified host
#[arg(short = 'c', long = "client", conflicts_with = "server")]
client: Option<String>,
/// Client transmits data (upload test)
#[arg(short = 't', long = "transmit")]
transmit: bool,
/// Client receives data (download test)
#[arg(short = 'r', long = "receive")]
receive: bool,
/// Use UDP instead of TCP
#[arg(short = 'u', long = "udp")]
udp: bool,
/// Target bandwidth (e.g., 100M, 1G, 500K)
#[arg(short = 'b', long = "bandwidth")]
bandwidth: Option<String>,
/// Listen/connect port (default: 2000)
#[arg(short = 'P', long = "port", default_value_t = BTEST_PORT)]
port: u16,
/// Authentication username
#[arg(short = 'a', long = "authuser")]
auth_user: Option<String>,
/// Authentication password
#[arg(short = 'p', long = "authpass")]
auth_pass: Option<String>,
/// NAT mode - send probe packet to open firewall
#[arg(short = 'n', long = "nat")]
nat: bool,
/// Verbose logging (repeat for more: -v, -vv, -vvv)
#[arg(short = 'v', long = "verbose", action = clap::ArgAction::Count)]
verbose: u8,
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let cli = Cli::parse();
// Set up logging based on verbosity
let filter = match cli.verbose {
0 => "info",
1 => "debug",
_ => "trace",
};
tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(filter)),
)
.with_target(false)
.init();
if cli.server {
// Server mode
tracing::info!("Starting btest server on port {}", cli.port);
server::run_server(cli.port, cli.auth_user, cli.auth_pass).await?;
} else if let Some(host) = cli.client {
// Client mode - must specify at least one direction
if !cli.transmit && !cli.receive {
eprintln!("Error: Client mode requires at least one of -t (transmit) or -r (receive)");
std::process::exit(1);
}
// Direction tells SERVER what to do (C client convention):
// client transmit → CMD_DIR_RX (server receives)
// client receive → CMD_DIR_TX (server transmits)
let direction = match (cli.transmit, cli.receive) {
(true, false) => CMD_DIR_RX,
(false, true) => CMD_DIR_TX,
(true, true) => CMD_DIR_BOTH,
_ => unreachable!(),
};
let bw = match &cli.bandwidth {
Some(b) => bandwidth::parse_bandwidth(b)?,
None => 0,
};
// For client: local_tx_speed controls upload, remote_tx_speed controls download
let (tx_speed, rx_speed) = match direction {
CMD_DIR_TX => (bw, 0),
CMD_DIR_RX => (0, bw),
CMD_DIR_BOTH => (bw, bw),
_ => (0, 0),
};
client::run_client(
&host,
cli.port,
direction,
cli.udp,
tx_speed,
rx_speed,
cli.auth_user,
cli.auth_pass,
cli.nat,
)
.await?;
} else {
eprintln!("Error: Must specify either -s (server) or -c <host> (client)");
eprintln!("Run with --help for usage information.");
std::process::exit(1);
}
Ok(())
}

205
src/protocol.rs Normal file
View File

@@ -0,0 +1,205 @@
use std::io;
use thiserror::Error;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
// --- Constants ---
pub const BTEST_PORT: u16 = 2000;
pub const BTEST_UDP_PORT_START: u16 = 2001;
pub const BTEST_PORT_CLIENT_OFFSET: u16 = 256;
pub const CMD_PROTO_UDP: u8 = 0x00;
pub const CMD_PROTO_TCP: u8 = 0x01;
pub const CMD_DIR_RX: u8 = 0x01;
pub const CMD_DIR_TX: u8 = 0x02;
pub const CMD_DIR_BOTH: u8 = 0x03;
pub const DEFAULT_TCP_TX_SIZE: u16 = 0x8000; // 32768
pub const DEFAULT_UDP_TX_SIZE: u16 = 0x05DC; // 1500
pub const HELLO: [u8; 4] = [0x01, 0x00, 0x00, 0x00];
pub const AUTH_OK: [u8; 4] = [0x01, 0x00, 0x00, 0x00];
pub const AUTH_REQUIRED: [u8; 4] = [0x02, 0x00, 0x00, 0x00];
pub const AUTH_FAILED: [u8; 4] = [0x00, 0x00, 0x00, 0x00];
pub const STATUS_MSG_TYPE: u8 = 0x07;
pub const STATUS_MSG_SIZE: usize = 12;
// --- Error Types ---
#[derive(Error, Debug)]
pub enum BtestError {
#[error("I/O error: {0}")]
Io(#[from] io::Error),
#[error("Protocol error: {0}")]
Protocol(String),
#[error("Authentication failed")]
AuthFailed,
#[error("Invalid command")]
InvalidCommand,
}
pub type Result<T> = std::result::Result<T, BtestError>;
// --- Command Structure ---
#[derive(Debug, Clone)]
pub struct Command {
pub proto: u8,
pub direction: u8,
pub random_data: u8,
pub tcp_conn_count: u8,
pub tx_size: u16,
pub client_buf_size: u16,
pub remote_tx_speed: u32,
pub local_tx_speed: u32,
}
impl Command {
pub fn new(proto: u8, direction: u8) -> Self {
let tx_size = if proto == CMD_PROTO_UDP {
DEFAULT_UDP_TX_SIZE
} else {
DEFAULT_TCP_TX_SIZE
};
Self {
proto,
direction,
random_data: 0,
tcp_conn_count: 0,
tx_size,
client_buf_size: 0,
remote_tx_speed: 0,
local_tx_speed: 0,
}
}
pub fn serialize(&self) -> [u8; 16] {
let mut buf = [0u8; 16];
buf[0] = self.proto;
buf[1] = self.direction;
buf[2] = self.random_data;
buf[3] = self.tcp_conn_count;
buf[4..6].copy_from_slice(&self.tx_size.to_le_bytes());
buf[6..8].copy_from_slice(&self.client_buf_size.to_le_bytes());
buf[8..12].copy_from_slice(&self.remote_tx_speed.to_le_bytes());
buf[12..16].copy_from_slice(&self.local_tx_speed.to_le_bytes());
buf
}
pub fn deserialize(buf: &[u8; 16]) -> Self {
Self {
proto: buf[0],
direction: buf[1],
random_data: buf[2],
tcp_conn_count: buf[3],
tx_size: u16::from_le_bytes([buf[4], buf[5]]),
client_buf_size: u16::from_le_bytes([buf[6], buf[7]]),
remote_tx_speed: u32::from_le_bytes([buf[8], buf[9], buf[10], buf[11]]),
local_tx_speed: u32::from_le_bytes([buf[12], buf[13], buf[14], buf[15]]),
}
}
pub fn is_udp(&self) -> bool {
self.proto == CMD_PROTO_UDP
}
// Direction bits are from SERVER's perspective:
// CMD_DIR_RX (0x01) = server receives
// CMD_DIR_TX (0x02) = server transmits
// Client inverts when building: client TX → CMD_DIR_RX, client RX → CMD_DIR_TX
/// Server should transmit (CMD_DIR_TX bit set)
pub fn server_tx(&self) -> bool {
self.direction & CMD_DIR_TX != 0
}
/// Server should receive (CMD_DIR_RX bit set)
pub fn server_rx(&self) -> bool {
self.direction & CMD_DIR_RX != 0
}
/// Client should transmit (inverse: CMD_DIR_RX bit = server receives our data)
pub fn client_tx(&self) -> bool {
self.direction & CMD_DIR_RX != 0
}
/// Client should receive (inverse: CMD_DIR_TX bit = server sends us data)
pub fn client_rx(&self) -> bool {
self.direction & CMD_DIR_TX != 0
}
}
// --- Status Message ---
#[derive(Debug, Clone, Default)]
pub struct StatusMessage {
pub seq: u32,
pub bytes_received: u32,
}
impl StatusMessage {
pub fn serialize(&self) -> [u8; STATUS_MSG_SIZE] {
let mut buf = [0u8; STATUS_MSG_SIZE];
buf[0] = STATUS_MSG_TYPE;
buf[1..5].copy_from_slice(&self.seq.to_be_bytes());
buf[5] = 0;
buf[6] = 0;
buf[7] = 0;
buf[8..12].copy_from_slice(&self.bytes_received.to_le_bytes());
buf
}
pub fn deserialize(buf: &[u8; STATUS_MSG_SIZE]) -> Self {
Self {
seq: u32::from_be_bytes([buf[1], buf[2], buf[3], buf[4]]),
bytes_received: u32::from_le_bytes([buf[8], buf[9], buf[10], buf[11]]),
}
}
}
// --- Protocol Helpers ---
pub async fn send_hello<W: AsyncWriteExt + Unpin>(writer: &mut W) -> Result<()> {
writer.write_all(&HELLO).await?;
writer.flush().await?;
Ok(())
}
pub async fn recv_hello<R: AsyncReadExt + Unpin>(reader: &mut R) -> Result<()> {
let mut buf = [0u8; 4];
reader.read_exact(&mut buf).await?;
if buf != HELLO {
return Err(BtestError::Protocol(format!(
"Expected HELLO {:02x?}, got {:02x?}",
HELLO, buf
)));
}
Ok(())
}
pub async fn send_command<W: AsyncWriteExt + Unpin>(
writer: &mut W,
cmd: &Command,
) -> Result<()> {
writer.write_all(&cmd.serialize()).await?;
writer.flush().await?;
Ok(())
}
pub async fn recv_command<R: AsyncReadExt + Unpin>(reader: &mut R) -> Result<Command> {
let mut buf = [0u8; 16];
reader.read_exact(&mut buf).await?;
let cmd = Command::deserialize(&buf);
if cmd.proto > 1 || cmd.direction == 0 || cmd.direction > 3 {
return Err(BtestError::InvalidCommand);
}
Ok(cmd)
}
pub async fn recv_response<R: AsyncReadExt + Unpin>(reader: &mut R) -> Result<[u8; 4]> {
let mut buf = [0u8; 4];
reader.read_exact(&mut buf).await?;
Ok(buf)
}

456
src/server.rs Normal file
View File

@@ -0,0 +1,456 @@
use std::net::SocketAddr;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use crate::auth;
use crate::bandwidth::{self, BandwidthState};
use crate::protocol::*;
pub async fn run_server(
port: u16,
auth_user: Option<String>,
auth_pass: Option<String>,
) -> Result<()> {
let addr = format!("0.0.0.0:{}", port);
let listener = TcpListener::bind(&addr).await?;
tracing::info!("btest server listening on {}", addr);
let udp_port_offset = Arc::new(std::sync::atomic::AtomicU16::new(0));
loop {
let (stream, peer) = listener.accept().await?;
tracing::info!("New connection from {}", peer);
let auth_user = auth_user.clone();
let auth_pass = auth_pass.clone();
let udp_offset = udp_port_offset.clone();
tokio::spawn(async move {
if let Err(e) = handle_client(stream, peer, auth_user, auth_pass, udp_offset).await {
tracing::error!("Client {} error: {}", peer, e);
}
});
}
}
async fn handle_client(
mut stream: TcpStream,
peer: SocketAddr,
auth_user: Option<String>,
auth_pass: Option<String>,
udp_port_offset: Arc<std::sync::atomic::AtomicU16>,
) -> Result<()> {
stream.set_nodelay(true)?;
send_hello(&mut stream).await?;
let cmd = recv_command(&mut stream).await?;
tracing::info!(
"Client {} command: proto={} dir={} tx_size={} remote_speed={} local_speed={}",
peer,
if cmd.is_udp() { "UDP" } else { "TCP" },
match cmd.direction {
CMD_DIR_RX => "RX",
CMD_DIR_TX => "TX",
CMD_DIR_BOTH => "BOTH",
_ => "?",
},
cmd.tx_size,
cmd.remote_tx_speed,
cmd.local_tx_speed,
);
auth::server_authenticate(
&mut stream,
auth_user.as_deref(),
auth_pass.as_deref(),
)
.await?;
if cmd.is_udp() {
run_udp_test_server(&mut stream, peer, &cmd, udp_port_offset).await
} else {
run_tcp_test_server(stream, cmd).await
}
}
// --- TCP Test Server ---
async fn run_tcp_test_server(stream: TcpStream, cmd: Command) -> Result<()> {
let state = BandwidthState::new();
let tx_size = cmd.tx_size as usize;
let server_should_tx = cmd.server_tx();
let server_should_rx = cmd.server_rx();
let tx_speed = cmd.remote_tx_speed;
let (reader, writer) = stream.into_split();
// IMPORTANT: Do NOT drop unused halves - dropping sends TCP FIN
let mut _writer_keepalive = None;
let mut _reader_keepalive = None;
let state_tx = state.clone();
let tx_handle = if server_should_tx {
Some(tokio::spawn(async move {
tcp_tx_loop(writer, tx_size, tx_speed, state_tx).await
}))
} else {
_writer_keepalive = Some(writer);
None
};
let state_rx = state.clone();
let rx_handle = if server_should_rx {
Some(tokio::spawn(async move {
tcp_rx_loop(reader, state_rx).await
}))
} else {
_reader_keepalive = Some(reader);
None
};
status_report_loop(&cmd, &state).await;
state.running.store(false, Ordering::SeqCst);
if let Some(h) = tx_handle { let _ = h.await; }
if let Some(h) = rx_handle { let _ = h.await; }
Ok(())
}
async fn tcp_tx_loop(
mut writer: tokio::net::tcp::OwnedWriteHalf,
tx_size: usize,
tx_speed: u32,
state: Arc<BandwidthState>,
) {
tokio::time::sleep(Duration::from_millis(100)).await;
let mut packet = vec![0u8; tx_size];
packet[0] = STATUS_MSG_TYPE;
let mut interval = bandwidth::calc_send_interval(tx_speed, tx_size as u16);
let mut next_send = Instant::now();
while state.running.load(Ordering::Relaxed) {
if writer.write_all(&packet).await.is_err() {
break;
}
state.tx_bytes.fetch_add(tx_size as u64, Ordering::Relaxed);
if state.tx_speed_changed.load(Ordering::Relaxed) {
state.tx_speed_changed.store(false, Ordering::Relaxed);
let new_speed = state.tx_speed.load(Ordering::Relaxed);
interval = bandwidth::calc_send_interval(new_speed, tx_size as u16);
next_send = Instant::now();
}
match interval {
Some(iv) => {
next_send += iv;
let now = Instant::now();
if next_send > now {
tokio::time::sleep(next_send - now).await;
}
}
None => {
tokio::task::yield_now().await;
}
}
}
}
async fn tcp_rx_loop(mut reader: tokio::net::tcp::OwnedReadHalf, state: Arc<BandwidthState>) {
let mut buf = vec![0u8; 65536];
while state.running.load(Ordering::Relaxed) {
match reader.read(&mut buf).await {
Ok(0) | Err(_) => break,
Ok(n) => {
state.rx_bytes.fetch_add(n as u64, Ordering::Relaxed);
}
}
}
}
// --- UDP Test Server ---
async fn run_udp_test_server(
stream: &mut TcpStream,
peer: SocketAddr,
cmd: &Command,
udp_port_offset: Arc<std::sync::atomic::AtomicU16>,
) -> Result<()> {
let offset = udp_port_offset.fetch_add(1, Ordering::SeqCst);
let server_udp_port = BTEST_UDP_PORT_START + offset;
let client_udp_port = server_udp_port + BTEST_PORT_CLIENT_OFFSET;
stream.write_all(&server_udp_port.to_be_bytes()).await?;
stream.flush().await?;
tracing::info!(
"UDP test: server_port={}, client_port={}, peer={}",
server_udp_port, client_udp_port, peer,
);
let udp = UdpSocket::bind(format!("0.0.0.0:{}", server_udp_port)).await?;
let client_udp_addr: SocketAddr =
format!("{}:{}", peer.ip(), client_udp_port).parse().unwrap();
udp.connect(client_udp_addr).await?;
let state = BandwidthState::new();
let tx_size = cmd.tx_size as usize;
let server_should_tx = cmd.server_tx();
let server_should_rx = cmd.server_rx();
let tx_speed = cmd.remote_tx_speed;
let udp = Arc::new(udp);
let state_tx = state.clone();
let udp_tx = udp.clone();
let tx_handle = if server_should_tx {
Some(tokio::spawn(async move {
udp_tx_loop(&udp_tx, tx_size, tx_speed, state_tx).await
}))
} else {
None
};
let state_rx = state.clone();
let udp_rx = udp.clone();
let rx_handle = if server_should_rx {
Some(tokio::spawn(async move {
udp_rx_loop(&udp_rx, state_rx).await
}))
} else {
None
};
// Status exchange using select! to match C pselect() behavior
udp_status_loop(stream, cmd, &state).await;
state.running.store(false, Ordering::SeqCst);
if let Some(h) = tx_handle { let _ = h.await; }
if let Some(h) = rx_handle { let _ = h.await; }
Ok(())
}
async fn udp_tx_loop(
socket: &UdpSocket,
tx_size: usize,
initial_tx_speed: u32,
state: Arc<BandwidthState>,
) {
let mut seq: u32 = 0;
let mut packet = vec![0u8; tx_size];
let mut interval = bandwidth::calc_send_interval(initial_tx_speed, tx_size as u16);
let mut next_send = Instant::now();
let mut consecutive_errors: u32 = 0;
while state.running.load(Ordering::Relaxed) {
packet[0..4].copy_from_slice(&seq.to_be_bytes());
match socket.send(&packet).await {
Ok(n) => {
seq = seq.wrapping_add(1);
state.tx_bytes.fetch_add(n as u64, Ordering::Relaxed);
consecutive_errors = 0;
}
Err(_) => {
consecutive_errors += 1;
if consecutive_errors > 1000 {
tracing::warn!("UDP TX: too many consecutive send errors, stopping");
break;
}
// Back off on ENOBUFS/EAGAIN
tokio::time::sleep(Duration::from_micros(200)).await;
continue;
}
}
// Pick up dynamic speed changes from status loop
if state.tx_speed_changed.load(Ordering::Relaxed) {
state.tx_speed_changed.store(false, Ordering::Relaxed);
let new_speed = state.tx_speed.load(Ordering::Relaxed);
interval = bandwidth::calc_send_interval(new_speed, tx_size as u16);
next_send = Instant::now();
tracing::debug!("TX speed adjusted to {} bps ({:.2} Mbps)",
new_speed, new_speed as f64 / 1_000_000.0);
}
match interval {
Some(iv) => {
next_send += iv;
let now = Instant::now();
if next_send > now {
tokio::time::sleep(next_send - now).await;
}
}
None => {
// Unlimited: yield every 64 packets to keep system responsive
if seq % 64 == 0 {
tokio::task::yield_now().await;
}
}
}
}
}
async fn udp_rx_loop(socket: &UdpSocket, state: Arc<BandwidthState>) {
let mut buf = vec![0u8; 65536];
let mut last_seq: Option<u32> = None;
while state.running.load(Ordering::Relaxed) {
match tokio::time::timeout(Duration::from_secs(5), socket.recv(&mut buf)).await {
Ok(Ok(n)) if n >= 4 => {
state.rx_bytes.fetch_add(n as u64, Ordering::Relaxed);
state.rx_packets.fetch_add(1, Ordering::Relaxed);
let seq = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
if let Some(last) = last_seq {
let expected = last.wrapping_add(1);
if seq > expected {
let lost = seq - expected;
state.rx_lost_packets.fetch_add(lost as u64, Ordering::Relaxed);
}
}
last_seq = Some(seq);
state.last_udp_seq.store(seq, Ordering::Relaxed);
}
Ok(Ok(_)) => {}
Ok(Err(e)) => {
tracing::debug!("UDP recv error: {}", e);
tokio::time::sleep(Duration::from_millis(10)).await;
}
Err(_) => {
tracing::debug!("UDP RX timeout");
}
}
}
}
// --- Status Reporting ---
async fn status_report_loop(cmd: &Command, state: &BandwidthState) {
let mut seq: u32 = 0;
let mut interval = tokio::time::interval(Duration::from_secs(1));
loop {
interval.tick().await;
if !state.running.load(Ordering::Relaxed) {
break;
}
seq += 1;
if cmd.server_tx() {
let tx = state.tx_bytes.swap(0, Ordering::Relaxed);
bandwidth::print_status(seq, "TX", tx, Duration::from_secs(1), None);
}
if cmd.server_rx() {
let rx = state.rx_bytes.swap(0, Ordering::Relaxed);
let lost = state.rx_lost_packets.swap(0, Ordering::Relaxed);
let lost_opt = if cmd.is_udp() { Some(lost) } else { None };
bandwidth::print_status(seq, "RX", rx, Duration::from_secs(1), lost_opt);
}
}
}
/// UDP status exchange loop - matches C pselect() behavior exactly:
/// 1. Wait up to 1 second for client status (like pselect with 1s timeout)
/// 2. If client sent status, read and process it
/// 3. ALWAYS send our status (unconditional, matching C line 1048)
/// 4. Reset counters and print stats
/// This sequential approach prevents the ticker from being starved.
async fn udp_status_loop(
stream: &mut TcpStream,
cmd: &Command,
state: &BandwidthState,
) {
let mut seq: u32 = 0;
let (mut reader, mut writer) = tokio::io::split(stream);
let mut status_buf = [0u8; STATUS_MSG_SIZE];
let mut next_status = Instant::now() + Duration::from_secs(1);
loop {
if !state.running.load(Ordering::Relaxed) {
break;
}
// Step 1: Wait for client status OR timeout (like C pselect)
let now = Instant::now();
let wait_time = if next_status > now {
next_status - now
} else {
Duration::ZERO
};
// Try to read client's status within the remaining time window
match tokio::time::timeout(wait_time, reader.read_exact(&mut status_buf)).await {
Ok(Ok(_)) => {
let client_status = StatusMessage::deserialize(&status_buf);
tracing::debug!(
"RECV status: raw={:02x?} seq={} bytes_received={}",
&status_buf, client_status.seq, client_status.bytes_received,
);
if client_status.bytes_received > 0 && cmd.server_tx() {
let new_speed =
((client_status.bytes_received as u64 * 8 * 3) / 2) as u32;
state.tx_speed.store(new_speed, Ordering::Relaxed);
state.tx_speed_changed.store(true, Ordering::Relaxed);
tracing::debug!(
"Speed adjust: client got {} bytes → our TX {:.2} Mbps",
client_status.bytes_received,
new_speed as f64 / 1_000_000.0,
);
}
if Instant::now() < next_status {
continue;
}
}
Ok(Err(e)) => {
tracing::debug!("Client TCP read error: {}", e);
state.running.store(false, Ordering::SeqCst);
break;
}
Err(_) => {
// Timeout - 1 second elapsed
}
}
// Step 2: ALWAYS send our status every 1 second
seq += 1;
next_status = Instant::now() + Duration::from_secs(1);
let rx_bytes = state.rx_bytes.swap(0, Ordering::Relaxed);
let tx_bytes = state.tx_bytes.swap(0, Ordering::Relaxed);
let lost = state.rx_lost_packets.swap(0, Ordering::Relaxed);
let status = StatusMessage {
seq,
bytes_received: rx_bytes as u32,
};
let serialized = status.serialize();
tracing::debug!(
"SEND status: raw={:02x?} seq={} bytes_received={} ({:.2} Mbps)",
&serialized, seq, rx_bytes, rx_bytes as f64 * 8.0 / 1_000_000.0,
);
if writer.write_all(&serialized).await.is_err() {
state.running.store(false, Ordering::SeqCst);
break;
}
let _ = writer.flush().await;
// Print local stats
if cmd.server_tx() {
bandwidth::print_status(seq, "TX", tx_bytes, Duration::from_secs(1), None);
}
if cmd.server_rx() {
bandwidth::print_status(seq, "RX", rx_bytes, Duration::from_secs(1), Some(lost));
}
}
}