From 4d928bd8def045b0a1ed43f389ad0b6dcd65ba37 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 11 Sep 2021 04:19:06 -0400 Subject: [PATCH] Rate limiter for TTS --- Cargo.lock | 309 ++++++++++++++++-- Cargo.toml | 2 + .../endpoints/tts/enqueue_infer_tts.rs | 8 + src/main.rs | 24 +- src/server_state.rs | 3 + src/util/redis/mod.rs | 3 +- src/util/redis/redis_rate_limiter.rs | 68 ++++ 7 files changed, 384 insertions(+), 33 deletions(-) create mode 100644 src/util/redis/redis_rate_limiter.rs diff --git a/Cargo.lock b/Cargo.lock index a93c002336..3c166d2b24 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -180,7 +180,7 @@ dependencies = [ "copyless", "futures-channel", "futures-util", - "smallvec", + "smallvec 1.6.1", "tokio 0.2.25", ] @@ -250,7 +250,7 @@ dependencies = [ "lazy_static", "log", "num_cpus", - "parking_lot", + "parking_lot 0.11.1", "threadpool", ] @@ -406,6 +406,12 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "23b62fc65de8e4e7f52534fb52b0f3ed04746ae267519eef2a83941e8085068b" +[[package]] +name = "ascii" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eab1c04a571841102f5345a8fc0f6bb3d31c315dec879b5c6e42e40ce7ffa34e" + [[package]] name = "async-channel" version = "1.6.1" @@ -495,7 +501,7 @@ dependencies = [ "async-global-executor", "async-io", "async-lock", - "crossbeam-utils", + "crossbeam-utils 0.8.4", "futures-channel", "futures-core", "futures-io", @@ -786,6 +792,16 @@ version = "1.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" +[[package]] +name = "bytes" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "206fdffcfa2df7cbe15601ef46c813fce0965eb3286db6b56c583b814b51c81c" +dependencies = [ + "byteorder", + "iovec", +] + [[package]] name = "bytes" version = "0.5.6" @@ -886,6 +902,28 @@ dependencies = [ "generic-array", ] +[[package]] +name = "cloudabi" +version = "0.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddfc5b9aa5d4507acaf872de71051dfd0e309860e88966e1051e462a077aac4f" +dependencies = [ + "bitflags", +] + +[[package]] +name = "combine" +version = "3.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3da6baa321ec19e1cc41d31bf599f00c783d0517095cdaf0332e3fe8d20680" +dependencies = [ + "ascii", + "byteorder", + "either", + "memchr", + "unreachable", +] + [[package]] name = "combine" version = "4.5.2" @@ -993,7 +1031,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06ed27e177f16d65f0f0c22a213e17c696ace5dd64b14258b52f9417ccb52db4" dependencies = [ "cfg-if 1.0.0", - "crossbeam-utils", + "crossbeam-utils 0.8.4", ] [[package]] @@ -1003,7 +1041,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0f6cb3c7f5b8e51bc3ebb73a2327ad4abdbd119dc13223f14f961d2f38486756" dependencies = [ "cfg-if 1.0.0", - "crossbeam-utils", + "crossbeam-utils 0.8.4", +] + +[[package]] +name = "crossbeam-utils" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3c7c73a2d1e9fc0886a08b93e98eb643461230d5f1925e4036204d5f2e261a8" +dependencies = [ + "autocfg 1.0.1", + "cfg-if 0.1.10", + "lazy_static", ] [[package]] @@ -1301,6 +1350,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fed34cd105917e91daa4da6b3728c47b068749d6a62c59811f06ed2ac71d9da7" +[[package]] +name = "futures" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a471a38ef8ed83cd6e40aa59c1ffe17db6855c18e3604d9c4ed8c08ebc28678" + [[package]] name = "futures" version = "0.3.15" @@ -1831,6 +1886,18 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7d73b3f436185384286bd8098d17ec07c9a7d2388a6599f824d8502b529702a" +[[package]] +name = "limitation" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e415413963930f7050f21994bfd336dc75ef8b5b3b142e99c55b72b61f726d5" +dependencies = [ + "chrono", + "futures 0.1.31", + "redis 0.13.0", + "time 0.1.44", +] + [[package]] name = "linked-hash-map" version = "0.5.4" @@ -1852,6 +1919,15 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "84f9a2d3e27ce99ce2c3aad0b09b1a7b916293ea9b2bf624c13fe646fadd8da4" +[[package]] +name = "lock_api" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4da24a77a3d8a6d4862d95f72e6fdb9c09a643ecdb402d754004a557f2bec75" +dependencies = [ + "scopeguard", +] + [[package]] name = "lock_api" version = "0.4.4" @@ -1915,6 +1991,12 @@ version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08" +[[package]] +name = "maybe-uninit" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60302e4db3a61da70c0cb7991976248362f30319e88850c487b9b95bbf059e00" + [[package]] name = "md-5" version = "0.9.1" @@ -2136,7 +2218,7 @@ dependencies = [ "num-traits", "rand 0.7.3", "serde", - "smallvec", + "smallvec 1.6.1", "zeroize", ] @@ -2241,6 +2323,17 @@ version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "427c3892f9e783d91cc128285287e70a59e206ca452770ece88a76f7a3eddd72" +[[package]] +name = "parking_lot" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f842b1982eb6c2fe34036a4fbfb06dd185a3f5c8edfaacdf7d1ea10b07de6252" +dependencies = [ + "lock_api 0.3.4", + "parking_lot_core 0.6.2", + "rustc_version", +] + [[package]] name = "parking_lot" version = "0.11.1" @@ -2248,8 +2341,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d7744ac029df22dca6284efe4e898991d28e3085c706c972bcd7da4a27a15eb" dependencies = [ "instant", - "lock_api", - "parking_lot_core", + "lock_api 0.4.4", + "parking_lot_core 0.8.3", +] + +[[package]] +name = "parking_lot_core" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b876b1b9e7ac6e1a74a6da34d25c42e17e8862aa409cbbbdcfc8d86c6f3bc62b" +dependencies = [ + "cfg-if 0.1.10", + "cloudabi", + "libc", + "redox_syscall 0.1.57", + "rustc_version", + "smallvec 0.6.14", + "winapi 0.3.9", ] [[package]] @@ -2261,8 +2369,8 @@ dependencies = [ "cfg-if 1.0.0", "instant", "libc", - "redox_syscall", - "smallvec", + "redox_syscall 0.2.8", + "smallvec 1.6.1", "winapi 0.3.9", ] @@ -2436,7 +2544,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "545c5bc2b880973c9c10e4067418407a0ccaa3091781d1671d46eb35107cb26f" dependencies = [ "log", - "parking_lot", + "parking_lot 0.11.1", "scheduled-thread-pool", ] @@ -2447,7 +2555,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "182473b876b0b93e353682ec58e207dd1cb4a62278bbe0045fe52b86b74363bb" dependencies = [ "r2d2", - "redis", + "redis 0.20.1", ] [[package]] @@ -2574,6 +2682,28 @@ dependencies = [ "rand_core 0.3.1", ] +[[package]] +name = "redis" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bb1079bc5692c03e1479bb7086ac870ac2aaf7115b5e72d314c01368f1a747e" +dependencies = [ + "bytes 0.4.12", + "combine 3.8.1", + "dtoa", + "futures 0.1.31", + "itoa", + "percent-encoding 2.1.0", + "sha1", + "tokio-codec", + "tokio-executor", + "tokio-io", + "tokio-sync", + "tokio-tcp", + "tokio-uds", + "url 2.2.2", +] + [[package]] name = "redis" version = "0.20.1" @@ -2581,7 +2711,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0a32cb439c4e89c1e6415e5b3b23d9d8cc6dc1bf5a8cade19adbd5418de803be" dependencies = [ "async-trait", - "combine", + "combine 4.5.2", "dtoa", "itoa", "percent-encoding 2.1.0", @@ -2589,6 +2719,12 @@ dependencies = [ "url 2.2.2", ] +[[package]] +name = "redox_syscall" +version = "0.1.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41cc0f7e4d5d4544e8861606a285bb08d3e70712ccc7d2b84d7c0ccfaf4b05ce" + [[package]] name = "redox_syscall" version = "0.2.8" @@ -2605,7 +2741,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "528532f3d801c87aec9def2add9ca802fe569e44a544afe633765267840abe64" dependencies = [ "getrandom 0.2.2", - "redox_syscall", + "redox_syscall 0.2.8", ] [[package]] @@ -2739,7 +2875,7 @@ dependencies = [ "base64 0.13.0", "cfg-if 1.0.0", "chrono", - "futures", + "futures 0.3.15", "hex", "hmac 0.9.0", "http", @@ -2788,7 +2924,7 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc6f74fd1204073fa02d5d5d68bec8021be4c38690b61264b2fdb48083d0e7d7" dependencies = [ - "parking_lot", + "parking_lot 0.11.1", ] [[package]] @@ -2974,6 +3110,15 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f173ac3d1a7e3b28003f40de0b5ce7fe2710f9b9dc3fc38664cebee46b3b6527" +[[package]] +name = "smallvec" +version = "0.6.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97fcaeba89edba30f044a10c6a3cc39df9c3f17d7cd829dd1446cab35f890e0" +dependencies = [ + "maybe-uninit", +] + [[package]] name = "smallvec" version = "1.6.1" @@ -3046,7 +3191,7 @@ dependencies = [ "crc", "crossbeam-channel", "crossbeam-queue", - "crossbeam-utils", + "crossbeam-utils 0.8.4", "digest", "either", "futures-channel", @@ -3061,14 +3206,14 @@ dependencies = [ "memchr", "num-bigint 0.3.2", "once_cell", - "parking_lot", + "parking_lot 0.11.1", "percent-encoding 2.1.0", "rand 0.7.3", "rsa", "serde", "sha-1", "sha2", - "smallvec", + "smallvec 1.6.1", "sqlformat", "sqlx-rt", "stringprep", @@ -3086,7 +3231,7 @@ dependencies = [ "cargo_metadata", "dotenv", "either", - "futures", + "futures 0.3.15", "heck", "hex", "lazy_static", @@ -3198,7 +3343,8 @@ dependencies = [ "diesel", "dotenv", "easyenv", - "futures", + "futures 0.1.31", + "futures 0.3.15", "hex", "hmac 0.10.1", "hostname", @@ -3207,6 +3353,7 @@ dependencies = [ "jwt", "lazy_static", "lfu", + "limitation", "log", "magic-crypt", "md-5", @@ -3303,7 +3450,7 @@ dependencies = [ "cfg-if 1.0.0", "libc", "rand 0.8.3", - "redox_syscall", + "redox_syscall 0.2.8", "remove_dir_all", "winapi 0.3.9", ] @@ -3454,12 +3601,44 @@ dependencies = [ "libc", "mio 0.7.11", "once_cell", - "parking_lot", + "parking_lot 0.11.1", "pin-project-lite 0.2.6", "signal-hook-registry", "winapi 0.3.9", ] +[[package]] +name = "tokio-codec" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25b2998660ba0e70d18684de5d06b70b70a3a747469af9dea7618cc59e75976b" +dependencies = [ + "bytes 0.4.12", + "futures 0.1.31", + "tokio-io", +] + +[[package]] +name = "tokio-executor" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb2d1b8f4548dbf5e1f7818512e9c406860678f29c300cdf0ebac72d1a3a1671" +dependencies = [ + "crossbeam-utils 0.7.2", + "futures 0.1.31", +] + +[[package]] +name = "tokio-io" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57fc868aae093479e3131e3d165c93b1c7474109d13c90ec0dda2a1bbfff0674" +dependencies = [ + "bytes 0.4.12", + "futures 0.1.31", + "log", +] + [[package]] name = "tokio-macros" version = "0.2.6" @@ -3481,6 +3660,49 @@ dependencies = [ "tokio 0.2.25", ] +[[package]] +name = "tokio-reactor" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09bc590ec4ba8ba87652da2068d150dcada2cfa2e07faae270a5e0409aa51351" +dependencies = [ + "crossbeam-utils 0.7.2", + "futures 0.1.31", + "lazy_static", + "log", + "mio 0.6.23", + "num_cpus", + "parking_lot 0.9.0", + "slab", + "tokio-executor", + "tokio-io", + "tokio-sync", +] + +[[package]] +name = "tokio-sync" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edfe50152bc8164fcc456dab7891fa9bf8beaf01c5ee7e1dd43a397c3cf87dee" +dependencies = [ + "fnv", + "futures 0.1.31", +] + +[[package]] +name = "tokio-tcp" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98df18ed66e3b72e742f185882a9e201892407957e45fbff8da17ae7a7c51f72" +dependencies = [ + "bytes 0.4.12", + "futures 0.1.31", + "iovec", + "mio 0.6.23", + "tokio-io", + "tokio-reactor", +] + [[package]] name = "tokio-tls" version = "0.3.1" @@ -3491,6 +3713,24 @@ dependencies = [ "tokio 0.2.25", ] +[[package]] +name = "tokio-uds" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab57a4ac4111c8c9dbcf70779f6fc8bc35ae4b2454809febac840ad19bd7e4e0" +dependencies = [ + "bytes 0.4.12", + "futures 0.1.31", + "iovec", + "libc", + "log", + "mio 0.6.23", + "mio-uds", + "tokio-codec", + "tokio-io", + "tokio-reactor", +] + [[package]] name = "tokio-util" version = "0.3.1" @@ -3560,12 +3800,12 @@ dependencies = [ "async-trait", "cfg-if 1.0.0", "enum-as-inner", - "futures", + "futures 0.3.15", "idna 0.2.3", "lazy_static", "log", "rand 0.7.3", - "smallvec", + "smallvec 1.6.1", "thiserror", "tokio 0.2.25", "url 2.2.2", @@ -3578,13 +3818,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "710f593b371175db53a26d0b38ed2978fafb9e9e8d3868b1acd753ea18df0ceb" dependencies = [ "cfg-if 0.1.10", - "futures", + "futures 0.3.15", "ipconfig", "lazy_static", "log", "lru-cache", "resolv-conf", - "smallvec", + "smallvec 1.6.1", "thiserror", "tokio 0.2.25", "trust-dns-proto", @@ -3675,6 +3915,15 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "unreachable" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "382810877fe448991dfc7f0dd6e3ae5d58088fd0ea5e35189655f84e6814fa56" +dependencies = [ + "void", +] + [[package]] name = "untrusted" version = "0.7.1" @@ -3773,6 +4022,12 @@ version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5fecdca9a5291cc2b8dcf7dc02453fee791a280f3743cb0905f8822ae463b3fe" +[[package]] +name = "void" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a02e4885ed3bc0f2de90ea6dd45ebcbb66dacffe03547fadbb0eeae2770887d" + [[package]] name = "waker-fn" version = "1.1.0" diff --git a/Cargo.toml b/Cargo.toml index 617394b102..e31f0c7e9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,6 +38,7 @@ derive_more = "0.99.14" dotenv = "0.15.0" easyenv = "0.0.5" futures = { version = "0.3.15", features = ["thread-pool"] } # feature for ThreadPool +futures_old_for_limiter = { package = "futures", version = "0.1.29" } # Older version of futures for `limiter` crate hex = "0.4.3" hostname = "0.3.1" lazy_static = "1.4.0" @@ -66,6 +67,7 @@ tokio = { version = "0.2.25", features = ["macros"] } sqlx = { version = "0.4", features = [ "mysql", "runtime-tokio-native-tls", "offline", "chrono" ] } # Redis Database +limitation = "0.1.1" # Redis rate limiter r2d2_redis = "0.14.0" # Files diff --git a/src/http_server/endpoints/tts/enqueue_infer_tts.rs b/src/http_server/endpoints/tts/enqueue_infer_tts.rs index 30eb0f307d..a41d83d1e9 100644 --- a/src/http_server/endpoints/tts/enqueue_infer_tts.rs +++ b/src/http_server/endpoints/tts/enqueue_infer_tts.rs @@ -44,6 +44,7 @@ pub struct InferTtsSuccessResponse { pub enum InferTtsError { BadInput(String), ServerError, + RateLimited, } impl ResponseError for InferTtsError { @@ -51,6 +52,7 @@ impl ResponseError for InferTtsError { match *self { InferTtsError::BadInput(_) => StatusCode::BAD_REQUEST, InferTtsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + InferTtsError::RateLimited => StatusCode::TOO_MANY_REQUESTS, } } @@ -58,6 +60,7 @@ impl ResponseError for InferTtsError { let error_reason = match self { InferTtsError::BadInput(reason) => reason.to_string(), InferTtsError::ServerError => "server error".to_string(), + InferTtsError::RateLimited => "rate limited".to_string(), }; to_simple_json_error(&error_reason, self.status_code()) @@ -69,6 +72,11 @@ pub async fn infer_tts_handler( request: web::Json, server_state: web::Data>) -> Result { + + if let Err(_err) = server_state.redis_rate_limiter.rate_limit_request(&http_request) { + return Err(InferTtsError::RateLimited); + } + let maybe_user_session = server_state .session_checker .maybe_get_user_session(&http_request, &server_state.mysql_pool) diff --git a/src/main.rs b/src/main.rs index 48d027e8e2..737633680e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,17 +36,21 @@ use crate::database::mediators::badge_granter::BadgeGranter; use crate::database::mediators::firehose_publisher::FirehosePublisher; use crate::http_server::endpoints::default_route_404::default_route_404; use crate::http_server::endpoints::events::list_events::list_events_handler; +use crate::http_server::endpoints::leaderboard::get_leaderboard::leaderboard_handler; use crate::http_server::endpoints::misc::enable_alpha_easy_handler::enable_alpha_easy_handler; use crate::http_server::endpoints::misc::enable_alpha_handler::enable_alpha_handler; use crate::http_server::endpoints::moderation::ip_bans::add_ip_ban::add_ip_ban_handler; use crate::http_server::endpoints::moderation::ip_bans::delete_ip_ban::delete_ip_ban_handler; use crate::http_server::endpoints::moderation::ip_bans::get_ip_ban::get_ip_ban_handler; use crate::http_server::endpoints::moderation::ip_bans::list_ip_bans::list_ip_bans_handler; +use crate::http_server::endpoints::moderation::jobs::get_tts_inference_queue_count::get_tts_inference_queue_count_handler; +use crate::http_server::endpoints::moderation::jobs::get_w2l_inference_queue_count::get_w2l_inference_queue_count_handler; use crate::http_server::endpoints::moderation::user_bans::ban_user::ban_user_handler; use crate::http_server::endpoints::moderation::user_bans::list_banned_users::list_banned_users_handler; use crate::http_server::endpoints::moderation::user_roles::list_roles::list_user_roles_handler; use crate::http_server::endpoints::moderation::user_roles::list_staff::list_staff_handler; use crate::http_server::endpoints::moderation::user_roles::set_user_role::set_user_role_handler; +use crate::http_server::endpoints::moderation::users::list_users::list_users_handler; use crate::http_server::endpoints::root_index::get_root_index; use crate::http_server::endpoints::tts::delete_tts_model::delete_tts_model_handler; use crate::http_server::endpoints::tts::delete_tts_result::delete_tts_inference_result_handler; @@ -95,6 +99,7 @@ use crate::threads::ip_banlist_set::IpBanlistSet; use crate::threads::poll_ip_banlist_thread::poll_ip_bans; use crate::util::buckets::bucket_client::BucketClient; use crate::util::encrypted_sort_id::SortKeyCrypto; +use crate::util::redis::redis_rate_limiter::RedisRateLimiter; use futures::Future; use futures::executor::ThreadPool; use log::{info}; @@ -105,10 +110,7 @@ use sqlx::MySqlPool; use sqlx::mysql::MySqlPoolOptions; use std::sync::Arc; use std::time::Duration; -use crate::http_server::endpoints::moderation::users::list_users::list_users_handler; -use crate::http_server::endpoints::leaderboard::get_leaderboard::leaderboard_handler; -use crate::http_server::endpoints::moderation::jobs::get_tts_inference_queue_count::get_tts_inference_queue_count_handler; -use crate::http_server::endpoints::moderation::jobs::get_w2l_inference_queue_count::get_w2l_inference_queue_count_handler; +use limitation::Limiter; // TODO TODO TODO TODO // TODO TODO TODO TODO @@ -184,11 +186,22 @@ async fn main() -> AnyhowResult<()> { firehose_publisher: firehose_publisher.clone(), // NB: Also safe }; - let redis_manager = RedisConnectionManager::new(redis_connection_string)?; + let redis_manager = RedisConnectionManager::new(redis_connection_string.clone())?; let redis_pool = r2d2::Pool::builder() .build(redis_manager)?; + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_MAX_REQUESTS", 3)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_WINDOW_SECONDS", 10)?; + + let limiter = Limiter::build(&redis_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + let redis_rate_limiter = RedisRateLimiter::new(limiter, limiter_enabled); + info!("Reading env vars and setting up utils..."); let bind_address = easyenv::get_env_string_or_default("BIND_ADDRESS", DEFAULT_BIND_ADDRESS); @@ -262,6 +275,7 @@ async fn main() -> AnyhowResult<()> { hostname: server_hostname, mysql_pool: pool, redis_pool, + redis_rate_limiter, firehose_publisher, badge_granter, cookie_manager, diff --git a/src/server_state.rs b/src/server_state.rs index 13c809a071..9e29cfb02b 100644 --- a/src/server_state.rs +++ b/src/server_state.rs @@ -5,6 +5,7 @@ use crate::http_server::web_utils::session_checker::SessionChecker; use crate::threads::ip_banlist_set::IpBanlistSet; use crate::util::buckets::bucket_client::BucketClient; use crate::util::encrypted_sort_id::SortKeyCrypto; +use crate::util::redis::redis_rate_limiter::RedisRateLimiter; use r2d2_redis::{r2d2, RedisConnectionManager}; use sqlx::MySqlPool; @@ -20,6 +21,8 @@ pub struct ServerState { pub redis_pool: r2d2::Pool, + pub redis_rate_limiter: RedisRateLimiter, + pub cookie_manager: CookieManager, pub session_checker: SessionChecker, diff --git a/src/util/redis/mod.rs b/src/util/redis/mod.rs index 7626fe14a6..fc70e62d66 100644 --- a/src/util/redis/mod.rs +++ b/src/util/redis/mod.rs @@ -1,2 +1,3 @@ pub mod redis_job_status_logger; -pub mod redis_keys; \ No newline at end of file +pub mod redis_keys; +pub mod redis_rate_limiter; \ No newline at end of file diff --git a/src/util/redis/redis_rate_limiter.rs b/src/util/redis/redis_rate_limiter.rs new file mode 100644 index 0000000000..f7442b403d --- /dev/null +++ b/src/util/redis/redis_rate_limiter.rs @@ -0,0 +1,68 @@ +use actix_http::http::{HeaderName, HeaderMap}; +use actix_web::HttpRequest; +use crate::http_server::web_utils::ip_address::get_request_ip; +use futures_old_for_limiter::Future; +use limitation::{Limiter, Error}; +use log::info; + +/// If this HTTP header is set, the rate limiter can be bypassed (eg. for debugging) +const RATE_LIMIT_BYPASS_HEADER : &'static str = "limitless"; + +#[derive(Clone)] // NB: Limiter is `Clone` +pub struct RedisRateLimiter { + limiter: Limiter, + limiter_enabled: bool, + rate_limit_bypass_header: HeaderName, +} + +pub enum RateLimiterError { + RateLimitExceededError, + ClientError, // NB: Fail open for this! +} + +impl RedisRateLimiter { + pub fn new(limiter: Limiter, limiter_enabled: bool) -> Self { + let rate_limit_bypass_header = HeaderName::from_static(RATE_LIMIT_BYPASS_HEADER); + RedisRateLimiter { + limiter, + limiter_enabled, + rate_limit_bypass_header, + } + } + + pub fn rate_limit_request(&self, request: &HttpRequest) -> Result<(), RateLimiterError> { + if !self.limiter_enabled { + return Ok(()) + } + + let headers = request.headers(); + if headers.contains_key(&self.rate_limit_bypass_header) { + info!("Bypassing rate limiter with special bypass/debug header."); + return Ok(()); + } + + let ip_address = get_request_ip(&request); + + let rate_limit_key = format!("rate_limit:{}", ip_address); + self.rate_limit_key(&rate_limit_key) + } + + pub fn rate_limit_key(&self, rate_limit_key: &str) -> Result<(), RateLimiterError> { + if !self.limiter_enabled { + return Ok(()) + } + + // NB: This library uses old-school futures (pre-async/await) + let permit = self.limiter.count(rate_limit_key); + + // TODO/FIXME: This could block. + match permit.wait() { + Ok(_) => Ok(()), + Err(err) => match err { + Error::Client(_) => Ok(()), // NB: Fail open for failure to connect to Redis + Error::Time(_) => Ok(()), // NB: Fail open for key parsing + Error::LimitExceeded(_) => Err(RateLimiterError::RateLimitExceededError), + }, + } + } +}