use rand::Rng;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use std::time::Duration;
pub mod fault_transport;
#[derive(Clone)]
pub struct NetworkSimulator {
inner: Arc<NetworkSimulatorInner>,
}
struct NetworkSimulatorInner {
udp_enabled: AtomicBool,
tcp_enabled: AtomicBool,
latency_ms: AtomicU32,
packet_loss_percent: AtomicU32,
bandwidth_limit: AtomicU32,
ip_change_counter: AtomicU32,
}
impl NetworkSimulator {
pub fn new() -> Self {
Self {
inner: Arc::new(NetworkSimulatorInner {
udp_enabled: AtomicBool::new(true),
tcp_enabled: AtomicBool::new(true),
latency_ms: AtomicU32::new(0),
packet_loss_percent: AtomicU32::new(0),
bandwidth_limit: AtomicU32::new(0),
ip_change_counter: AtomicU32::new(0),
}),
}
}
pub fn is_udp_enabled(&self) -> bool {
self.inner.udp_enabled.load(Ordering::SeqCst)
}
pub fn set_udp_enabled(&self, enabled: bool) {
self.inner.udp_enabled.store(enabled, Ordering::SeqCst);
log::info!(
"NetworkSimulator: UDP {}",
if enabled { "enabled" } else { "disabled" }
);
}
pub fn disable_udp(&self) {
self.set_udp_enabled(false);
}
pub fn enable_udp(&self) {
self.set_udp_enabled(true);
}
pub fn is_tcp_enabled(&self) -> bool {
self.inner.tcp_enabled.load(Ordering::SeqCst)
}
pub fn set_tcp_enabled(&self, enabled: bool) {
self.inner.tcp_enabled.store(enabled, Ordering::SeqCst);
log::info!(
"NetworkSimulator: TCP {}",
if enabled { "enabled" } else { "disabled" }
);
}
pub fn latency(&self) -> Duration {
Duration::from_millis(self.inner.latency_ms.load(Ordering::SeqCst) as u64)
}
pub fn set_latency(&self, latency: Duration) {
let ms = latency.as_millis().min(u32::MAX as u128) as u32;
self.inner.latency_ms.store(ms, Ordering::SeqCst);
log::info!("NetworkSimulator: latency set to {}ms", ms);
}
pub async fn apply_latency(&self) {
let latency = self.latency();
if !latency.is_zero() {
tokio::time::sleep(latency).await;
}
}
pub fn packet_loss_percent(&self) -> u32 {
self.inner.packet_loss_percent.load(Ordering::SeqCst)
}
pub fn set_packet_loss(&self, percent: u32) {
let percent = percent.min(100);
self.inner
.packet_loss_percent
.store(percent, Ordering::SeqCst);
log::info!("NetworkSimulator: packet loss set to {}%", percent);
}
pub fn should_drop_packet(&self) -> bool {
let loss = self.packet_loss_percent();
if loss == 0 {
return false;
}
if loss >= 100 {
return true;
}
rand::rngs::OsRng.gen::<u32>() % 100 < loss
}
pub fn bandwidth_limit(&self) -> u32 {
self.inner.bandwidth_limit.load(Ordering::SeqCst)
}
pub fn set_bandwidth_limit(&self, bytes_per_sec: u32) {
self.inner
.bandwidth_limit
.store(bytes_per_sec, Ordering::SeqCst);
log::info!(
"NetworkSimulator: bandwidth limit set to {} B/s",
bytes_per_sec
);
}
pub fn ip_change_counter(&self) -> u32 {
self.inner.ip_change_counter.load(Ordering::SeqCst)
}
pub fn trigger_ip_change(&self) {
let old = self.inner.ip_change_counter.fetch_add(1, Ordering::SeqCst);
log::info!(
"NetworkSimulator: IP change triggered (counter: {} -> {})",
old,
old + 1
);
}
pub fn preset_high_latency(&self) {
self.set_latency(Duration::from_millis(200));
self.set_packet_loss(5);
}
pub fn preset_lossy(&self) {
self.set_latency(Duration::from_millis(50));
self.set_packet_loss(20);
}
pub fn preset_mobile_roaming(&self) {
self.set_latency(Duration::from_millis(100));
self.set_packet_loss(2);
self.trigger_ip_change();
}
pub fn reset(&self) {
self.set_udp_enabled(true);
self.set_tcp_enabled(true);
self.set_latency(Duration::ZERO);
self.set_packet_loss(0);
self.set_bandwidth_limit(0);
}
}
impl Default for NetworkSimulator {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for NetworkSimulator {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("NetworkSimulator")
.field("udp_enabled", &self.is_udp_enabled())
.field("tcp_enabled", &self.is_tcp_enabled())
.field("latency_ms", &self.latency().as_millis())
.field("packet_loss_percent", &self.packet_loss_percent())
.field("bandwidth_limit", &self.bandwidth_limit())
.field("ip_changes", &self.ip_change_counter())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_network_simulator_basics() {
let sim = NetworkSimulator::new();
assert!(sim.is_udp_enabled());
assert!(sim.is_tcp_enabled());
assert_eq!(sim.latency(), Duration::ZERO);
assert_eq!(sim.packet_loss_percent(), 0);
sim.disable_udp();
assert!(!sim.is_udp_enabled());
sim.set_latency(Duration::from_millis(100));
assert_eq!(sim.latency(), Duration::from_millis(100));
sim.set_packet_loss(10);
assert_eq!(sim.packet_loss_percent(), 10);
let c1 = sim.ip_change_counter();
sim.trigger_ip_change();
assert_eq!(sim.ip_change_counter(), c1 + 1);
sim.reset();
assert!(sim.is_udp_enabled());
assert_eq!(sim.latency(), Duration::ZERO);
}
#[test]
fn test_presets() {
let sim = NetworkSimulator::new();
sim.preset_high_latency();
assert_eq!(sim.latency(), Duration::from_millis(200));
assert_eq!(sim.packet_loss_percent(), 5);
sim.preset_lossy();
assert_eq!(sim.latency(), Duration::from_millis(50));
assert_eq!(sim.packet_loss_percent(), 20);
}
#[test]
fn test_packet_drop_decision() {
let sim = NetworkSimulator::new();
sim.set_packet_loss(0);
for _ in 0..100 {
assert!(!sim.should_drop_packet());
}
sim.set_packet_loss(100);
for _ in 0..100 {
assert!(sim.should_drop_packet());
}
}
}