p2p_core%fix: accept up to 8000 headers in headers2, bound header counts by the input, export the header count limits

pkgs/p2p_core/src/msg/headers.rs

@@ -11,12 +11,21 @@ use crate::prelude::*;
 use crate::version::ProtocolVersion;
 
 use dash_primitives::{hash_impl, BlockHash, BlockHeader, MerkleRoot};
-use dash_types::codec::{BaseCodec, DecodeError, EncodeBuf};
+use dash_types::codec::{self, BaseCodec, DecodeError, EncodeBuf};
 use dash_types::type_id::TypeId;
 use dash_types::CompactSize;
 
-/// Maximum headers per message.
-const MAX_HEADERS: usize = 2_000;
+/// The fewest bytes a `headers` entry takes: a header and its tx count.
+const MIN_HEADERS_ENTRY: usize = 81;
+
+/// Reads an entry count of at most `limit`, reporting one the remaining
+/// input cannot hold as [`DecodeError::Eof`] before anything is allocated
+/// for it.
+pub(super) fn decode_count(data: &mut &[u8], limit: usize, min_entry: usize) -> Result<usize, DecodeError> {
+  let count = CompactSize::decode(data)?.into_len(limit)?;
+  codec::ensure(data, count * min_entry)?;
+  Ok(count)
+}
 
 /// Requests block headers starting from a locator.
 #[derive(Clone, Debug, Eq, Hash, PartialEq, TypeId)]
@@ -51,7 +60,7 @@ impl_p2p!(Headers);
 
 impl BaseCodec for Headers {
   fn decode(data: &mut &[u8]) -> Result<Self, DecodeError> {
-    let count = CompactSize::decode(data)?.into_len(MAX_HEADERS)?;
+    let count = decode_count(data, Self::MAX_COUNT, MIN_HEADERS_ENTRY)?;
     let mut headers = Vec::with_capacity(count);
     for _ in 0..count {
       headers.push(BlockHeader {
@@ -83,3 +92,27 @@ impl BaseCodec for Headers {
 }
 
 hash_impl!(Headers);
+
+impl Headers {
+  /// The most headers a `headers` message carries.
+  pub const MAX_COUNT: usize = 2_000;
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  use rstest::rstest;
+
+  #[rstest]
+  fn counts_past_the_input_fail_before_any_header() {
+    let payload = [&[2][..], &[0; MIN_HEADERS_ENTRY]].concat();
+    assert_eq!(
+      Headers::decode(&mut &payload[..]),
+      Err(DecodeError::Eof {
+        needed: 2 * MIN_HEADERS_ENTRY,
+        remaining: MIN_HEADERS_ENTRY,
+      })
+    );
+  }
+}

pkgs/p2p_core/src/msg/headers2.rs

@@ -6,6 +6,7 @@
 
 //! Compressed header messages.
 
+use super::headers::{decode_count, Headers};
 use crate::codec::{codec_p2p, impl_p2p};
 use crate::prelude::*;
 use crate::version::ProtocolVersion;
@@ -15,8 +16,9 @@ use dash_types::codec::{BaseCodec, DecodeError, EncodeBuf};
 use dash_types::type_id::{TypeId, Unencodable};
 use dash_types::{CompactSize, Hashable};
 
-/// Maximum headers per message.
-const MAX_HEADERS: usize = 2_000;
+/// The fewest bytes a compressed header takes: flags, merkle root, a time
+/// delta and the nonce.
+const MIN_HEADERS2_ENTRY: usize = 39;
 
 /// Maximum entries in the MRU version cache.
 const MAX_VERSION_CACHE: usize = 7;
@@ -255,7 +257,7 @@ impl_p2p!(Headers2);
 
 impl BaseCodec for Headers2 {
   fn decode(data: &mut &[u8]) -> Result<Self, DecodeError> {
-    let count = CompactSize::decode(data)?.into_len(MAX_HEADERS)?;
+    let count = decode_count(data, Self::MAX_COUNT, MIN_HEADERS2_ENTRY)?;
     let mut state = CompressionState::new();
     let mut headers = Vec::with_capacity(count);
     for _ in 0..count {
@@ -275,11 +277,27 @@ impl BaseCodec for Headers2 {
 
 hash_impl!(Headers2);
 
+impl Headers2 {
+  /// The most headers a `headers2` message carries.
+  pub const MAX_COUNT: usize = 8_000;
+
+  /// The most headers a peer may send in one `headers2` message, given the
+  /// protocol version both peers share.
+  pub fn max_count(common: ProtocolVersion) -> usize {
+    if common >= ProtocolVersion::INCREASE_MAX_HEADERS2 {
+      Self::MAX_COUNT
+    } else {
+      Headers::MAX_COUNT
+    }
+  }
+}
+
 #[cfg(all(test, feature = "serde"))]
 mod tests {
   use super::*;
 
   use dash_dev::{assert_serde_rt, check_wire, Corpus};
+  use dash_types::Numeric;
   use rstest::rstest;
 
   #[rstest]
@@ -295,4 +313,46 @@ mod tests {
     let items = corpus.entries::<Headers2>("headers2", check_wire);
     assert_serde_rt("headers2", &items);
   }
+
+  #[rstest]
+  fn counts_past_2000_decode() {
+    let headers = (0..2_001)
+      .map(|i| BlockHeader {
+        version: 0x2000_0000,
+        prev_hash: BlockHash::from_lendian([0; 32]),
+        merkle_root: MerkleRoot::from_lendian([1; 32]),
+        time: 1_700_000_000 + i,
+        bits: 0x1e0f_ffff,
+        nonce: i,
+      })
+      .collect();
+    let msg = Headers2 { headers };
+    let mut buf = Vec::new();
+    msg.encode(&mut buf);
+    assert_eq!(Headers2::decode(&mut &buf[..]), Ok(msg));
+  }
+
+  #[rstest]
+  fn counts_past_the_limit_fail_before_any_header() {
+    let payload = [0xfd, 0x41, 0x1f];
+    assert_eq!(
+      Headers2::decode(&mut &payload[..]),
+      Err(DecodeError::CompactSizeExceedsLimit {
+        limit: Headers2::MAX_COUNT,
+        value: 8_001,
+      })
+    );
+  }
+
+  #[rstest]
+  fn counts_past_the_input_fail_before_any_header() {
+    let payload = [&[2][..], &[0; MIN_HEADERS2_ENTRY]].concat();
+    assert_eq!(
+      Headers2::decode(&mut &payload[..]),
+      Err(DecodeError::Eof {
+        needed: 2 * MIN_HEADERS2_ENTRY,
+        remaining: MIN_HEADERS2_ENTRY,
+      })
+    );
+  }
 }

pkgs/p2p_core/src/version.rs

@@ -33,4 +33,6 @@ impl ProtocolVersion {
   pub const MNLISTDIFF_VERSION_ORDER: Self = Self(70229);
   /// Chainlock signatures in MN list diff boundary.
   pub const MNLISTDIFF_CHAINLOCKS: Self = Self(70230);
+  /// `headers2` limit raised to 8000 boundary.
+  pub const INCREASE_MAX_HEADERS2: Self = Self(70235);
 }