Add rate limiter for total number of subscriber per ip (#435)

danimhr · web-flow · commit 0be61f39d199 · 2025-03-18T15:08:14.000+01:00
diff --git a/auction-server/src/api.rs b/auction-server/src/api.rs
@@ -369,6 +369,8 @@ pub enum RestError {
     InvalidFirstSigner(String),
     /// Duplicate bid
     DuplicateBid,
+    /// Too many open websocket connections
+    TooManyOpenWebsocketConnections,
 }
 
 
@@ -482,6 +484,10 @@ impl RestError {
                 StatusCode::BAD_REQUEST,
                 "Duplicate bid".to_string(),
             ),
+            RestError::TooManyOpenWebsocketConnections => (
+                StatusCode::TOO_MANY_REQUESTS,
+                "Too many open websocket connections".to_string(),
+            ),
         }
     }
 }
diff --git a/auction-server/src/api/ws.rs b/auction-server/src/api/ws.rs
@@ -1,6 +1,7 @@
 use {
     super::{
         Auth,
+        RestError,
         WrappedRouter,
     },
     crate::{
@@ -32,6 +33,7 @@ use {
             State,
             WebSocketUpgrade,
         },
+        http::HeaderMap,
         response::IntoResponse,
         Router,
     },
@@ -68,8 +70,12 @@ use {
         StreamExt,
     },
     std::{
-        collections::HashSet,
+        collections::{
+            HashMap,
+            HashSet,
+        },
         future::Future,
+        net::IpAddr,
         sync::{
             atomic::{
                 AtomicUsize,
@@ -82,6 +88,7 @@ use {
     time::OffsetDateTime,
     tokio::sync::{
         broadcast,
+        RwLock,
         Semaphore,
     },
     tracing::{
@@ -91,26 +98,102 @@ use {
 };
 
 pub struct WsState {
-    pub subscriber_counter: AtomicUsize,
-    pub broadcast_sender:   broadcast::Sender<UpdateEvent>,
-    pub broadcast_receiver: broadcast::Receiver<UpdateEvent>,
+    pub requester_ip_header_name: String,
+    subscriber_counter:           AtomicUsize,
+    subscriber_per_ip:            RwLock<HashMap<IpAddr, HashSet<SubscriberId>>>,
+    pub broadcast_sender:         broadcast::Sender<UpdateEvent>,
+    pub broadcast_receiver:       broadcast::Receiver<UpdateEvent>,
+}
+
+const MAXIMUM_SUBSCRIBERS_PER_IP: usize = 10;
+
+impl WsState {
+    pub fn new(requester_ip_header_name: String, broadcast_channel_size: usize) -> Self {
+        let (broadcast_sender, broadcast_receiver) = broadcast::channel(broadcast_channel_size);
+        Self {
+            requester_ip_header_name,
+            subscriber_counter: AtomicUsize::new(0),
+            subscriber_per_ip: RwLock::new(HashMap::new()),
+            broadcast_sender,
+            broadcast_receiver,
+        }
+    }
+
+    /// If the specified IP address has too many open websocket connections, this function will
+    /// return none. Otherwise, it will return the new subscriber id.
+    pub async fn get_new_subscriber_id(&self, ip: Option<IpAddr>) -> Option<SubscriberId> {
+        let id = self.subscriber_counter.fetch_add(1, Ordering::SeqCst);
+        if let Some(ip) = ip {
+            let mut write_gaurd = self.subscriber_per_ip.write().await;
+            let ids = write_gaurd.entry(ip).or_insert_with(HashSet::new);
+            if ids.len() >= MAXIMUM_SUBSCRIBERS_PER_IP {
+                return None;
+            }
+            ids.insert(id);
+        }
+        Some(id)
+    }
+
+    pub async fn remove_subscriber(&self, id: SubscriberId, ip: Option<IpAddr>) {
+        if let Some(ip) = ip {
+            let mut write_gaurd = self.subscriber_per_ip.write().await;
+            if let Some(ids) = write_gaurd.get_mut(&ip) {
+                ids.remove(&id);
+                if ids.is_empty() {
+                    write_gaurd.remove(&ip);
+                }
+            }
+        }
+    }
 }
 
 pub async fn ws_route_handler(
     auth: Auth,
     ws: WebSocketUpgrade,
     State(store): State<Arc<StoreNew>>,
+    headers: HeaderMap,
 ) -> impl IntoResponse {
-    ws.on_upgrade(move |socket| websocket_handler(socket, store, auth))
+    let ws_state = &store.store.ws;
+    let requester_ip = headers
+        .get(ws_state.requester_ip_header_name.as_str())
+        .and_then(|value| value.to_str().ok())
+        .and_then(|value| value.split(',').next()) // Only take the first ip if there are multiple
+        .and_then(|value| value.parse().ok());
+
+    if requester_ip.is_none() {
+        tracing::warn!("Failed to get requester IP address");
+    }
+
+    match ws_state.get_new_subscriber_id(requester_ip).await {
+        Some(subscriber_id) => ws.on_upgrade(move |socket| {
+            websocket_handler(socket, store, subscriber_id, auth, requester_ip)
+        }),
+        None => RestError::TooManyOpenWebsocketConnections.into_response(),
+    }
 }
 
-async fn websocket_handler(stream: WebSocket, state: Arc<StoreNew>, auth: Auth) {
+async fn websocket_handler(
+    stream: WebSocket,
+    state: Arc<StoreNew>,
+    subscriber_id: SubscriberId,
+    auth: Auth,
+    requester_ip: Option<IpAddr>,
+) {
     let ws_state = &state.store.ws;
-    let id = ws_state.subscriber_counter.fetch_add(1, Ordering::SeqCst);
     let (sender, receiver) = stream.split();
     let new_receiver = ws_state.broadcast_receiver.resubscribe();
-    let mut subscriber = Subscriber::new(id, state, new_receiver, receiver, sender, auth);
+    let mut subscriber = Subscriber::new(
+        subscriber_id,
+        state.clone(),
+        new_receiver,
+        receiver,
+        sender,
+        auth,
+    );
     subscriber.run().await;
+    ws_state
+        .remove_subscriber(subscriber_id, requester_ip)
+        .await;
 }
 
 #[derive(Clone, PartialEq, Debug)]
diff --git a/auction-server/src/config/server.rs b/auction-server/src/config/server.rs
@@ -6,6 +6,7 @@ use {
 const DEFAULT_LISTEN_ADDR: &str = "127.0.0.1:9000";
 const DEFAULT_METRICS_ADDR: &str = "127.0.0.1:9001";
 const DEFAULT_DATABASE_MAX_CONNECTIONS: &str = "10";
+const DEAFULT_REQUESTER_IP_HEADER_NAME: &str = "X-Forwarded-For";
 
 #[derive(Args, Clone, Debug)]
 #[command(next_help_heading = "Server Options")]
@@ -30,4 +31,9 @@ pub struct Options {
     #[arg(default_value = DEFAULT_METRICS_ADDR)]
     #[arg(env = "METRICS_ADDR")]
     pub metrics_addr:             SocketAddr,
+    /// The header name to use for the requester IP address.
+    #[arg(long = "requester-ip-header-name")]
+    #[arg(default_value = DEAFULT_REQUESTER_IP_HEADER_NAME)]
+    #[arg(env = "REQUESTER_IP_HEADER_NAME")]
+    pub requester_ip_header_name: String,
 }
diff --git a/auction-server/src/opportunity/service/mod.rs b/auction-server/src/opportunity/service/mod.rs
@@ -305,7 +305,6 @@ pub mod tests {
             opportunity::repository::MockDatabase,
             server::setup_metrics_recorder,
         },
-        std::sync::atomic::AtomicUsize,
         tokio::sync::broadcast::Receiver,
     };
 
@@ -325,20 +324,14 @@ pub mod tests {
                 ordered_fee_tokens:      vec![],
             };
 
-            let (broadcast_sender, broadcast_receiver) = tokio::sync::broadcast::channel(100);
-
             let mut chains_svm = HashMap::new();
             chains_svm.insert(chain_id.clone(), config_svm);
 
             let store = Arc::new(Store {
                 db:               DB::connect_lazy("https://test").unwrap(),
                 chains_evm:       HashMap::new(),
                 chains_svm:       HashMap::new(),
-                ws:               ws::WsState {
-                    subscriber_counter: AtomicUsize::new(0),
-                    broadcast_sender,
-                    broadcast_receiver,
-                },
+                ws:               ws::WsState::new("X-Forwarded-For".to_string(), 100),
                 secret_key:       "test".to_string(),
                 access_tokens:    RwLock::new(HashMap::new()),
                 metrics_recorder: setup_metrics_recorder().unwrap(),
diff --git a/auction-server/src/server.rs b/auction-server/src/server.rs
@@ -71,7 +71,6 @@ use {
         sync::{
             atomic::{
                 AtomicBool,
-                AtomicUsize,
                 Ordering,
             },
             Arc,
@@ -292,9 +291,6 @@ pub async fn start_server(run_options: RunOptions) -> Result<()> {
 
     let chains_svm = setup_chain_store_svm(config_map)?;
 
-    let (broadcast_sender, broadcast_receiver) =
-        tokio::sync::broadcast::channel(NOTIFICATIONS_CHAN_LEN);
-
     let pool = create_pg_pool(
         &run_options.server.database_url,
         run_options.server.database_max_connections,
@@ -321,11 +317,10 @@ pub async fn start_server(run_options: RunOptions) -> Result<()> {
         db:               pool.clone(),
         chains_evm:       chains_evm.clone(),
         chains_svm:       chains_svm.clone(),
-        ws:               ws::WsState {
-            subscriber_counter: AtomicUsize::new(0),
-            broadcast_sender,
-            broadcast_receiver,
-        },
+        ws:               ws::WsState::new(
+            run_options.server.requester_ip_header_name.clone(),
+            NOTIFICATIONS_CHAN_LEN,
+        ),
         secret_key:       run_options.secret_key.clone(),
         access_tokens:    RwLock::new(access_tokens),
         metrics_recorder: setup_metrics_recorder()?,

Original file line number	Diff line number	Diff line change
`@@ -369,6 +369,8 @@ pub enum RestError {`
`369`	`369`	`InvalidFirstSigner(String),`
`370`	`370`	`/// Duplicate bid`
`371`	`371`	`DuplicateBid,`
	`372`	`+ /// Too many open websocket connections`
	`373`	`+ TooManyOpenWebsocketConnections,`
`372`	`374`	`}`
`373`	`375`
`374`	`376`
`@@ -482,6 +484,10 @@ impl RestError {`
`482`	`484`	`StatusCode::BAD_REQUEST,`
`483`	`485`	`"Duplicate bid".to_string(),`
`484`	`486`	`),`
	`487`	`+ RestError::TooManyOpenWebsocketConnections => (`
	`488`	`+ StatusCode::TOO_MANY_REQUESTS,`
	`489`	`+ "Too many open websocket connections".to_string(),`
	`490`	`+ ),`
`485`	`491`	`}`
`486`	`492`	`}`
`487`	`493`	`}`