From 39ae7cd3dc84617fc7b2ca786e022cf24e59b0e6 Mon Sep 17 00:00:00 2001 From: gigirassy Date: Fri, 19 Jun 2026 18:23:09 +0200 Subject: [PATCH] Add src/main.rs --- src/main.rs | 229 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 229 insertions(+) create mode 100644 src/main.rs diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..7f6543c --- /dev/null +++ b/src/main.rs @@ -0,0 +1,229 @@ +#![forbid(unsafe_code)] + +use mimalloc::MiMalloc; +use std::{ + env, + hash::{Hash, Hasher}, + io, + net::SocketAddr, + sync::Arc, + time::{SystemTime, UNIX_EPOCH}, +}; +use tokio::{ + io::AsyncWriteExt, + net::{TcpListener, TcpStream}, + sync::{OwnedSemaphorePermit, Semaphore}, + time::{self, Duration, Instant, MissedTickBehavior}, +}; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +#[derive(Clone, Copy)] +struct Config { + delay: Duration, + min_line_len: usize, + max_line_len: usize, + max_connections: usize, +} + +impl Config { + fn from_env() -> Self { + let delay_ms = env_usize("TARPIT_DELAY_MS", 10_000); + let min_line_len = env_usize("TARPIT_MIN_LINE", 8); + let max_line_len = env_usize("TARPIT_MAX_LINE", 64); + let max_connections = env_usize("TARPIT_MAX_CONNECTIONS", 4096); + + let min_line_len = min_line_len.max(1); + let max_line_len = max_line_len.max(min_line_len); + + Self { + delay: Duration::from_millis(delay_ms as u64), + min_line_len, + max_line_len, + max_connections: max_connections.max(1), + } + } + + fn validate(&self) -> io::Result<()> { + if self.min_line_len > self.max_line_len { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "TARPIT_MIN_LINE must be <= TARPIT_MAX_LINE", + )); + } + + if self.max_connections == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "TARPIT_MAX_CONNECTIONS must be > 0", + )); + } + + Ok(()) + } +} + +#[derive(Clone, Copy)] +struct XorShift64 { + state: u64, +} + +impl XorShift64 { + fn new(seed: u64) -> Self { + // Avoid the all-zero state. + Self { + state: if seed == 0 { 0x9e3779b97f4a7c15 } else { seed }, + } + } + + fn next_u64(&mut self) -> u64 { + let mut x = self.state; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.state = x; + x.wrapping_mul(0x2545F4914F6CDD1D) + } + + fn gen_range_inclusive(&mut self, start: usize, end_inclusive: usize) -> usize { + debug_assert!(start <= end_inclusive); + let span = (end_inclusive - start + 1) as u64; + (start as u64 + (self.next_u64() % span)) as usize + } +} + +fn seed_for_peer(peer: SocketAddr) -> u64 { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos() as u64) + .unwrap_or(0); + + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + peer.hash(&mut hasher); + + now ^ hasher.finish() +} + +#[tokio::main] +async fn main() -> io::Result<()> { + let bind_addr = env::var("TARPIT_BIND").unwrap_or_else(|_| "0.0.0.0:2222".to_string()); + let cfg = Config::from_env(); + cfg.validate()?; + + let listener = TcpListener::bind(&bind_addr).await?; + let limiter = Arc::new(Semaphore::new(cfg.max_connections)); + + eprintln!( + "listening on {bind_addr}, delay={}ms, line={}..{} bytes, max_conn={}", + cfg.delay.as_millis(), + cfg.min_line_len, + cfg.max_line_len, + cfg.max_connections + ); + + accept_loop(listener, limiter, cfg).await +} + +async fn accept_loop( + listener: TcpListener, + limiter: Arc, + cfg: Config, +) -> io::Result<()> { + let shutdown = shutdown_signal(); + tokio::pin!(shutdown); + + loop { + tokio::select! { + _ = &mut shutdown => { + eprintln!("shutdown requested"); + break; + } + + accepted = listener.accept() => { + let (stream, peer) = match accepted { + Ok(v) => v, + Err(e) => { + eprintln!("accept error: {e}"); + continue; + } + }; + + let permit = match limiter.clone().acquire_owned().await { + Ok(p) => p, + Err(_) => break, + }; + + tokio::spawn(async move { + if let Err(e) = handle_client(stream, peer, cfg, permit).await { + eprintln!("{peer}: {e}"); + } + }); + } + } + } + + Ok(()) +} + +async fn handle_client( + mut stream: TcpStream, + _peer: SocketAddr, + cfg: Config, + _permit: OwnedSemaphorePermit, +) -> io::Result<()> { + let _ = stream.set_nodelay(true); + + let mut rng = XorShift64::new(seed_for_peer(_peer)); + let mut buf = vec![0u8; cfg.max_line_len + 2]; + + let mut interval = time::interval_at(Instant::now() + cfg.delay, cfg.delay); + interval.set_missed_tick_behavior(MissedTickBehavior::Delay); + + loop { + interval.tick().await; + + let line_len = rng.gen_range_inclusive(cfg.min_line_len, cfg.max_line_len); + fill_line(&mut buf[..line_len], &mut rng); + buf[line_len] = b'\r'; + buf[line_len + 1] = b'\n'; + + stream.write_all(&buf[..line_len + 2]).await?; + } +} + +fn fill_line(dst: &mut [u8], rng: &mut XorShift64) { + const ALPHABET: &[u8] = b"abcdefghijklmnopqrstuvwxyz"; + + for byte in dst.iter_mut() { + let idx = rng.gen_range_inclusive(0, ALPHABET.len() - 1); + *byte = ALPHABET[idx]; + } +} + +async fn shutdown_signal() { + #[cfg(unix)] + { + use tokio::signal::unix::{signal, SignalKind}; + + let mut sigterm = signal(SignalKind::terminate()) + .expect("failed to install SIGTERM handler"); + + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + _ = sigterm.recv() => {} + } + } + + #[cfg(not(unix))] + { + let _ = tokio::signal::ctrl_c().await; + } +} + +fn env_usize(name: &str, default: usize) -> usize { + env::var(name) + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(default) +} \ No newline at end of file