Rate limit per app/api-key on the authed endpoint (#405)

* configurable rate limits per authed app

* cleanup, resolve comments:

---------

Co-authored-by: Haardik H <[email protected]>

Author: avalonche

crates/websocket-proxy/src/auth.rs

@@ -1,19 +1,21 @@
 use crate::auth::AuthenticationParseError::{
     DuplicateAPIKeyArgument, DuplicateApplicationArgument, MissingAPIKeyArgument,
-    MissingApplicationArgument, NoData, TooManyComponents,
+    MissingApplicationArgument, MissingRateLimitArgument, NoData, TooManyComponents,
 };
 use std::collections::{HashMap, HashSet};
 
 #[derive(Clone, Debug)]
 pub struct Authentication {
     key_to_application: HashMap<String, String>,
+    app_to_rate_limit: HashMap<String, usize>,
 }
 
 #[derive(Debug, PartialEq)]
 pub enum AuthenticationParseError {
     NoData(),
     MissingApplicationArgument(String),
     MissingAPIKeyArgument(String),
+    MissingRateLimitArgument(String),
     TooManyComponents(String),
     DuplicateApplicationArgument(String),
     DuplicateAPIKeyArgument(String),
@@ -23,9 +25,10 @@ impl std::fmt::Display for AuthenticationParseError {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
             NoData() => write!(f, "No API Keys Provided"),
-            MissingApplicationArgument(arg) => write!(f, "Missing application argument: [{arg}]"),
-            MissingAPIKeyArgument(app) => write!(f, "Missing API Key argument: [{app}]"),
-            TooManyComponents(app) => write!(f, "Too many components: [{app}]"),
+            MissingApplicationArgument(arg) => write!(f, "Missing application argument: [{}]", arg),
+            MissingAPIKeyArgument(app) => write!(f, "Missing API Key argument: [{}]", app),
+            MissingRateLimitArgument(app) => write!(f, "Missing rate limit argument: [{}]", app),
+            TooManyComponents(app) => write!(f, "Too many components: [{}]", app),
             DuplicateApplicationArgument(app) => {
                 write!(f, "Duplicate application argument: [{app}]")
             }
@@ -42,6 +45,7 @@ impl TryFrom<Vec<String>> for Authentication {
     fn try_from(args: Vec<String>) -> Result<Self, Self::Error> {
         let mut applications = HashSet::new();
         let mut key_to_application: HashMap<String, String> = HashMap::new();
+        let mut app_to_rate_limit: HashMap<String, usize> = HashMap::new();
 
         if args.is_empty() {
             return Err(NoData());
@@ -61,6 +65,13 @@ impl TryFrom<Vec<String>> for Authentication {
                 return Err(MissingAPIKeyArgument(app.to_string()));
             }
 
+            let rate_limit = parts
+                .next()
+                .ok_or(MissingRateLimitArgument(app.to_string()))?;
+            if rate_limit.is_empty() {
+                return Err(MissingRateLimitArgument(app.to_string()));
+            }
+
             if parts.count() > 0 {
                 return Err(TooManyComponents(app.to_string()));
             }
@@ -75,29 +86,42 @@ impl TryFrom<Vec<String>> for Authentication {
 
             applications.insert(app.to_string());
             key_to_application.insert(key.to_string(), app.to_string());
+            app_to_rate_limit.insert(app.to_string(), rate_limit.parse().unwrap());
         }
 
-        Ok(Self { key_to_application })
+        Ok(Self {
+            key_to_application,
+            app_to_rate_limit,
+        })
     }
 }
 
 impl Authentication {
     pub fn none() -> Self {
         Self {
             key_to_application: HashMap::new(),
+            app_to_rate_limit: HashMap::new(),
         }
     }
 
     #[allow(dead_code)]
-    pub fn new(api_keys: HashMap<String, String>) -> Self {
+    pub fn new(
+        api_keys: HashMap<String, String>,
+        app_to_rate_limit: HashMap<String, usize>,
+    ) -> Self {
         Self {
             key_to_application: api_keys,
+            app_to_rate_limit,
         }
     }
 
     pub fn get_application_for_key(&self, api_key: &String) -> Option<&String> {
         self.key_to_application.get(api_key)
     }
+
+    pub fn get_rate_limits(&self) -> HashMap<String, usize> {
+        self.app_to_rate_limit.clone()
+    }
 }
 
 #[cfg(test)]
@@ -107,69 +131,90 @@ mod tests {
     #[test]
     fn test_parsing() {
         let auth = Authentication::try_from(vec![
-            "app1:key1".to_string(),
-            "app2:key2".to_string(),
-            "app3:key3".to_string(),
+            "app1:key1:10".to_string(),
+            "app2:key2:10".to_string(),
+            "app3:key3:10".to_string(),
         ])
         .unwrap();
 
         assert_eq!(auth.key_to_application.len(), 3);
         assert_eq!(auth.key_to_application["key1"], "app1");
         assert_eq!(auth.key_to_application["key2"], "app2");
         assert_eq!(auth.key_to_application["key3"], "app3");
+        assert_eq!(auth.app_to_rate_limit.len(), 3);
+        assert_eq!(auth.app_to_rate_limit["app1"], 10);
+        assert_eq!(auth.app_to_rate_limit["app2"], 10);
+        assert_eq!(auth.app_to_rate_limit["app3"], 10);
 
         let auth = Authentication::try_from(vec![
-            "app1:key1".to_string(),
+            "app1:key1:10".to_string(),
             "".to_string(),
-            "app3:key3".to_string(),
+            "app3:key3:10".to_string(),
         ]);
         assert!(auth.is_err());
         assert_eq!(auth.unwrap_err(), MissingApplicationArgument("".into()));
 
         let auth = Authentication::try_from(vec![
-            "app1:key1".to_string(),
+            "app1:key1:10".to_string(),
             "app2".to_string(),
-            "app3:key3".to_string(),
+            "app3:key3:10".to_string(),
         ]);
         assert!(auth.is_err());
         assert_eq!(auth.unwrap_err(), MissingAPIKeyArgument("app2".into()));
 
         let auth = Authentication::try_from(vec![
-            "app1:key1".to_string(),
-            ":".to_string(),
+            "app1:key1:10".to_string(),
+            "app2:key2:10".to_string(),
             "app3:key3".to_string(),
         ]);
         assert!(auth.is_err());
+        assert_eq!(auth.unwrap_err(), MissingRateLimitArgument("app3".into()));
+
+        let auth = Authentication::try_from(vec![
+            "app1:key1:10".to_string(),
+            ":".to_string(),
+            "app3:key3:10".to_string(),
+        ]);
+        assert!(auth.is_err());
         assert_eq!(auth.unwrap_err(), MissingApplicationArgument(":".into()));
 
         let auth = Authentication::try_from(vec![
-            "app1:key1".to_string(),
+            "app1:key1:10".to_string(),
             "app2:".to_string(),
-            "app3:key3".to_string(),
+            "app3:key3:10".to_string(),
         ]);
         assert!(auth.is_err());
         assert_eq!(auth.unwrap_err(), MissingAPIKeyArgument("app2".into()));
 
         let auth = Authentication::try_from(vec![
-            "app1:key1".to_string(),
-            "app2:key2:unexpected2".to_string(),
-            "app3:key3".to_string(),
+            "app1:key1:10".to_string(),
+            "app2:key2:10".to_string(),
+            "app3:key3:".to_string(),
+        ]);
+        assert!(auth.is_err());
+        assert_eq!(auth.unwrap_err(), MissingRateLimitArgument("app3".into()));
+
+        let auth = Authentication::try_from(vec![
+            "app1:key1:10".to_string(),
+            "app2:key2:10:unexpected2".to_string(),
+            "app3:key3:10".to_string(),
         ]);
         assert!(auth.is_err());
         assert_eq!(auth.unwrap_err(), TooManyComponents("app2".into()));
 
         let auth = Authentication::try_from(vec![
-            "app1:key1".to_string(),
-            "app1:key3".to_string(),
-            "app2:key2".to_string(),
+            "app1:key1:10".to_string(),
+            "app1:key3:10".to_string(),
+            "app2:key2:10".to_string(),
         ]);
         assert!(auth.is_err());
         assert_eq!(
             auth.unwrap_err(),
             DuplicateApplicationArgument("app1".into())
         );
 
-        let auth = Authentication::try_from(vec!["app1:key1".to_string(), "app2:key1".to_string()]);
+        let auth =
+            Authentication::try_from(vec!["app1:key1:10".to_string(), "app2:key1:10".to_string()]);
         assert!(auth.is_err());
         assert_eq!(auth.unwrap_err(), DuplicateAPIKeyArgument("app2".into()));
     }

crates/websocket-proxy/src/main.rs

@@ -15,6 +15,7 @@ use metrics_exporter_prometheus::PrometheusBuilder;
 use rate_limit::{InMemoryRateLimit, RateLimit, RedisRateLimit};
 use registry::Registry;
 use server::Server;
+use std::collections::HashMap;
 use std::io::Write;
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -66,9 +67,10 @@ struct Args {
         long,
         env,
         default_value = "10",
-        help = "Maximum number of concurrently connected clients per IP"
+        help = "Maximum number of concurrently connected clients per IP. 0 here means no limit."
     )]
     per_ip_connection_limit: usize,
+
     #[arg(
         long,
         env,
@@ -96,7 +98,7 @@ struct Args {
     #[arg(long, env, default_value = "true")]
     metrics: bool,
 
-    /// API Keys, if not provided will be an unauthenticated endpoint, should be in the format <app1>:<apiKey1>,<app2>:<apiKey2>,..
+    /// API Keys, if not provided will be an unauthenticated endpoint, should be in the format <app1>:<apiKey1>:<rateLimit1>,<app2>:<apiKey2>:<rateLimit2>,..
     #[arg(long, env, value_delimiter = ',', help = "API keys to allow")]
     api_keys: Vec<String>,
 
@@ -337,13 +339,20 @@ async fn main() {
         args.client_pong_timeout_ms,
     );
 
+    let app_rate_limits = if let Some(auth) = &authentication {
+        auth.get_rate_limits()
+    } else {
+        HashMap::new()
+    };
+
     let rate_limiter = match &args.redis_url {
         Some(redis_url) => {
             info!(message = "Using Redis rate limiter", redis_url = redis_url);
             match RedisRateLimit::new(
                 redis_url,
                 args.instance_connection_limit,
                 args.per_ip_connection_limit,
+                app_rate_limits.clone(),
                 &args.redis_key_prefix,
             ) {
                 Ok(limiter) => {
@@ -359,6 +368,7 @@ async fn main() {
                     Arc::new(InMemoryRateLimit::new(
                         args.instance_connection_limit,
                         args.per_ip_connection_limit,
+                        app_rate_limits.clone(),
                     )) as Arc<dyn RateLimit>
                 }
             }
@@ -368,6 +378,7 @@ async fn main() {
             Arc::new(InMemoryRateLimit::new(
                 args.instance_connection_limit,
                 args.per_ip_connection_limit,
+                app_rate_limits,
             )) as Arc<dyn RateLimit>
         }
     };