From e21e5fffa58ced1ab564d3da63a91d4537fd0f99 Mon Sep 17 00:00:00 2001 From: Joshua Barretto Date: Tue, 29 Apr 2025 20:37:34 +0100 Subject: [PATCH] Stats collection --- .gitignore | 1 + Cargo.lock | 49 ++++++++++++++++++++++ Cargo.toml | 3 +- src/main.rs | 119 ++++++++++++++++++++++++++++++---------------------- 4 files changed, 122 insertions(+), 50 deletions(-) diff --git a/.gitignore b/.gitignore index 9f59e66..a1f8c38 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ /target count.txt +stats.txt diff --git a/Cargo.lock b/Cargo.lock index c8c95c2..65482da 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 390adf7..3993f85 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/main.rs b/src/main.rs index 3bcd53f..4f499db 100644 --- a/src/main.rs +++ b/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) -> Rng { } const COUNT_FILE: &str = "count.txt"; +const STATS_FILE: &str = "stats.txt"; -const SLOW_CHUNK_SIZE: usize = 50; -const SLOW_DURATION: Duration = Duration::from_millis(250); +const SLOW_CHUNK_SIZE: usize = 100; +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| async move { - // Create a RNG for this path (deterministic, to simulate static pages) - let mut rng = create_rng(id.bytes()); + get( + |Path(id): Path, ConnectInfo(sock): ConnectInfo| async move { + // 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 - let count = counter.fetch_add(1, Ordering::Relaxed); + stats_tx.send(RequestStats { sock }).unwrap(); - // Create a RNG for this session (non-deterministic) - let mut session_rng = create_rng(count.to_le_bytes()); + // Count the request. Also doubles as the non-deterministic seed + let count = counter.fetch_add(1, Ordering::Relaxed); - // Artificially slow down connections as rudimentary DDoS protection, and to use up client resources - tokio::time::sleep(Duration::from_millis(session_rng.random_range(200..1000))) - .await; + // Create a RNG for this session (non-deterministic) + let mut session_rng = create_rng(count.to_le_bytes()); - // Choose a bullshit generator from our collection for this page - let generator = generators.choose(&mut rng).unwrap(); + // Artificially slow down connections as rudimentary DDoS protection, and to use up client resources + tokio::time::sleep(Duration::from_millis(session_rng.random_range(200..1000))) + .await; - let title = generator - .word_stream(rng.random_range(2..10), &mut rng.clone()) - .join(" "); + // Choose a bullshit generator from our collection for this page + let generator = generators.choose(&mut rng).unwrap(); - 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 - .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!(" {}", 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) { ".
" } else { " " }; - content += &word - } - content - }); + let stats = format!("Served rubbish to {count} clients so far"); - let html = format!( - " + let content = generator + .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!(" {}", 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) { ".
" } else { " " }; + content += &word + } + content + }); + + let html = format!( + " {title} @@ -166,32 +178,41 @@ async fn main() { " - ); + ); - SlowBody { - bytes: html.into(), - interval: interval(SLOW_DURATION), - } - }), + SlowBody { + bytes: html.into(), + 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::>(); + 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::()) .await .unwrap(); } else { println!("WARNING: TLS disabled."); axum_server::bind(sock) - .serve(app.into_make_service()) + .serve(app.into_make_service_with_connect_info::()) .await .unwrap() }