From 4fb917eb0826377fecca5949135197cf4dfd18ad Mon Sep 17 00:00:00 2001 From: Adam Doyle Date: Wed, 9 Oct 2024 20:59:09 -0400 Subject: [PATCH] add rate limiting --- Cargo.lock | 140 ++++++++++++++++++++++++++++++++++++++++++++- server/Cargo.toml | 1 + server/src/main.rs | 25 +++++++- 3 files changed, 162 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f7f8b10..c7624e5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -236,7 +236,7 @@ dependencies = [ "sync_wrapper 1.0.1", "tokio", "tokio-tungstenite", - "tower", + "tower 0.5.1", "tower-layer", "tower-service", "tracing", @@ -616,6 +616,19 @@ dependencies = [ "syn", ] +[[package]] +name = "dashmap" +version = "5.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" +dependencies = [ + "cfg-if", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "dashmap" version = "6.1.0" @@ -768,6 +781,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "forwarded-header-value" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8835f84f38484cc86f110a805655697908257fb9a7af005234060891557198e9" +dependencies = [ + "nonempty", + "thiserror", +] + [[package]] name = "futures" version = "0.3.31" @@ -839,6 +862,12 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + [[package]] name = "futures-util" version = "0.3.31" @@ -932,6 +961,26 @@ dependencies = [ "web-sys", ] +[[package]] +name = "governor" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68a7f542ee6b35af73b06abc0dad1c1bae89964e4e253bc4b587b91c9637867b" +dependencies = [ + "cfg-if", + "dashmap 5.5.3", + "futures", + "futures-timer", + "no-std-compat", + "nonzero_ext", + "parking_lot", + "portable-atomic", + "quanta", + "rand", + "smallvec", + "spinning_top", +] + [[package]] name = "guardian" version = "1.2.0" @@ -1701,6 +1750,12 @@ version = "0.1.0-gamma2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "550cf708f0dd373c47d2d6092fc3c38826147afa57cf2f646ed7306b6b2a7f62" +[[package]] +name = "no-std-compat" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b93853da6d84c2e3c7d730d6473e8817692dd89be387eb01b94d7f108ecb5b8c" + [[package]] name = "nom" version = "7.1.3" @@ -1711,6 +1766,18 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nonempty" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9e591e719385e6ebaeb5ce5d3887f7d5676fceca6411d1925ccc95745f3d6f7" + +[[package]] +name = "nonzero_ext" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" + [[package]] name = "nu-ansi-term" version = "0.46.0" @@ -1924,6 +1991,12 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" +[[package]] +name = "portable-atomic" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc9c68a3f6da06753e9335d63e27f6b9754dd1920d941135b7ea8224f141adb2" + [[package]] name = "powerfmt" version = "0.2.0" @@ -2037,6 +2110,21 @@ dependencies = [ "yansi", ] +[[package]] +name = "quanta" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5167a477619228a0b284fac2674e3c388cba90631d7b7de620e6f1fcd08da5" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi", + "web-sys", + "winapi", +] + [[package]] name = "quote" version = "1.0.37" @@ -2098,6 +2186,15 @@ dependencies = [ "getrandom", ] +[[package]] +name = "raw-cpuid" +version = "11.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ab240315c661615f2ee9f0f2cd32d5a7343a84d5ebcccb99d46e6637565e7b0" +dependencies = [ + "bitflags", +] + [[package]] name = "reactive_graph" version = "0.1.0-gamma2" @@ -2328,8 +2425,9 @@ dependencies = [ "serde", "serde_json", "tokio", - "tower", + "tower 0.5.1", "tower-http", + "tower_governor", "tracing", "tracing-subscriber", "uuid", @@ -2343,7 +2441,7 @@ checksum = "08a98570188735ac32205912d3460d7e12cfcf969bcd7d2b01c4c4f795711e3e" dependencies = [ "bytes", "const_format", - "dashmap", + "dashmap 6.1.0", "futures", "gloo-net", "http", @@ -2475,6 +2573,15 @@ dependencies = [ "tungstenite", ] +[[package]] +name = "spinning_top" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d96d2d1d716fb500937168cc09353ffdc7a012be8475ac7308e1bdf0e3923300" +dependencies = [ + "lock_api", +] + [[package]] name = "strsim" version = "0.11.1" @@ -2787,6 +2894,17 @@ dependencies = [ "winnow", ] +[[package]] +name = "tower" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c" +dependencies = [ + "tower-layer", + "tower-service", + "tracing", +] + [[package]] name = "tower" version = "0.5.1" @@ -2842,6 +2960,22 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" +[[package]] +name = "tower_governor" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "313fa625fea5790ed56360a30ea980e41229cf482b4835801a67ef1922bf63b9" +dependencies = [ + "axum", + "forwarded-header-value", + "governor", + "http", + "pin-project", + "thiserror", + "tower 0.4.13", + "tracing", +] + [[package]] name = "tracing" version = "0.1.40" diff --git a/server/Cargo.toml b/server/Cargo.toml index e54b8dc..3f10263 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -17,6 +17,7 @@ serde_json = "1" tokio = { version = "1", features = ["full"] } tower = { version = "0", features = ["util"] } tower-http = { version = "0", features = ["fs", "trace", "compression-full"] } +tower_governor = { version = "0", features = ["axum", "tracing"] } tracing = "0" tracing-subscriber = { version = "0", features = ["env-filter"] } uuid = { version = "1", features = ["v7", "serde", "fast-rng"] } diff --git a/server/src/main.rs b/server/src/main.rs index 45a8e4f..67a53ed 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -9,9 +9,11 @@ use std::{ collections::{HashMap, HashSet}, net::SocketAddr, sync::{Arc, RwLock}, + time::Duration, }; use tokio::sync::{broadcast, mpsc}; use tower::ServiceBuilder; +use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer}; use tower_http::{compression::CompressionLayer, services::ServeDir}; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use user_handler::UserHandler; @@ -43,6 +45,24 @@ async fn main() -> Result<()> { bind_addr = String::from("127.0.0.1:3030"); } + // Set up rate limiting/governor + let governor_conf = Arc::new( + GovernorConfigBuilder::default() + .per_second(2) + .burst_size(5) + .finish() + .unwrap(), + ); + + let governor_limiter = governor_conf.limiter().clone(); + let interval = Duration::from_secs(60); + // a separate background task to clean up + std::thread::spawn(move || loop { + std::thread::sleep(interval); + tracing::info!("rate limiting storage size: {}", governor_limiter.len()); + governor_limiter.retain_recent(); + }); + // Set up state let (broadcast_tx, _rx) = broadcast::channel(1000); let (users_tx, mut users_rx) = mpsc::channel(1000); @@ -108,7 +128,10 @@ async fn main() -> Result<()> { .route("/websocket", get(websocket_connection_handler)) .nest_service("/", ServeDir::new("dist")) .layer(ServiceBuilder::new().layer(CompressionLayer::new())) - .with_state(app_state); + .with_state(app_state) + .layer(GovernorLayer { + config: governor_conf, + }); // send it let listener = tokio::net::TcpListener::bind(bind_addr.to_owned())