Просмотр исходного кода

event_graph: Cap peer request sizes

x 1 месяц назад
Родитель
Сommit
680364dfc5
3 измененных файлов с 105 добавлено и 11 удалено
  1. 31 3
      src/event_graph/mod.rs
  2. 51 4
      src/event_graph/proto.rs
  3. 23 4
      src/event_graph/tests.rs

+ 31 - 3
src/event_graph/mod.rs

@@ -50,7 +50,11 @@ pub mod event;
 pub use event::{display_order, Event, Header};
 
 pub mod proto;
-use proto::{EventRep, EventReq, HeaderRep, HeaderReq, StaticPut, SyncDirection, TipRep, TipReq};
+use proto::{
+    cap_layer_tips, count_layer_tips, EventRep, EventReq, HeaderRep, HeaderReq, StaticPut,
+    SyncDirection, TipRep, TipReq, MAX_EVENT_REP_EVENTS, MAX_EVENT_REQ_IDS, MAX_HEADER_REP_HEADERS,
+    MAX_HEADER_REQ_TIPS, MAX_RANGE_PAGE_SIZE, MAX_TIP_REP_TIPS,
+};
 
 pub mod rln;
 use rln::{IdentityState, RlnState, ZkKeys};
@@ -1199,9 +1203,11 @@ impl EventGraph {
                 return Err(Error::DagSyncFailed)
             }
 
-            let batch: Vec<blake3::Hash> = want.iter().copied().collect();
+            let batch: Vec<blake3::Hash> = want.iter().take(MAX_EVENT_REQ_IDS).copied().collect();
+            for id in &batch {
+                want.remove(id);
+            }
             let mut pending: HashSet<blake3::Hash> = batch.iter().copied().collect();
-            want.clear();
 
             // Ask every peer for the same batch. We keep consuming
             // responses until the batch is complete or every peer has
@@ -1373,6 +1379,7 @@ impl EventGraph {
         dir: SyncDirection,
         limit: usize,
     ) -> Result<Vec<Event>> {
+        let limit = limit.min(MAX_RANGE_PAGE_SIZE);
         let mut out = vec![];
         let store = self.dag_store.read().await;
         let slots: Vec<_> = match dir {
@@ -1789,6 +1796,10 @@ impl EventGraph {
         dag_name: &str,
         tips: &LayerUTips,
     ) -> Result<Vec<Header>> {
+        if count_layer_tips(tips) > MAX_HEADER_REQ_TIPS {
+            return Err(Error::DagSyncFailed)
+        }
+
         let dag_ts = u64::from_str(dag_name)?;
         let store = self.dag_store.read().await;
         let slot = store.get_slot(&dag_ts).ok_or(Error::DagSyncFailed)?;
@@ -1814,6 +1825,9 @@ impl EventGraph {
             let (id, v) = item?;
             let h = blake3::Hash::from_bytes((&id as &[u8]).try_into()?);
             if !ancestors.contains(&h) {
+                if out.len() >= MAX_HEADER_REP_HEADERS {
+                    break
+                }
                 out.push(deserialize_async(&v).await?);
             }
         }
@@ -2608,6 +2622,9 @@ async fn request_tips(
         .await
         .map_err(|_| Error::EventNotFound("tip timeout".into()))?;
     sub.unsubscribe().await;
+    if count_layer_tips(&r.0) > MAX_TIP_REP_TIPS {
+        return Err(Error::DagSyncFailed)
+    }
     Ok(r.0.clone())
 }
 
@@ -2618,12 +2635,16 @@ async fn request_header(
     timeout: u64,
 ) -> Result<Vec<Header>> {
     let sub = peer.subscribe_msg::<HeaderRep>().await?;
+    let tips = cap_layer_tips(&tips, MAX_HEADER_REQ_TIPS);
     peer.send(&HeaderReq(name, tips)).await?;
     let r = sub
         .receive_with_timeout(timeout)
         .await
         .map_err(|_| Error::EventNotFound("hdr timeout".into()))?;
     sub.unsubscribe().await;
+    if r.0.len() > MAX_HEADER_REP_HEADERS {
+        return Err(Error::DagSyncFailed)
+    }
     Ok(r.0.to_vec())
 }
 
@@ -2633,6 +2654,10 @@ async fn request_event(
     cid: usize,
     timeout: u64,
 ) -> (Result<(Vec<Event>, Vec<Vec<u8>>)>, usize, Arc<Channel>) {
+    if ids.len() > MAX_EVENT_REQ_IDS {
+        return (Err(Error::DagSyncFailed), cid, peer)
+    }
+
     let sub = match peer.subscribe_msg::<EventRep>().await {
         Ok(s) => s,
         Err(e) => return (Err(e), cid, peer),
@@ -2645,6 +2670,9 @@ async fn request_event(
     match sub.receive_with_timeout(timeout).await {
         Ok(r) => {
             sub.unsubscribe().await;
+            if r.0.len() > MAX_EVENT_REP_EVENTS || r.1.len() > MAX_EVENT_REP_EVENTS {
+                return (Err(Error::DagSyncFailed), cid, peer)
+            }
             (Ok((r.0.clone(), r.1.clone())), cid, peer)
         }
         Err(_) => (Err(Error::EventNotFound("ev timeout".into())), cid, peer),

+ 51 - 4
src/event_graph/proto.rs

@@ -111,6 +111,44 @@ const MAX_PARENT_FETCH_DEPTH: usize = 1000;
 /// unbounded memory growth under sustained load.
 const BROADCASTER_CAPACITY: usize = 256;
 
+/// Maximum event bodies a peer may request in one `EventReq`.
+pub const MAX_EVENT_REQ_IDS: usize = 128;
+/// Maximum event bodies accepted in one `EventRep`.
+pub const MAX_EVENT_REP_EVENTS: usize = MAX_EVENT_REQ_IDS;
+/// Maximum tips a peer may include in one `HeaderReq`.
+pub const MAX_HEADER_REQ_TIPS: usize = 1024;
+/// Maximum headers returned in one `HeaderRep`.
+pub const MAX_HEADER_REP_HEADERS: usize = 4096;
+/// Maximum tips returned in one `TipRep`.
+pub const MAX_TIP_REP_TIPS: usize = 1024;
+/// Maximum events served for one paginated range request.
+pub const MAX_RANGE_PAGE_SIZE: usize = 100;
+
+pub(crate) fn count_layer_tips(tips: &LayerUTips) -> usize {
+    tips.values().map(HashSet::len).sum()
+}
+
+pub(crate) fn cap_layer_tips(tips: &LayerUTips, limit: usize) -> LayerUTips {
+    let mut out = BTreeMap::new();
+    let mut remaining = limit;
+
+    for (layer, hashes) in tips {
+        if remaining == 0 {
+            break
+        }
+
+        let mut sorted: Vec<_> = hashes.iter().copied().collect();
+        sorted.sort_unstable_by(|a, b| a.as_bytes().cmp(b.as_bytes()));
+        let take = sorted.len().min(remaining);
+        if take > 0 {
+            out.insert(*layer, sorted.into_iter().take(take).collect());
+            remaining -= take;
+        }
+    }
+
+    out
+}
+
 struct MovingWindow {
     times: VecDeque<NanoTimestamp>,
     expiry_time: NanoTimestamp,
@@ -773,6 +811,10 @@ impl ProtocolEventGraph {
             if !self.event_graph.is_synced() {
                 continue
             }
+            if ids.len() > MAX_EVENT_REQ_IDS {
+                self.clone().strike().await?;
+                continue
+            }
 
             // Only serve IDs we've previously broadcast (prevents
             // arbitrary DAG enumeration by malicious peers). The
@@ -844,6 +886,10 @@ impl ProtocolEventGraph {
                 continue
             }
             let (dag_name, tips) = (&v.0, &v.1);
+            if count_layer_tips(tips) > MAX_HEADER_REQ_TIPS {
+                self.clone().strike().await?;
+                continue
+            }
             let dag_ts = match u64::from_str(dag_name) {
                 Ok(v) => v,
                 Err(_) => continue,
@@ -906,6 +952,8 @@ impl ProtocolEventGraph {
                 }
             };
 
+            let layers = cap_layer_tips(&layers, MAX_TIP_REP_TIPS);
+
             let mut b = self.event_graph.broadcasted_ids.write().await;
             for tips in layers.values() {
                 for t in tips {
@@ -929,10 +977,9 @@ impl ProtocolEventGraph {
             if !self.event_graph.is_synced() {
                 continue
             }
-            let events = self
-                .event_graph
-                .fetch_page(req.cursor_ts, req.direction.clone(), req.limit as usize)
-                .await?;
+            let limit = (req.limit as usize).min(MAX_RANGE_PAGE_SIZE);
+            let events =
+                self.event_graph.fetch_page(req.cursor_ts, req.direction.clone(), limit).await?;
             self.channel.send(&RangeRep(events)).await?;
         }
     }

+ 23 - 4
src/event_graph/tests.rs

@@ -33,13 +33,14 @@ use crate::{
         compute_unreferenced_tips,
         event::Header,
         filter_requested_event_rep, merge_static_sync_event_rep,
-        proto::{EventPut, SyncDirection},
+        proto::{cap_layer_tips, count_layer_tips, EventPut, SyncDirection, MAX_RANGE_PAGE_SIZE},
         test_helpers::{
             archive_config, bounded_dag_store_config, init_logger, make_eg, make_network,
             run_multi_node_test, shutdown_network, TestIdentity,
         },
         util::next_hour_timestamp,
-        DagStore, Event, EventGraphPtr, TimeIndex, NULL_ID, NULL_PARENTS, N_EVENT_PARENTS,
+        DagStore, Event, EventGraphPtr, LayerUTips, TimeIndex, NULL_ID, NULL_PARENTS,
+        N_EVENT_PARENTS,
     },
     system::{sleep, timeout::timeout},
 };
@@ -141,6 +142,18 @@ fn evgr_static_sync_merge_tracks_partial_requested_batches() {
     assert_eq!(fetched_ids, HashSet::from([event_a.id(), event_b.id()]));
 }
 
+#[test]
+fn evgr_layer_tip_cap_is_bounded() {
+    let mut tips = LayerUTips::new();
+    tips.entry(0).or_default().insert(blake3::hash(b"tip-0"));
+    tips.entry(1).or_default().insert(blake3::hash(b"tip-1"));
+    tips.entry(1).or_default().insert(blake3::hash(b"tip-2"));
+
+    let capped = cap_layer_tips(&tips, 2);
+    assert_eq!(count_layer_tips(&capped), 2);
+    assert!(capped.get(&0).is_some_and(|layer| layer.len() == 1));
+}
+
 #[test]
 fn evgr_time_index_queries_and_saturating_cursor() {
     // Forward, backward, newest, oldest queries plus the saturating
@@ -343,8 +356,8 @@ fn evgr_fetch_page_both_directions() {
         let eg = make_eg().await;
         let dag_name = eg.current_genesis.read().await.header.timestamp.to_string();
         let base = UNIX_EPOCH.elapsed().unwrap().as_millis() as u64;
-        for i in 0..10u64 {
-            let ev = Event::with_timestamp(base + i, vec![i as u8], &eg).await;
+        for i in 0..(MAX_RANGE_PAGE_SIZE as u64 + 10) {
+            let ev = Event::with_timestamp(base + i, vec![(i % 251) as u8], &eg).await;
             eg.header_dag_insert(vec![ev.header.clone()], &dag_name).await.unwrap();
             eg.dag_insert(slice::from_ref(&ev), &dag_name).await.unwrap();
         }
@@ -360,6 +373,12 @@ fn evgr_fetch_page_both_directions() {
         for w in page.windows(2) {
             assert!(w[0].header.timestamp <= w[1].header.timestamp);
         }
+
+        let capped = eg
+            .fetch_page(u64::MAX, SyncDirection::Backward, MAX_RANGE_PAGE_SIZE + 10)
+            .await
+            .unwrap();
+        assert_eq!(capped.len(), MAX_RANGE_PAGE_SIZE);
     })
 }