daniel543x/artcraft
1
0 Fork 0
mirror of https://github.com/storytold/artcraft.git synced 2026-11-09 19:41:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

Rate limiter for TTS

Browse Source
This commit is contained in:
Brandon Thomas committed 2021-09-11 04:19:06 -04:00
1 parent 71dc65a5e4
commit 4d928bd8de
7 files changed
+384 -33
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+282 -27
View File
@@ -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"
Cargo.toml
+2
View File
@@ -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
src/http_server/endpoints/tts/enqueue_infer_tts.rs
+8
View File
@@ -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<InferTtsRequest>,
server_state: web::Data<Arc<ServerState>>) -> Result<HttpResponse, InferTtsError>
{
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)
src/main.rs
+19 -5
View File
@@ -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,
src/server_state.rs
+3
View File
@@ -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<RedisConnectionManager>,
pub redis_rate_limiter: RedisRateLimiter,
pub cookie_manager: CookieManager,
pub session_checker: SessionChecker,
src/util/redis/mod.rs
+2 -1
View File
@@ -1,2 +1,3 @@
pub mod redis_job_status_logger;
pub mod redis_keys;
pub mod redis_keys;
pub mod redis_rate_limiter;
src/util/redis/redis_rate_limiter.rs
+68
View File
@@ -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),
},
}
}
}