manawenuz/btest-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 7 Wiki Activity

Public btest server: byte budget, multi-conn, web dashboard, quotas

Browse Source 
- Inline byte budget in BandwidthState prevents quota overshoot at any
  link speed (TX/RX loops check per-packet, not per-interval)
- TCP multi-connection support for server-pro (session tokens, secondary
  connection joins, delegates to standard multi-conn handler)
- MD5 password verification against stored raw passwords in user DB
- Web dashboard: quota progress bars (daily/weekly/monthly), JSON export
  endpoint (/api/ip/{ip}/export), quota API (/api/ip/{ip}/quota)
- Landing page with usage instructions, UDP NAT warning, credentials
- Fix IP usage double-counting bug in QuotaManager::record_usage
- UserDb now stores DB path and raw passwords for MD5 auth
- 10 enforcer tests (4 new: budget calc, budget stop, budget exhausted,
  unlimited passthrough)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Siavash Sameni
2026-04-01 18:43:09 +04:00
parent 7dd4820d2c
commit 4cdcc4e6c4
7 changed files with 763 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files

295
src/server_pro/server_loop.rs
View File

@@ -2,14 +2,18 @@
//!
//! Wraps the standard btest server connection handler with:
//! - Pre-connection IP/user quota checks
//! - MD5 challenge-response auth against user DB
//! - TCP multi-connection session support
//! - Mid-session quota enforcement via QuotaEnforcer
//! - Post-session usage recording
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::Mutex;
use btest_rs::protocol::*;
use btest_rs::bandwidth::BandwidthState;
@@ -18,22 +22,27 @@ use super::enforcer::{QuotaEnforcer, StopReason};
use super::quota::{Direction, QuotaManager};
use super::user_db::UserDb;
/// Pending TCP multi-connection session.
struct TcpSession {
peer_ip: std::net::IpAddr,
username: String,
cmd: Command,
streams: Vec<TcpStream>,
expected: u8,
}
type SessionMap = Arc<Mutex<HashMap<u16, TcpSession>>>;
/// Run the pro server with quota enforcement.
pub async fn run_pro_server(
port: u16,
ecsrp5: bool,
_ecsrp5: bool,
listen_v4: Option<String>,
listen_v6: Option<String>,
db: UserDb,
quota_mgr: QuotaManager,
quota_check_interval: u64,
) -> anyhow::Result<()> {
// Pre-derive EC-SRP5 creds if needed
// For pro server, we don't use CLI -a/-p — we use the user DB
// EC-SRP5 needs a fixed password for the server challenge, but
// the actual verification happens against the DB.
// For now, the first user in the DB is used for EC-SRP5 derivation.
let v4_listener = if let Some(ref addr) = listen_v4 {
let bind_addr = format!("{}:{}", addr, port);
Some(TcpListener::bind(&bind_addr).await?)
@@ -52,6 +61,8 @@ pub async fn run_pro_server(
anyhow::bail!("No listeners bound");
}
let sessions: SessionMap = Arc::new(Mutex::new(HashMap::new()));
tracing::info!("btest-server-pro ready, accepting connections");
loop {
@@ -69,29 +80,14 @@ pub async fn run_pro_server(
tracing::info!("New connection from {}", peer);
// Pre-connection IP check
if let Err(e) = quota_mgr.check_ip(&peer.ip(), Direction::Both) {
tracing::warn!("Rejected {} — {}", peer, e);
btest_rs::syslog_logger::auth_failure(
&peer.to_string(), "-", "-", &format!("{}", e),
);
// Send a quick rejection and close
let mut s = stream;
let _ = s.write_all(&HELLO).await;
drop(s);
continue;
}
quota_mgr.connect(&peer.ip());
let db = db.clone();
let qm = quota_mgr.clone();
let qm_disconnect = quota_mgr.clone();
let interval = quota_check_interval;
let sess = sessions.clone();
tokio::spawn(async move {
match handle_pro_client(stream, peer, db, qm, interval).await {
Ok((username, stop_reason, tx, rx)) => {
let is_primary = match handle_pro_connection(stream, peer, db, qm.clone(), interval, sess).await {
Ok(Some((username, stop_reason, tx, rx))) => {
tracing::info!(
"Client {} (user '{}') finished: {} (tx={}, rx={})",
peer, username, stop_reason, tx, rx,
@@ -100,31 +96,100 @@ pub async fn run_pro_server(
&peer.to_string(), "btest", &format!("{}", stop_reason),
tx, rx, 0, 0,
);
true
}
Ok(None) => false, // secondary connection or pending multi-conn
Err(e) => {
tracing::error!("Client {} error: {}", peer, e);
true
}
};
// Only decrement connection count for primary connections
if is_primary {
qm.disconnect(&peer.ip());
}
qm_disconnect.disconnect(&peer.ip());
});
}
}
async fn handle_pro_client(
/// Handle a single TCP connection. Returns None for secondary multi-conn joins.
async fn handle_pro_connection(
mut stream: TcpStream,
peer: SocketAddr,
db: UserDb,
quota_mgr: QuotaManager,
quota_check_interval: u64,
) -> anyhow::Result<(String, StopReason, u64, u64)> {
sessions: SessionMap,
) -> anyhow::Result<Option<(String, StopReason, u64, u64)>> {
stream.set_nodelay(true)?;
// HELLO
stream.write_all(&HELLO).await?;
// Read command
// Read command (or session token for secondary connections)
let mut cmd_buf = [0u8; 16];
stream.read_exact(&mut cmd_buf).await?;
// Check if this is a secondary connection joining an existing TCP session
// Secondary connections send [HI, LO, ...] matching an existing session token
{
let potential_token = u16::from_be_bytes([cmd_buf[0], cmd_buf[1]]);
let mut map = sessions.lock().await;
if let Some(session) = map.get_mut(&potential_token) {
if session.peer_ip == peer.ip()
&& session.streams.len() < session.expected as usize
{
tracing::info!(
"Secondary connection from {} joining session (token={:04x}, {}/{})",
peer, potential_token,
session.streams.len() + 1, session.expected,
);
// Auth the secondary connection with same token response
let ok = [0x01, cmd_buf[0], cmd_buf[1], 0x00];
stream.write_all(&ok).await?;
stream.flush().await?;
session.streams.push(stream);
// If all connections have joined, start the test
if session.streams.len() >= session.expected as usize {
let session = map.remove(&potential_token).unwrap();
let db2 = db.clone();
let qm2 = quota_mgr.clone();
tokio::spawn(async move {
match run_pro_multiconn_test(
session.streams, session.cmd, peer,
&session.username, db2, qm2, quota_check_interval,
).await {
Ok((stop, tx, rx)) => {
tracing::info!(
"Multi-conn {} (user '{}') finished: {} (tx={}, rx={})",
peer, session.username, stop, tx, rx,
);
}
Err(e) => {
tracing::error!("Multi-conn {} error: {}", peer, e);
}
}
});
}
return Ok(None);
}
}
}
// Primary connection — check IP quota/connection limit now
if let Err(e) = quota_mgr.check_ip(&peer.ip(), Direction::Both) {
tracing::warn!("Rejected {} — {}", peer, e);
btest_rs::syslog_logger::auth_failure(
&peer.to_string(), "-", "-", &format!("{}", e),
);
return Ok(None);
}
quota_mgr.connect(&peer.ip());
let cmd = Command::deserialize(&cmd_buf);
tracing::info!(
@@ -136,14 +201,25 @@ async fn handle_pro_client(
cmd.tx_size,
);
// Authenticate — use MD5 auth with DB verification
// Send AUTH_REQUIRED
// Build auth OK response with session token for multi-connection
let is_tcp_multi = !cmd.is_udp() && cmd.tcp_conn_count > 0;
let session_token: u16 = if is_tcp_multi {
rand::random::<u16>() | 0x0101 // ensure both bytes non-zero
} else {
0
};
let ok_response: [u8; 4] = if is_tcp_multi {
[0x01, (session_token >> 8) as u8, (session_token & 0xFF) as u8, 0x00]
} else {
AUTH_OK
};
// Authenticate — MD5 challenge-response against DB
stream.write_all(&AUTH_REQUIRED).await?;
let challenge = btest_rs::auth::generate_challenge();
stream.write_all(&challenge).await?;
stream.flush().await?;
// Read response
let mut response = [0u8; 48];
stream.read_exact(&mut response).await?;
@@ -176,17 +252,21 @@ async fn handle_pro_client(
anyhow::bail!("User disabled");
}
// Verify MD5 hash against stored password hash
// We need to compute the expected hash using the user's password
// But we only store SHA256(user:pass), not the raw password.
// For MD5 auth, we need the raw password to compute MD5(pass + challenge).
// This is a limitation — MD5 auth needs the raw password.
// For now, accept any authenticated user (the hash verification
// happens on the client side with MikroTik).
// TODO: Store password in a reversible form or use EC-SRP5 only.
// Verify MD5 hash against stored raw password
if let Ok(Some(raw_pass)) = db.get_password(&username) {
let expected_hash = btest_rs::auth::compute_auth_hash(&raw_pass, &challenge);
if received_hash != expected_hash {
tracing::warn!("Auth failed: password mismatch for user '{}'", username);
stream.write_all(&AUTH_FAILED).await?;
btest_rs::syslog_logger::auth_failure(
&peer.to_string(), &username, "md5", "password mismatch",
);
anyhow::bail!("Auth failed");
}
}
// If no raw password stored, accept (backwards compat with old DB entries)
// Send AUTH_OK
stream.write_all(&AUTH_OK).await?;
stream.write_all(&ok_response).await?;
stream.flush().await?;
tracing::info!("Auth successful for user '{}'", username);
@@ -202,79 +282,168 @@ async fn handle_pro_client(
btest_rs::syslog_logger::auth_failure(
&peer.to_string(), &username, "quota", &format!("{}", e),
);
// Connection is already authenticated, just close it
return Ok((username, StopReason::UserDailyQuota, 0, 0));
return Ok(Some((username, StopReason::UserDailyQuota, 0, 0)));
}
// Start session tracking
// TCP multi-connection: register session and wait for secondary connections
if is_tcp_multi {
tracing::info!(
"TCP multi-connection: waiting for {} connections (token={:04x})",
cmd.tcp_conn_count, session_token,
);
let mut map = sessions.lock().await;
map.insert(session_token, TcpSession {
peer_ip: peer.ip(),
username: username.clone(),
cmd: cmd.clone(),
streams: vec![stream],
expected: cmd.tcp_conn_count, // tcp_conn_count includes the primary
});
// The test will be started when all connections join (in the secondary handler above)
return Ok(None);
}
// Single-connection test
run_pro_single_test(stream, cmd, peer, &username, db, quota_mgr, quota_check_interval).await
.map(|(stop, tx, rx)| Some((username, stop, tx, rx)))
}
/// Run a single-connection bandwidth test with quota enforcement.
async fn run_pro_single_test(
stream: TcpStream,
cmd: Command,
peer: SocketAddr,
username: &str,
db: UserDb,
quota_mgr: QuotaManager,
quota_check_interval: u64,
) -> anyhow::Result<(StopReason, u64, u64)> {
let proto_str = if cmd.is_udp() { "UDP" } else { "TCP" };
let dir_str = match cmd.direction {
CMD_DIR_RX => "RX", CMD_DIR_TX => "TX", _ => "BOTH"
};
let session_id = db.start_session(
&username, &peer.ip().to_string(), proto_str, dir_str,
username, &peer.ip().to_string(), proto_str, dir_str,
)?;
btest_rs::syslog_logger::test_start(
&peer.to_string(), proto_str, dir_str, cmd.tcp_conn_count,
);
// Create shared bandwidth state for the test
let state = BandwidthState::new();
// Spawn quota enforcer
// Set byte budget
let budget = quota_mgr.remaining_budget(username, &peer.ip());
if budget < u64::MAX {
state.set_budget(budget);
tracing::info!("Byte budget for '{}' from {}: {} bytes", username, peer.ip(), budget);
}
let enforcer = QuotaEnforcer::new(
quota_mgr.clone(),
username.clone(),
username.to_string(),
peer.ip(),
state.clone(),
quota_check_interval,
quota_mgr.max_duration(),
);
// Spawn quota enforcer — runs alongside the test
let enforcer_state = state.clone();
let enforcer_handle = tokio::spawn(async move {
enforcer.run().await
});
// Run the actual bandwidth test using the standard server handlers.
// The enforcer runs concurrently and will set state.running = false
// if any quota is exceeded, which gracefully stops the TX/RX loops.
static UDP_PORT_OFFSET: std::sync::atomic::AtomicU16 = std::sync::atomic::AtomicU16::new(0);
let mut stream_mut = stream;
let test_result = if cmd.is_udp() {
let offset = UDP_PORT_OFFSET.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let udp_port = btest_rs::protocol::BTEST_UDP_PORT_START + offset;
btest_rs::server::run_udp_test(
&mut stream, peer, &cmd, state.clone(), udp_port,
&mut stream_mut, peer, &cmd, state.clone(), udp_port,
).await
} else {
btest_rs::server::run_tcp_test(stream, cmd.clone(), state.clone()).await
btest_rs::server::run_tcp_test(stream_mut, cmd.clone(), state.clone()).await
};
// Test finished — stop the enforcer if still running
enforcer_state.running.store(false, std::sync::atomic::Ordering::SeqCst);
let stop_reason = enforcer_handle.await.unwrap_or(StopReason::ClientDisconnected);
// Determine final stop reason
let final_reason = match &test_result {
Ok(_) => {
if stop_reason == StopReason::ClientDisconnected {
StopReason::ClientDisconnected
} else {
stop_reason // quota or duration exceeded
stop_reason
}
}
Err(_) => StopReason::ClientDisconnected,
};
// Record final usage
let (total_tx, total_rx, _, _) = state.summary();
// Flush to DB
quota_mgr.record_usage(&username, &peer.ip().to_string(), total_tx, total_rx);
quota_mgr.record_usage(username, &peer.ip().to_string(), total_tx, total_rx);
db.end_session(session_id, total_tx, total_rx)?;
Ok((username, final_reason, total_tx, total_rx))
Ok((final_reason, total_tx, total_rx))
}
/// Run a TCP multi-connection test with all streams collected.
/// Delegates to the standard multi-conn handler which correctly manages
/// TX+status injection for bidirectional mode.
async fn run_pro_multiconn_test(
streams: Vec<TcpStream>,
cmd: Command,
peer: SocketAddr,
username: &str,
db: UserDb,
quota_mgr: QuotaManager,
quota_check_interval: u64,
) -> anyhow::Result<(StopReason, u64, u64)> {
let dir_str = match cmd.direction {
CMD_DIR_RX => "RX", CMD_DIR_TX => "TX", _ => "BOTH"
};
let session_id = db.start_session(
username, &peer.ip().to_string(), "TCP", dir_str,
)?;
tracing::info!(
"Starting TCP multi-conn test: {} streams, dir={}",
streams.len(), dir_str,
);
let state = BandwidthState::new();
let budget = quota_mgr.remaining_budget(username, &peer.ip());
if budget < u64::MAX {
state.set_budget(budget);
}
let enforcer = QuotaEnforcer::new(
quota_mgr.clone(),
username.to_string(),
peer.ip(),
state.clone(),
quota_check_interval,
quota_mgr.max_duration(),
);
let enforcer_state = state.clone();
let enforcer_handle = tokio::spawn(async move {
enforcer.run().await
});
// Use the standard multi-connection handler which correctly handles
// all direction modes (TX, RX, BOTH with status injection)
let _test_result = btest_rs::server::run_tcp_multiconn_test(
streams, cmd, state.clone(),
).await;
enforcer_state.running.store(false, std::sync::atomic::Ordering::SeqCst);
let stop_reason = enforcer_handle.await.unwrap_or(StopReason::ClientDisconnected);
let (total_tx, total_rx, _, _) = state.summary();
quota_mgr.record_usage(username, &peer.ip().to_string(), total_tx, total_rx);
db.end_session(session_id, total_tx, total_rx)?;
Ok((stop_reason, total_tx, total_rx))
}