Stats collection
This commit is contained in:
parent
c7073c8fbe
commit
e21e5fffa5
4 changed files with 122 additions and 50 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -1,2 +1,3 @@
|
|||
/target
|
||||
count.txt
|
||||
stats.txt
|
||||
|
|
|
|||
49
Cargo.lock
generated
49
Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
119
src/main.rs
119
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_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<String>| async move {
|
||||
// Create a RNG for this path (deterministic, to simulate static pages)
|
||||
let mut rng = create_rng(id.bytes());
|
||||
get(
|
||||
|Path(id): Path<String>, ConnectInfo(sock): ConnectInfo<SocketAddr>| 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!(" <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 stats = format!("Served rubbish to {count} clients so far");
|
||||
|
||||
let html = format!(
|
||||
"<!DOCTYPE html>
|
||||
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!(" <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 {
|
||||
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::<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()
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue