|
7 | 7 | use std::collections::{BTreeMap, HashMap}; |
8 | 8 |
|
9 | 9 | use tracing::warn; |
| 10 | +use willow_common::MAX_AUTHORS_PER_SYNC; |
10 | 11 | use willow_state::{ |
11 | 12 | apply_incremental, Event, EventDag, EventHash, EventKind, HeadsSummary, InsertError, |
12 | 13 | PendingBuffer, ServerState, Snapshot, DEFAULT_PENDING_MAX_AGE_MS, DEFAULT_PENDING_MAX_ENTRIES, |
@@ -273,6 +274,21 @@ impl WorkerRole for ReplayRole { |
273 | 274 | fn handle_request(&mut self, req: WorkerRequest) -> WorkerResponse { |
274 | 275 | match req { |
275 | 276 | WorkerRequest::Sync { server_id, heads } => { |
| 277 | + // Reject peer-supplied summaries that would force O(N) |
| 278 | + // BTreeMap construction and DAG walks (see |
| 279 | + // `MAX_AUTHORS_PER_SYNC`). Mirrors the storage cap added in |
| 280 | + // PR #507 / b075140; gated before any allocation so a hostile |
| 281 | + // request fails fast. |
| 282 | + if heads.heads.len() > MAX_AUTHORS_PER_SYNC { |
| 283 | + return WorkerResponse::Denied { |
| 284 | + reason: format!( |
| 285 | + "too many heads in sync request: {} > {}", |
| 286 | + heads.heads.len(), |
| 287 | + MAX_AUTHORS_PER_SYNC |
| 288 | + ), |
| 289 | + }; |
| 290 | + } |
| 291 | + |
276 | 292 | let data = match self.servers.get(&server_id) { |
277 | 293 | Some(d) => d, |
278 | 294 | None => { |
@@ -1458,6 +1474,80 @@ mod tests { |
1458 | 1474 | } |
1459 | 1475 | } |
1460 | 1476 |
|
| 1477 | + // ── Issue #514: oversize HeadsSummary rejection ────────────────────── |
| 1478 | + // |
| 1479 | + // Mirrors the storage cap added by PR #507 / b075140. Without a guard |
| 1480 | + // here, a malicious peer could send a multi-thousand-entry HeadsSummary |
| 1481 | + // and force replay to do per-author BTreeMap inserts and DAG walks for |
| 1482 | + // every entry — same DoS shape as the storage path the sibling PR fixed. |
| 1483 | + |
| 1484 | + /// Build a `HeadsSummary` with `n` distinct random authors. Mirrors the |
| 1485 | + /// helper in `crates/storage/src/store.rs` (sibling cap test infra). |
| 1486 | + fn heads_summary_with_authors(n: usize) -> HeadsSummary { |
| 1487 | + use willow_state::AuthorHead; |
| 1488 | + let mut heads = BTreeMap::new(); |
| 1489 | + for _ in 0..n { |
| 1490 | + let id = Identity::generate(); |
| 1491 | + heads.insert( |
| 1492 | + id.endpoint_id(), |
| 1493 | + AuthorHead { |
| 1494 | + seq: 1, |
| 1495 | + hash: EventHash::ZERO, |
| 1496 | + }, |
| 1497 | + ); |
| 1498 | + } |
| 1499 | + HeadsSummary { heads } |
| 1500 | + } |
| 1501 | + |
| 1502 | + /// A peer-supplied `HeadsSummary` with more than `MAX_AUTHORS_PER_SYNC` |
| 1503 | + /// entries must be rejected by `handle_request(Sync)` before any |
| 1504 | + /// per-author BTreeMap construction or DAG walk occurs. |
| 1505 | + #[test] |
| 1506 | + fn sync_request_rejects_oversize_heads() { |
| 1507 | + use willow_common::MAX_AUTHORS_PER_SYNC; |
| 1508 | + let mut role = ReplayRole::new(ReplayConfig::default()); |
| 1509 | + let (_, _) = setup_server(&mut role, "srv-1"); |
| 1510 | + |
| 1511 | + let oversize = heads_summary_with_authors(MAX_AUTHORS_PER_SYNC + 1); |
| 1512 | + let resp = role.handle_request(WorkerRequest::Sync { |
| 1513 | + server_id: "srv-1".to_string(), |
| 1514 | + heads: oversize, |
| 1515 | + }); |
| 1516 | + |
| 1517 | + match resp { |
| 1518 | + WorkerResponse::Denied { reason } => { |
| 1519 | + assert!( |
| 1520 | + reason.contains("too many heads"), |
| 1521 | + "denial reason should mention the cap; got: {reason}" |
| 1522 | + ); |
| 1523 | + } |
| 1524 | + other => panic!("expected Denied for oversize heads, got: {other:?}"), |
| 1525 | + } |
| 1526 | + } |
| 1527 | + |
| 1528 | + /// `handle_request(Sync)` must accept exactly `MAX_AUTHORS_PER_SYNC` |
| 1529 | + /// entries — the cap is inclusive on the legal side. |
| 1530 | + #[test] |
| 1531 | + fn sync_request_accepts_exact_cap_heads() { |
| 1532 | + use willow_common::MAX_AUTHORS_PER_SYNC; |
| 1533 | + let mut role = ReplayRole::new(ReplayConfig::default()); |
| 1534 | + let (_, _) = setup_server(&mut role, "srv-1"); |
| 1535 | + |
| 1536 | + let at_cap = heads_summary_with_authors(MAX_AUTHORS_PER_SYNC); |
| 1537 | + let resp = role.handle_request(WorkerRequest::Sync { |
| 1538 | + server_id: "srv-1".to_string(), |
| 1539 | + heads: at_cap, |
| 1540 | + }); |
| 1541 | + |
| 1542 | + // The peer's heads mention authors we don't know, so events_since |
| 1543 | + // returns the genesis event; our store has 1 author the peer doesn't, |
| 1544 | + // so they_are_behind is true. Either branch is acceptable — the only |
| 1545 | + // forbidden outcome is Denied for at-cap input. |
| 1546 | + if let WorkerResponse::Denied { reason } = resp { |
| 1547 | + panic!("at-cap heads must not be denied; got reason: {reason}"); |
| 1548 | + } |
| 1549 | + } |
| 1550 | + |
1461 | 1551 | /// When the configured `pending_max_entries` is exceeded, the oldest |
1462 | 1552 | /// entries are evicted and `pending_count()` reflects the cap. |
1463 | 1553 | #[test] |
|
0 commit comments