event_graph: rate limit message broadcast

src/event_graph/proto.rs

@@ -20,7 +20,7 @@ use std::{
     collections::{BTreeMap, HashSet, VecDeque},
     sync::{
         atomic::{AtomicUsize, Ordering::SeqCst},
-        Arc, Mutex as SyncMutex,
+        Arc,
     },
 };
 
@@ -29,7 +29,7 @@ use log::{debug, error, trace, warn};
 use smol::Executor;
 
 use super::{Event, EventGraphPtr, NULL_ID};
-use crate::{impl_p2p_message, net::*, util::time::NanoTimestamp, Error, Result};
+use crate::{impl_p2p_message, net::*, system::msleep, util::time::NanoTimestamp, Error, Result};
 
 /// Malicious behaviour threshold. If the threshold is reached, we will
 /// drop the peer from our P2P connection.
@@ -40,6 +40,47 @@ const WINDOW_MAXSIZE: usize = 200;
 /// Rolling length of the window
 const WINDOW_EXPIRY_TIME: NanoTimestamp = NanoTimestamp(60);
 
+/// Rolling length of the window
+const RATELIMIT_EXPIRY_TIME: NanoTimestamp = NanoTimestamp(10);
+/// Ratelimit kicks in above this count
+const RATELIMIT_MIN_COUNT: usize = 6;
+/// Sample point used to calculate sleep time when ratelimit is active
+const RATELIMIT_SAMPLE_IDX: usize = 10;
+/// Sleep for this amount of time when `count == RATE_LIMIT_SAMPLE_IDX`.
+const RATELIMIT_SAMPLE_SLEEP: usize = 1000;
+
+struct MovingWindow {
+    times: VecDeque<NanoTimestamp>,
+    expiry_time: NanoTimestamp,
+}
+
+impl MovingWindow {
+    fn new(expiry_time: NanoTimestamp) -> Self {
+        Self { times: VecDeque::new(), expiry_time }
+    }
+
+    /// Clean out expired timestamps from the window.
+    fn clean(&mut self) {
+        while let Some(ts) = self.times.front() {
+            if ts.elapsed().unwrap() < self.expiry_time {
+                break
+            }
+            let _ = self.times.pop_front();
+        }
+    }
+
+    /// Add new timestamp
+    fn ticktock(&mut self) {
+        self.clean();
+        self.times.push_back(NanoTimestamp::current_time());
+    }
+
+    #[inline]
+    fn count(&self) -> usize {
+        self.times.len()
+    }
+}
+
 /// P2P protocol implementation for the Event Graph.
 pub struct ProtocolEventGraph {
     /// Pointer to the connected peer
@@ -60,8 +101,11 @@ pub struct ProtocolEventGraph {
     malicious_count: AtomicUsize,
     /// P2P jobs manager pointer
     jobsman: ProtocolJobsManagerPtr,
-    /// Rolling window of event timestamps on this channel
-    bantimes: SyncMutex<VecDeque<NanoTimestamp>>,
+    /// To apply the rate-limit, we don't broadcast directly but instead send into the
+    /// sending queue.
+    broadcaster_push: smol::channel::Sender<EventPut>,
+    /// Receive send requests and rate-limit broadcasting them.
+    broadcaster_pull: smol::channel::Receiver<EventPut>,
 }
 
 /// A P2P message representing publishing an event on the network
@@ -96,6 +140,7 @@ impl ProtocolBase for ProtocolEventGraph {
         self.jobsman.clone().spawn(self.clone().handle_event_put(), ex.clone()).await;
         self.jobsman.clone().spawn(self.clone().handle_event_req(), ex.clone()).await;
         self.jobsman.clone().spawn(self.clone().handle_tip_req(), ex.clone()).await;
+        self.jobsman.clone().spawn(self.clone().broadcast_rate_limiter(), ex.clone()).await;
         Ok(())
     }
 
@@ -119,6 +164,8 @@ impl ProtocolEventGraph {
         let tip_req_sub = channel.subscribe_msg::<TipReq>().await?;
         let _tip_rep_sub = channel.subscribe_msg::<TipRep>().await?;
 
+        let (broadcaster_push, broadcaster_pull) = smol::channel::unbounded();
+
         Ok(Arc::new(Self {
             channel: channel.clone(),
             event_graph,
@@ -129,7 +176,8 @@ impl ProtocolEventGraph {
             _tip_rep_sub,
             malicious_count: AtomicUsize::new(0),
             jobsman: ProtocolJobsManager::new("ProtocolEventGraph", channel.clone()),
-            bantimes: SyncMutex::new(VecDeque::new()),
+            broadcaster_push,
+            broadcaster_pull,
         }))
     }
 
@@ -157,6 +205,9 @@ impl ProtocolEventGraph {
     /// This is triggered whenever someone broadcasts (or relays) a new
     /// event on the network.
     async fn handle_event_put(self: Arc<Self>) -> Result<()> {
+        // Rolling window of event timestamps on this channel
+        let mut bantimes = MovingWindow::new(WINDOW_EXPIRY_TIME);
+
         loop {
             let event = match self.ev_put_sub.receive().await {
                 Ok(v) => v.0.clone(),
@@ -188,23 +239,8 @@ impl ProtocolEventGraph {
 
             // There's a new unique event.
             // Apply ban logic to stop network floods.
-            let is_malicious = {
-                let mut bantimes = self.bantimes.lock().unwrap();
-
-                // Clean out expired timestamps from the window.
-                while let Some(ts) = bantimes.front() {
-                    if ts.elapsed().unwrap() < WINDOW_EXPIRY_TIME {
-                        break
-                    }
-                    let _ = bantimes.pop_front();
-                }
-
-                // Add new timestamp
-                bantimes.push_back(NanoTimestamp::current_time());
-
-                bantimes.len() > WINDOW_MAXSIZE
-            };
-            if is_malicious {
+            bantimes.ticktock();
+            if bantimes.count() > WINDOW_MAXSIZE {
                 self.channel.ban().await;
                 // This error is actually unused. We could return Ok here too.
                 return Err(Error::MaliciousFlood)
@@ -383,11 +419,7 @@ impl ProtocolEventGraph {
                 continue
             }
 
-            // Relay the event to other peers.
-            self.event_graph
-                .p2p
-                .broadcast_with_exclude(&EventPut(event), &[self.channel.address().clone()])
-                .await;
+            self.broadcaster_push.send(EventPut(event)).await.expect("push broadcaster closed");
         }
     }
 
@@ -525,4 +557,35 @@ impl ProtocolEventGraph {
             self.channel.send(&TipRep(layers)).await?;
         }
     }
+
+    /// We need to rate limit message propagation so malicious nodes don't get us banned
+    /// for flooding. We do that by aggregating messages here into a queue then apply
+    /// rate limit logic before broadcasting.
+    async fn broadcast_rate_limiter(self: Arc<Self>) -> Result<()> {
+        let mut ratelimit = MovingWindow::new(RATELIMIT_EXPIRY_TIME);
+
+        loop {
+            let event_put = self.broadcaster_pull.recv().await.expect("pull broadcaster closed");
+
+            ratelimit.ticktock();
+            if ratelimit.count() > RATELIMIT_MIN_COUNT {
+                let sleep_time =
+                    ((ratelimit.count() - RATELIMIT_MIN_COUNT) * RATELIMIT_SAMPLE_SLEEP /
+                        (RATELIMIT_SAMPLE_IDX - RATELIMIT_MIN_COUNT)) as u64;
+                debug!(
+                    target: "event_graph::protocol::broadcast_rate_limiter",
+                    "Activated rate limit: sleeping {sleep_time} ms [count={}]",
+                    ratelimit.count()
+                );
+                // Apply the ratelimit
+                msleep(sleep_time).await;
+            }
+
+            // Relay the event to other peers.
+            self.event_graph
+                .p2p
+                .broadcast_with_exclude(&event_put, &[self.channel.address().clone()])
+                .await;
+        }
+    }
 }