Update src/main.rs

src/main.rs

@@ -2,17 +2,21 @@
 
 use mimalloc::MiMalloc;
 use std::{
+    collections::HashMap,
     env,
     hash::{Hash, Hasher},
     io,
-    net::SocketAddr,
-    sync::Arc,
-    time::{SystemTime, UNIX_EPOCH},
+    net::{IpAddr, SocketAddr},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+    time::{Duration as StdDuration, SystemTime, UNIX_EPOCH},
 };
 use tokio::{
     io::AsyncWriteExt,
     net::{TcpListener, TcpStream},
-    sync::{OwnedSemaphorePermit, Semaphore},
+    sync::{Mutex, OwnedSemaphorePermit, Semaphore},
     time::{self, Duration, Instant, MissedTickBehavior},
 };
 
@@ -53,17 +57,15 @@ impl Config {
             ));
         }
 
-        if self.max_connections == 0 {
-            return Err(io::Error::new(
-                io::ErrorKind::InvalidInput,
-                "TARPIT_MAX_CONNECTIONS must be > 0",
-            ));
-        }
-
         Ok(())
     }
 }
 
+struct State {
+    connections: AtomicUsize,
+    totals: Mutex<HashMap<IpAddr, StdDuration>>,
+}
+
 #[derive(Clone, Copy)]
 struct XorShift64 {
     state: u64,
@@ -71,9 +73,12 @@ struct XorShift64 {
 
 impl XorShift64 {
     fn new(seed: u64) -> Self {
-        // Avoid the all-zero state.
         Self {
-            state: if seed == 0 { 0x9e3779b97f4a7c15 } else { seed },
+            state: if seed == 0 {
+                0x9e3779b97f4a7c15
+            } else {
+                seed
+            },
         }
     }
 
@@ -86,9 +91,8 @@ impl XorShift64 {
         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;
+    fn gen_range_inclusive(&mut self, start: usize, end: usize) -> usize {
+        let span = (end - start + 1) as u64;
         (start as u64 + (self.next_u64() % span)) as usize
     }
 }
@@ -108,12 +112,18 @@ fn seed_for_peer(peer: SocketAddr) -> u64 {
 #[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));
 
+    let state = Arc::new(State {
+        connections: AtomicUsize::new(0),
+        totals: Mutex::new(HashMap::new()),
+    });
+
     eprintln!(
         "listening on {bind_addr}, delay={}ms, line={}..{} bytes, max_conn={}",
         cfg.delay.as_millis(),
@@ -122,12 +132,13 @@ async fn main() -> io::Result<()> {
         cfg.max_connections
     );
 
-    accept_loop(listener, limiter, cfg).await
+    accept_loop(listener, limiter, state, cfg).await
 }
 
 async fn accept_loop(
     listener: TcpListener,
     limiter: Arc<Semaphore>,
+    state: Arc<State>,
     cfg: Config,
 ) -> io::Result<()> {
     let shutdown = shutdown_signal();
@@ -154,8 +165,10 @@ async fn accept_loop(
                     Err(_) => break,
                 };
 
+                let state = state.clone();
+
                 tokio::spawn(async move {
-                    if let Err(e) = handle_client(stream, peer, cfg, permit).await {
+                    if let Err(e) = handle_client(stream, peer, state, cfg, permit).await {
                         eprintln!("{peer}: {e}");
                     }
                 });
@@ -168,28 +181,70 @@ async fn accept_loop(
 
 async fn handle_client(
     mut stream: TcpStream,
-    _peer: SocketAddr,
+    peer: SocketAddr,
+    state: Arc<State>,
     cfg: Config,
     _permit: OwnedSemaphorePermit,
 ) -> io::Result<()> {
     let _ = stream.set_nodelay(true);
 
-    let mut rng = XorShift64::new(seed_for_peer(_peer));
+    let ip = peer.ip();
+
+    let current = state.connections.fetch_add(1, Ordering::SeqCst) + 1;
+
+    eprintln!(
+        "[+] {} connected ({} active)",
+        peer,
+        current,
+    );
+
+    let connected_at = Instant::now();
+
+    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 result = async {
+        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';
+            let line_len =
+                rng.gen_range_inclusive(cfg.min_line_len, cfg.max_line_len);
 
-        stream.write_all(&buf[..line_len + 2]).await?;
+            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?;
+        }
+
+        #[allow(unreachable_code)]
+        Ok::<(), io::Error>(())
     }
+    .await;
+
+    let elapsed = connected_at.elapsed();
+
+    let total = {
+        let mut totals = state.totals.lock().await;
+        let entry = totals.entry(ip).or_insert(StdDuration::ZERO);
+        *entry += elapsed;
+        *entry
+    };
+
+    let current = state.connections.fetch_sub(1, Ordering::SeqCst) - 1;
+
+    eprintln!(
+        "[-] {} disconnected ({} active) session={:.1?} total={:.1?}",
+        peer,
+        current,
+        elapsed,
+        total
+    );
+
+    result
 }
 
 fn fill_line(dst: &mut [u8], rng: &mut XorShift64) {