diff --git a/crates/service/web/storyteller_web/config/storyteller-web.development.env b/crates/service/web/storyteller_web/config/storyteller-web.development.env index 01ffcc593e..c890b5ef3d 100644 --- a/crates/service/web/storyteller_web/config/storyteller-web.development.env +++ b/crates/service/web/storyteller_web/config/storyteller-web.development.env @@ -5,6 +5,9 @@ COOKIE_SECURE=false COOKIE_DOMAIN=.jungle.horse WEBSITE_HOMEPAGE_REDIRECT=http://jungle.horse +# Test the rate limiter bypass (obviously these are dummy values) +AI_STREAMER_USERNAMES="PewDiePie,xQc" + # Stripe configs # All secrets are development and okay to check in STRIPE_CHECKOUT_SUCCESS_URL_PATH=/checkout_success diff --git a/crates/service/web/storyteller_web/src/configs/app_startup/mod.rs b/crates/service/web/storyteller_web/src/configs/app_startup/mod.rs new file mode 100644 index 0000000000..824be656d0 --- /dev/null +++ b/crates/service/web/storyteller_web/src/configs/app_startup/mod.rs @@ -0,0 +1,2 @@ +pub mod redis_rate_limiters; +pub mod username_set; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/configs/app_startup/redis_rate_limiters.rs b/crates/service/web/storyteller_web/src/configs/app_startup/redis_rate_limiters.rs new file mode 100644 index 0000000000..1dc2b4d5cd --- /dev/null +++ b/crates/service/web/storyteller_web/src/configs/app_startup/redis_rate_limiters.rs @@ -0,0 +1,102 @@ +use std::time::Duration; + +use log::info; + +use config::common_env::CommonEnv; +use errors::AnyhowResult; +use limitation::Limiter; +use crate::configs::app_startup::username_set::UsernameSet; + +use crate::http_server::web_utils::redis_rate_limiter::RedisRateLimiter; +use crate::server_state::RedisRateLimiters; + + +/// Build the various rate limiters +pub fn configure_redis_rate_limiters(common_env: &CommonEnv) -> AnyhowResult { + // TODO(bt,2023-12-13): Kill CommonEnv. + + info!("Setting up Redis rate limiters..."); + + let logged_out_redis_rate_limiter = { + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_OUT_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_OUT_MAX_REQUESTS", 3)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_OUT_WINDOW_SECONDS", 10)?; + + let limiter = Limiter::build(&common_env.redis_0_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + RedisRateLimiter::new(limiter, "logged_out", limiter_enabled) + }; + + let logged_in_redis_rate_limiter = { + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_IN_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_IN_MAX_REQUESTS", 3)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_IN_WINDOW_SECONDS", 10)?; + + let limiter = Limiter::build(&common_env.redis_0_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + RedisRateLimiter::new(limiter, "logged_in", limiter_enabled) + }; + + let api_high_priority_redis_rate_limiter = { + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_API_HIGH_PRIORITY_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_API_HIGH_PRIORITY_MAX_REQUESTS", 30)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_API_HIGH_PRIORITY_WINDOW_SECONDS", 30)?; + + let limiter = Limiter::build(&common_env.redis_0_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + RedisRateLimiter::new(limiter, "api_high_priority", limiter_enabled) + }; + + // This is for the following users: + // https://www.notion.so/storytellerai/dd88609558a24196b4ddeeef6079da98 + let api_ai_streamers_redis_rate_limiter = { + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_API_AI_STREAMERS_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_API_AI_STREAMERS_MAX_REQUESTS", 30)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_API_AI_STREAMERS_WINDOW_SECONDS", 30)?; + + let limiter = Limiter::build(&common_env.redis_0_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + RedisRateLimiter::new(limiter, "api_ai_streamers", limiter_enabled) + }; + + let ai_streamer_usernames = + UsernameSet::from_comma_separated(&easyenv::get_env_string_or_default("AI_STREAMER_USERNAMES", "")); + + info!("AI Streamers that can bypass the normal rate limiter ({}): {:?}", + ai_streamer_usernames.len(), + ai_streamer_usernames.list_names()); + + let model_upload_rate_limiter = { + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_MODEL_UPLOAD_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_MAX_REQUESTS", 3)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_WINDOW_SECONDS", 10)?; + + let limiter = Limiter::build(&common_env.redis_0_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + RedisRateLimiter::new(limiter, "model_upload", limiter_enabled) + }; + + Ok(RedisRateLimiters { + logged_out: logged_out_redis_rate_limiter, + logged_in: logged_in_redis_rate_limiter, + api_high_priority: api_high_priority_redis_rate_limiter, + api_ai_streamers: api_ai_streamers_redis_rate_limiter, + api_ai_streamer_username_set: ai_streamer_usernames, + model_upload: model_upload_rate_limiter, + }) +} diff --git a/crates/service/web/storyteller_web/src/configs/app_startup/username_set.rs b/crates/service/web/storyteller_web/src/configs/app_startup/username_set.rs new file mode 100644 index 0000000000..f9858a201b --- /dev/null +++ b/crates/service/web/storyteller_web/src/configs/app_startup/username_set.rs @@ -0,0 +1,102 @@ +use std::collections::{BTreeSet, HashSet}; + +#[derive(Clone)] +pub struct UsernameSet { + /// "usernames" in our service are always lowercase. + /// "display names" are case sensitive. + usernames: HashSet, +} + +impl UsernameSet { + pub fn from_comma_separated(comma_separated_list: &str) -> Self { + let usernames = comma_separated_list.split(",") + .into_iter() + .map(|item| item.trim()) + .map(|item| item.to_lowercase()) + .filter(|item| !item.is_empty()) + .collect::>(); + Self { + usernames + } + } + + pub fn len(&self) -> usize { + self.usernames.len() + } + + pub fn list_names(&self) -> BTreeSet<&str> { + // NB: BTreeSet for sorting, which is good for logging at app startup + self.usernames.iter().map(|item| item.as_str()).collect::<_>() + } + + pub fn username_is_in_set(&self, username: &str) -> bool { + // We lowercase the input username just in case a display name is passed + let username = username.to_lowercase(); + self.usernames.contains(&username) + } +} + +#[cfg(test)] +mod tests { + use crate::configs::app_startup::username_set::UsernameSet; + fn build_set() -> UsernameSet { + UsernameSet::from_comma_separated("bob,john,ALICE, BoB,jOhn, ,,, , aLiCe, joHn, bob,Alice") + } + #[test] + fn test_deduplication() { + let set = build_set(); + assert_eq!(set.len(), 3); + } + + #[test] + fn test_in_set() { + let set = build_set(); + assert!(set.username_is_in_set("bob")); + assert!(set.username_is_in_set("alice")); + assert!(set.username_is_in_set("john")); + } + + #[test] + fn test_in_set_case_insensitive() { + let set = build_set(); + + assert!(set.username_is_in_set("Bob")); + assert!(set.username_is_in_set("BOB")); + + assert!(set.username_is_in_set("AlIcE")); + assert!(set.username_is_in_set("ALICE")); + + assert!(set.username_is_in_set("jOHN")); + assert!(set.username_is_in_set("JOHN")); + } + + #[test] + fn test_not_in_set() { + let set = build_set(); + + // NB: The build_set() function might have attempted to introduce these, + // but the struct constructor should filter blank spaces out + assert!(!set.username_is_in_set("")); + assert!(!set.username_is_in_set(" ")); + + // Obviously not in set + assert!(!set.username_is_in_set("mary")); + assert!(!set.username_is_in_set("bill")); + + // These usernames are in the set, but not with padding + assert!(!set.username_is_in_set(" alice ")); + assert!(!set.username_is_in_set("bob ")); + assert!(!set.username_is_in_set(" john")); + } + + #[test] + fn test_username_return() { + let set = build_set(); + let usernames = set.list_names(); + assert_eq!(usernames.len(), 3); + assert!(usernames.contains(&"alice")); + assert!(usernames.contains(&"bob")); + assert!(usernames.contains(&"john")); + assert!(!usernames.contains(&"mary")); + } +} diff --git a/crates/service/web/storyteller_web/src/configs/mod.rs b/crates/service/web/storyteller_web/src/configs/mod.rs index 0a735a5b5c..806e1e3209 100644 --- a/crates/service/web/storyteller_web/src/configs/mod.rs +++ b/crates/service/web/storyteller_web/src/configs/mod.rs @@ -1,2 +1,3 @@ +pub mod app_startup; pub mod plans; pub mod static_api_tokens; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs index b01dd37a5d..cb417f5906 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs @@ -242,7 +242,12 @@ pub async fn enqueue_infer_tts_handler( // ==================== RATE LIMIT ==================== // + // TODO(bt,2023-12-13): Unit test this! if !disable_rate_limiter { + // For AI live streamers + // https://www.notion.so/storytellerai/dd88609558a24196b4ddeeef6079da98 + let mut is_ai_streamer = false; + let mut rate_limiter = match maybe_user_session { None => &server_state.redis_rate_limiters.logged_out, Some(ref user) => { @@ -250,6 +255,15 @@ pub async fn enqueue_infer_tts_handler( warn!("User is not authorized to use TTS model because they are banned."); return Err(InferTtsError::NotAuthorized); } + + if server_state + .redis_rate_limiters + .api_ai_streamer_username_set + .username_is_in_set(&user.user.username) + { + is_ai_streamer = true; + } + &server_state.redis_rate_limiters.logged_in }, }; @@ -264,6 +278,11 @@ pub async fn enqueue_infer_tts_handler( rate_limiter = &server_state.redis_rate_limiters.api_high_priority; } + if is_ai_streamer { + info!("Using AI streamer rate limiter"); + rate_limiter = &server_state.redis_rate_limiters.api_ai_streamers; + } + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { return Err(InferTtsError::RateLimited); } diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index e77ed85093..57f09a82c8 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -75,6 +75,7 @@ use crate::billing::internal_product_to_stripe_lookup_impl::InternalProductToStr use crate::billing::internal_session_cache_purge_impl::InternalSessionCachePurgeImpl; use crate::billing::stripe_internal_subscription_product_lookup_impl::StripeInternalSubscriptionProductLookupImpl; use crate::billing::stripe_internal_user_lookup_impl::StripeInternalUserLookupImpl; +use crate::configs::app_startup::redis_rate_limiters::configure_redis_rate_limiters; use crate::configs::static_api_tokens::{StaticApiTokenConfig, StaticApiTokens, StaticApiTokenSet}; use crate::http_server::middleware::pushback_filter_middleware::PushbackFilter; use crate::http_server::web_utils::redis_rate_limiter::RedisRateLimiter; @@ -173,65 +174,6 @@ async fn main() -> AnyhowResult<()> { easyenv::get_env_num("REDIS_CACHE_TTL_SECONDS", 60)?, ); - info!("Setting up Redis rate limiters..."); - - // Old env vars: - // - // "LIMITER_ENABLED" - // "LIMITER_MAX_REQUESTS" - // "LIMITER_WINDOW_SECONDS" - // - let logged_out_redis_rate_limiter = { - let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_OUT_ENABLED", true); - let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_OUT_MAX_REQUESTS", 3)?; - let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_OUT_WINDOW_SECONDS", 10)?; - - let limiter = Limiter::build(&common_env.redis_0_connection_string) - .limit(limiter_max_requests) - .period(Duration::from_secs(limiter_window_seconds)) - .finish()?; - - RedisRateLimiter::new(limiter, "logged_out", limiter_enabled) - }; - - let logged_in_redis_rate_limiter = { - let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_IN_ENABLED", true); - let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_IN_MAX_REQUESTS", 3)?; - let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_IN_WINDOW_SECONDS", 10)?; - - let limiter = Limiter::build(&common_env.redis_0_connection_string) - .limit(limiter_max_requests) - .period(Duration::from_secs(limiter_window_seconds)) - .finish()?; - - RedisRateLimiter::new(limiter, "logged_in", limiter_enabled) - }; - - let api_high_priority_redis_rate_limiter = { - let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_API_HIGH_PRIORITY_ENABLED", true); - let limiter_max_requests = easyenv::get_env_num("LIMITER_API_HIGH_PRIORITY_MAX_REQUESTS", 30)?; - let limiter_window_seconds = easyenv::get_env_num("LIMITER_API_HIGH_PRIORITY_WINDOW_SECONDS", 30)?; - - let limiter = Limiter::build(&common_env.redis_0_connection_string) - .limit(limiter_max_requests) - .period(Duration::from_secs(limiter_window_seconds)) - .finish()?; - - RedisRateLimiter::new(limiter, "api_high_priority", limiter_enabled) - }; - - let model_upload_rate_limiter = { - let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_MODEL_UPLOAD_ENABLED", true); - let limiter_max_requests = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_MAX_REQUESTS", 3)?; - let limiter_window_seconds = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_WINDOW_SECONDS", 10)?; - - let limiter = Limiter::build(&common_env.redis_0_connection_string) - .limit(limiter_max_requests) - .period(Duration::from_secs(limiter_window_seconds)) - .finish()?; - - RedisRateLimiter::new(limiter, "model_upload", limiter_enabled) - }; info!("Connecting to elasticsearch..."); @@ -466,12 +408,7 @@ async fn main() -> AnyhowResult<()> { elasticsearch, redis_pool, redis_ttl_cache, - redis_rate_limiters: RedisRateLimiters { - logged_out: logged_out_redis_rate_limiter, - logged_in: logged_in_redis_rate_limiter, - api_high_priority: api_high_priority_redis_rate_limiter, - model_upload: model_upload_rate_limiter, - }, + redis_rate_limiters: configure_redis_rate_limiters(&common_env)?, firehose_publisher, badge_granter, avt_cookie_manager, diff --git a/crates/service/web/storyteller_web/src/server_state.rs b/crates/service/web/storyteller_web/src/server_state.rs index 41ddadf1d3..a8f86133e2 100644 --- a/crates/service/web/storyteller_web/src/server_state.rs +++ b/crates/service/web/storyteller_web/src/server_state.rs @@ -20,6 +20,7 @@ use url_config::third_party_url_redirector::ThirdPartyUrlRedirector; use users_component::cookies::anonymous_visitor_tracking::avt_cookie_manager::AvtCookieManager; use users_component::cookies::session::session_cookie_manager::SessionCookieManager; use users_component::utils::session_checker::SessionChecker; +use crate::configs::app_startup::username_set::UsernameSet; use crate::http_server::endpoints::categories::tts::list_fully_computed_assigned_tts_categories::list_fully_computed_assigned_tts_categories::ModelTokensByCategoryToken; use crate::http_server::endpoints::leaderboard::get_leaderboard::LeaderboardInfo; @@ -140,6 +141,12 @@ pub struct RedisRateLimiters { /// constructions) pub api_high_priority: RedisRateLimiter, + /// API rate limit for AI streamers + pub api_ai_streamers: RedisRateLimiter, + + /// Usernames of AI streamers + pub api_ai_streamer_username_set: UsernameSet, + /// A rate limiter for TTS and W2L uploads pub model_upload: RedisRateLimiter, }