Add src/main.rs

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<Semaphore>,
+    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::<usize>().ok())
+        .unwrap_or(default)
+}