Refactor rate limiting

Author: danyalprout

.github/workflows/ci.yaml

@@ -40,9 +40,6 @@ jobs:
       - name: Check for common mistakes
         run: cargo check
 
-      - name: Check documentation
-        run: cargo doc --no-deps --document-private-items
-
   security-audit:
     name: Security Audit
     runs-on: ubuntu-latest

Cargo.lock

@@ -19,6 +19,8 @@ tokio-util = "0.7.12"
 reqwest = "0.12.12"
 metrics = "0.24.1"
 metrics-derive = "0.1"
+thiserror = "2.0.11"
+serde_json = "1.0.138"
 
 [features]
 integration = []

Cargo.toml

@@ -0,0 +1,28 @@
+use crate::rate_limit::Ticket;
+use axum::extract::ws::WebSocket;
+use axum::Error;
+use std::net::IpAddr;
+
+pub struct ClientConnection {
+    client_addr: IpAddr,
+    _ticket: Ticket,
+    websocket: WebSocket,
+}
+
+impl ClientConnection {
+    pub fn new(client_addr: IpAddr, ticket: Ticket, websocket: WebSocket) -> Self {
+        Self {
+            client_addr,
+            _ticket: ticket,
+            websocket,
+        }
+    }
+
+    pub async fn send(&mut self, data: String) -> Result<(), Error> {
+        self.websocket.send(data.into_bytes().into()).await
+    }
+
+    pub fn id(&self) -> String {
+        self.client_addr.to_string()
+    }
+}

src/client.rs

@@ -1,5 +1,6 @@
 mod test {
     use crate::metrics::Metrics;
+    use crate::rate_limit::InMemoryRateLimit;
     use crate::registry::Registry;
     use crate::server::Server;
     use futures::StreamExt;
@@ -36,14 +37,22 @@ mod test {
         }
         fn new(addr: SocketAddr) -> TestHarness {
             let (sender, _) = broadcast::channel(5);
-            let registry = Registry::new(sender.clone(), 3, Arc::new(Metrics::default()));
+            let metrics = Arc::new(Metrics::default());
+            let registry = Registry::new(sender.clone(), metrics.clone());
+            let rate_limited = Arc::new(InMemoryRateLimit::new(3, 10));
 
             Self {
                 received_messages: Arc::new(Mutex::new(HashMap::new())),
                 clients_failed_to_connect: Arc::new(Mutex::new(HashMap::new())),
                 current_client_id: 0,
                 cancel_token: CancellationToken::new(),
-                server: Server::new(addr.into(), registry),
+                server: Server::new(
+                    addr.into(),
+                    registry,
+                    metrics,
+                    rate_limited,
+                    "header".to_string(),
+                ),
                 server_addr: addr,
                 client_id_to_handle: HashMap::new(),
                 sender,

src/integration.rs

@@ -1,11 +1,14 @@
+mod client;
 #[cfg(all(feature = "integration", test))]
 mod integration;
 mod metrics;
+mod rate_limit;
 mod registry;
 mod server;
 mod subscriber;
 
 use crate::metrics::Metrics;
+use crate::rate_limit::InMemoryRateLimit;
 use crate::registry::Registry;
 use crate::server::Server;
 use crate::subscriber::WebsocketSubscriber;
@@ -49,7 +52,23 @@ struct Args {
         default_value = "100",
         help = "Maximum number of concurrently connected clients"
     )]
-    maximum_concurrent_connections: usize,
+    global_connections_limit: usize,
+
+    #[arg(
+        long,
+        env,
+        default_value = "10",
+        help = "Maximum number of concurrently connected clients"
+    )]
+    per_ip_connections_limit: usize,
+
+    #[arg(
+        long,
+        env,
+        default_value = "X-Forwarded-For",
+        help = "Header to use to determine the clients origin IP"
+    )]
+    ip_addr_http_header: String,
 
     #[arg(long, env, default_value = "info")]
     log_level: Level,
@@ -134,9 +153,20 @@ async fn main() {
     );
     let subscriber_task = subscriber.run(token.clone());
 
-    let registry = Registry::new(sender, args.maximum_concurrent_connections, metrics.clone());
+    let registry = Registry::new(sender, metrics.clone());
+
+    let rate_limiter = Arc::new(InMemoryRateLimit::new(
+        args.global_connections_limit,
+        args.per_ip_connections_limit,
+    ));
 
-    let server = Server::new(args.listen_addr, registry.clone());
+    let server = Server::new(
+        args.listen_addr,
+        registry.clone(),
+        metrics,
+        rate_limiter,
+        args.ip_addr_http_header,
+    );
     let server_task = server.listen(token.clone());
 
     let mut interrupt = signal(SignalKind::interrupt()).unwrap();