add rate limiting

This commit is contained in:
Adam 2024-10-09 20:59:09 -04:00
parent 786b198f4a
commit 4fb917eb08
3 changed files with 162 additions and 4 deletions

140
Cargo.lock generated
View file

@ -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"

View file

@ -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"] }

View file

@ -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())