Stats collection

.gitignore

@@ -1,2 +1,3 @@
 /target
 count.txt
+stats.txt

Cargo.lock

@@ -222,6 +222,7 @@ dependencies = [
  "axum-server",
  "chrono",
  "clap",
+ "flume",
  "hashbrown",
  "http-body",
  "itertools 0.14.0",
@@ -426,6 +427,18 @@ dependencies = [
  "windows-sys 0.59.0",
 ]
 
+[[package]]
+name = "flume"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "da0e4dd2a88388a1f4ccc7c9ce104604dab68d9f408dc34cd45823d5a9069095"
+dependencies = [
+ "futures-core",
+ "futures-sink",
+ "nanorand",
+ "spin",
+]
+
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -509,8 +522,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592"
 dependencies = [
  "cfg-if",
+ "js-sys",
  "libc",
  "wasi 0.11.0+wasi-snapshot-preview1",
+ "wasm-bindgen",
 ]
 
 [[package]]
@@ -782,6 +797,16 @@ version = "0.4.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab"
 
+[[package]]
+name = "lock_api"
+version = "0.4.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17"
+dependencies = [
+ "autocfg",
+ "scopeguard",
+]
+
 [[package]]
 name = "log"
 version = "0.4.27"
@@ -832,6 +857,15 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "nanorand"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a51313c5820b0b02bd422f4b44776fbf47961755c74ce64afc73bfad10226c3"
+dependencies = [
+ "getrandom 0.2.16",
+]
+
 [[package]]
 name = "nom"
 version = "7.1.3"
@@ -1077,6 +1111,12 @@ version = "1.0.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
 
+[[package]]
+name = "scopeguard"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
+
 [[package]]
 name = "serde"
 version = "1.0.219"
@@ -1162,6 +1202,15 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "spin"
+version = "0.9.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67"
+dependencies = [
+ "lock_api",
+]
+
 [[package]]
 name = "strsim"
 version = "0.11.1"

Cargo.toml

@@ -5,10 +5,11 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-axum = "0.8.3"
+axum = { version = "0.8.3", features = ["tokio"] }
 axum-server = { version = "0.7.2", features = ["tls-rustls"] }
 chrono = "0.4.40"
 clap = { version = "4.5.37", features = ["derive"] }
+flume = "0.11.1"
 hashbrown = "0.15.2"
 http-body = "1.0.1"
 itertools = "0.14.0"

src/main.rs

@@ -2,7 +2,7 @@ mod generator;
 
 use axum::{
     body::{Body, Bytes},
-    extract::Path,
+    extract::{ConnectInfo, Path},
     response::{Html, IntoResponse, Response},
     routing::get,
     Router,
@@ -17,6 +17,7 @@ use rand_chacha::{rand_core::SeedableRng, ChaCha8Rng};
 use std::{
     borrow::Cow,
     convert::Infallible,
+    net::SocketAddr,
     path::PathBuf,
     pin::Pin,
     sync::{
@@ -53,9 +54,10 @@ fn create_rng(seed_bytes: impl IntoIterator<Item = u8>) -> Rng {
 }
 
 const COUNT_FILE: &str = "count.txt";
+const STATS_FILE: &str = "stats.txt";
 
-const SLOW_CHUNK_SIZE: usize = 50;
+const SLOW_CHUNK_SIZE: usize = 100;
-const SLOW_DURATION: Duration = Duration::from_millis(250);
+const SLOW_DURATION: Duration = Duration::from_millis(100);
 
 struct SlowBody {
     bytes: Bytes,
@@ -90,6 +92,10 @@ impl IntoResponse for SlowBody {
     }
 }
 
+struct RequestStats {
+    sock: SocketAddr,
+}
+
 #[tokio::main]
 async fn main() {
     let args = Args::parse();
@@ -110,50 +116,56 @@ async fn main() {
         },
     ));
 
+    let (stats_tx, stats_rx) = flume::unbounded();
+
     let app = {
         let counter = counter.clone();
+        let stats_tx = stats_tx.clone();
         Router::new().route(
             "/{id}",
-            get(|Path(id): Path<String>| async move {
+            get(
-                // Create a RNG for this path (deterministic, to simulate static pages)
+                |Path(id): Path<String>, ConnectInfo(sock): ConnectInfo<SocketAddr>| async move {
-                let mut rng = create_rng(id.bytes());
+                    // Create a RNG for this path (deterministic, to simulate static pages)
+                    let mut rng = create_rng(id.bytes());
 
-                // Count the request. Also doubles as the non-deterministic seed
+                    stats_tx.send(RequestStats { sock }).unwrap();
-                let count = counter.fetch_add(1, Ordering::Relaxed);
 
-                // Create a RNG for this session (non-deterministic)
+                    // Count the request. Also doubles as the non-deterministic seed
-                let mut session_rng = create_rng(count.to_le_bytes());
+                    let count = counter.fetch_add(1, Ordering::Relaxed);
 
-                // Artificially slow down connections as rudimentary DDoS protection, and to use up client resources
+                    // Create a RNG for this session (non-deterministic)
-                tokio::time::sleep(Duration::from_millis(session_rng.random_range(200..1000)))
+                    let mut session_rng = create_rng(count.to_le_bytes());
-                    .await;
 
-                // Choose a bullshit generator from our collection for this page
+                    // Artificially slow down connections as rudimentary DDoS protection, and to use up client resources
-                let generator = generators.choose(&mut rng).unwrap();
+                    tokio::time::sleep(Duration::from_millis(session_rng.random_range(200..1000)))
+                        .await;
 
-                let title = generator
+                    // Choose a bullshit generator from our collection for this page
-                    .word_stream(rng.random_range(2..10), &mut rng.clone())
+                    let generator = generators.choose(&mut rng).unwrap();
-                    .join(" ");
 
-                let stats = format!("Served rubbish to {count} clients so far");
+                    let title = generator
+                        .word_stream(rng.random_range(2..10), &mut rng.clone())
+                        .join(" ");
 
-                let content = generator
+                    let stats = format!("Served rubbish to {count} clients so far");
-                    .word_stream(rng.random_range(50..5_000), &mut rng.clone())
-                    .fold(String::new(), |mut content, word| {
-                        // Small chance of every word becoming a link back into the void
-                        if rng.random_bool(0.05) {
-                            let url = generator.word_stream(3, &mut rng.clone()).join("-");
-                            content += &format!(" <a href=\"{}\">{}</a>", url, word);
-                        } else {
-                            // Also, a chance for every word to end with a newline. This should probably be controlled by the generator.
-                            content += if rng.random_bool(0.01) { ".<br>" } else { " " };
-                            content += &word
-                        }
-                        content
-                    });
 
-                let html = format!(
+                    let content = generator
-                    "<!DOCTYPE html>
+                        .word_stream(rng.random_range(50..5_000), &mut rng.clone())
+                        .fold(String::new(), |mut content, word| {
+                            // Small chance of every word becoming a link back into the void
+                            if rng.random_bool(0.05) {
+                                let url = generator.word_stream(3, &mut rng.clone()).join("-");
+                                content += &format!(" <a href=\"{}\">{}</a>", url, word);
+                            } else {
+                                // Also, a chance for every word to end with a newline. This should probably be controlled by the generator.
+                                content += if rng.random_bool(0.01) { ".<br>" } else { " " };
+                                content += &word
+                            }
+                            content
+                        });
+
+                    let html = format!(
+                        "<!DOCTYPE html>
                     <html>
                     <head>
                     <title>{title}</title>
@@ -166,32 +178,41 @@ async fn main() {
 
                     </body>
                     </html>"
-                );
+                    );
 
-                SlowBody {
+                    SlowBody {
-                    bytes: html.into(),
+                        bytes: html.into(),
-                    interval: interval(SLOW_DURATION),
+                        interval: interval(SLOW_DURATION),
-                }
+                    }
-            }),
+                },
+            ),
         )
     };
 
     let mut interval = tokio::time::interval(Duration::from_secs(20));
+    let mut worst_offenders = HashMap::<_, u64>::default();
     tokio::spawn(async move {
         let mut last = 0;
         loop {
             interval.tick().await;
 
-            let count = counter.load(Ordering::Relaxed);
+            while let Ok(stats) = stats_rx.try_recv() {
+                *worst_offenders.entry(stats.sock.ip()).or_default() += 1;
+            }
 
+            let count = counter.load(Ordering::Relaxed);
             if count != last {
                 last = count;
-                println!(
-                    "{} Served bollocks to {} clients!",
-                    chrono::offset::Local::now(),
-                    count,
-                );
                 let _ = std::fs::write(COUNT_FILE, &format!("{count}"));
+
+                let mut worst_offenders = worst_offenders.iter().collect::<Vec<_>>();
+                worst_offenders.sort_by_key(|(_, n)| *n);
+                let stats = worst_offenders
+                    .iter()
+                    .enumerate()
+                    .map(|(i, (sock, n))| format!("#{:>4} | {:>4} | {}", i + 1, n, sock))
+                    .join("\n");
+                let _ = std::fs::write(STATS_FILE, &stats);
             }
         }
     });
@@ -207,13 +228,13 @@ async fn main() {
         println!("Enabling TLS...");
         let config = RustlsConfig::from_pem_file(cert, key).await.unwrap();
         bind_rustls(sock, config)
-            .serve(app.into_make_service())
+            .serve(app.into_make_service_with_connect_info::<SocketAddr>())
             .await
             .unwrap();
     } else {
         println!("WARNING: TLS disabled.");
         axum_server::bind(sock)
-            .serve(app.into_make_service())
+            .serve(app.into_make_service_with_connect_info::<SocketAddr>())
             .await
             .unwrap()
     }