Streamable size limits

crates/chia-protocol/src/full_node_protocol.rs

@@ -70,6 +70,7 @@ pub struct RequestBlocks {
 pub struct RespondBlocks {
     start_height: u32,
     end_height: u32,
+    #[chia(max_length = 64)]
     blocks: Vec<FullBlock>,
 }
 

crates/chia-protocol/src/wallet_protocol.rs

@@ -72,14 +72,17 @@ pub struct RejectHeaderRequest {
 pub struct RequestRemovals {
     height: u32,
     header_hash: Bytes32,
+    #[chia(max_length = 30000)]
     coin_names: Option<Vec<Bytes32>>,
 }
 
 #[streamable(message)]
 pub struct RespondRemovals {
     height: u32,
     header_hash: Bytes32,
+    #[chia(max_length = 30000)]
     coins: Vec<(Bytes32, Option<Coin>)>,
+    #[chia(max_length = 30000)]
     proofs: Option<Vec<(Bytes32, Bytes)>>,
 }
 
@@ -93,14 +96,17 @@ pub struct RejectRemovalsRequest {
 pub struct RequestAdditions {
     height: u32,
     header_hash: Option<Bytes32>,
+    #[chia(max_length = 30000)]
     puzzle_hashes: Option<Vec<Bytes32>>,
 }
 
 #[streamable(message)]
 pub struct RespondAdditions {
     height: u32,
     header_hash: Bytes32,
+    #[chia(max_length = 30000)]
     coins: Vec<(Bytes32, Vec<Coin>)>,
+    #[chia(max_length = 30000)]
     proofs: Option<Vec<(Bytes32, Bytes, Option<Bytes>)>>,
 }
 
@@ -114,6 +120,7 @@ pub struct RejectAdditionsRequest {
 pub struct RespondBlockHeaders {
     start_height: u32,
     end_height: u32,
+    #[chia(max_length = 64)]
     header_blocks: Vec<HeaderBlock>,
 }
 
@@ -146,32 +153,39 @@ pub struct RejectHeaderBlocks {
 pub struct RespondHeaderBlocks {
     start_height: u32,
     end_height: u32,
+    #[chia(max_length = 64)]
     header_blocks: Vec<HeaderBlock>,
 }
 
 #[streamable(message)]
 pub struct RegisterForPhUpdates {
+    #[chia(max_length = 1600000)]
     puzzle_hashes: Vec<Bytes32>,
     min_height: u32,
 }
 
 #[streamable(message)]
 pub struct RespondToPhUpdates {
+    #[chia(max_length = 1600000)]
     puzzle_hashes: Vec<Bytes32>,
     min_height: u32,
+    #[chia(max_length = 1600000)]
     coin_states: Vec<CoinState>,
 }
 
 #[streamable(message)]
 pub struct RegisterForCoinUpdates {
+    #[chia(max_length = 1600000)]
     coin_ids: Vec<Bytes32>,
     min_height: u32,
 }
 
 #[streamable(message)]
 pub struct RespondToCoinUpdates {
+    #[chia(max_length = 1600000)]
     coin_ids: Vec<Bytes32>,
     min_height: u32,
+    #[chia(max_length = 1600000)]
     coin_states: Vec<CoinState>,
 }
 
@@ -180,6 +194,7 @@ pub struct CoinStateUpdate {
     height: u32,
     fork_height: u32,
     peak_hash: Bytes32,
+    #[chia(max_length = 30000)]
     items: Vec<CoinState>,
 }
 
@@ -190,6 +205,7 @@ pub struct RequestChildren {
 
 #[streamable(message)]
 pub struct RespondChildren {
+    #[chia(max_length = 30000)]
     coin_states: Vec<CoinState>,
 }
 
@@ -217,21 +233,25 @@ pub struct RespondFeeEstimates {
 
 #[streamable(message)]
 pub struct RequestRemovePuzzleSubscriptions {
+    #[chia(max_length = 1600000)]
     puzzle_hashes: Option<Vec<Bytes32>>,
 }
 
 #[streamable(message)]
 pub struct RespondRemovePuzzleSubscriptions {
+    #[chia(max_length = 1600000)]
     puzzle_hashes: Vec<Bytes32>,
 }
 
 #[streamable(message)]
 pub struct RequestRemoveCoinSubscriptions {
+    #[chia(max_length = 1600000)]
     coin_ids: Option<Vec<Bytes32>>,
 }
 
 #[streamable(message)]
 pub struct RespondRemoveCoinSubscriptions {
+    #[chia(max_length = 1600000)]
     coin_ids: Vec<Bytes32>,
 }
 
@@ -245,6 +265,7 @@ pub struct CoinStateFilters {
 
 #[streamable(message)]
 pub struct RequestPuzzleState {
+    #[chia(max_length = 32700)]
     puzzle_hashes: Vec<Bytes32>,
     previous_height: Option<u32>,
     header_hash: Bytes32,
@@ -254,10 +275,12 @@ pub struct RequestPuzzleState {
 
 #[streamable(message)]
 pub struct RespondPuzzleState {
+    #[chia(max_length = 32700)]
     puzzle_hashes: Vec<Bytes32>,
     height: u32,
     header_hash: Bytes32,
     is_finished: bool,
+    #[chia(max_length = 32700)]
     coin_states: Vec<CoinState>,
 }
 
@@ -268,6 +291,7 @@ pub struct RejectPuzzleState {
 
 #[streamable(message)]
 pub struct RequestCoinState {
+    #[chia(max_length = 32700)]
     coin_ids: Vec<Bytes32>,
     previous_height: Option<u32>,
     header_hash: Bytes32,
@@ -276,7 +300,9 @@ pub struct RequestCoinState {
 
 #[streamable(message)]
 pub struct RespondCoinState {
+    #[chia(max_length = 32700)]
     coin_ids: Vec<Bytes32>,
+    #[chia(max_length = 32700)]
     coin_states: Vec<CoinState>,
 }
 
@@ -334,11 +360,13 @@ pub struct RemovedMempoolItem {
 
 #[streamable(message)]
 pub struct MempoolItemsAdded {
+    #[chia(max_length = 30000)]
     transaction_ids: Vec<Bytes32>,
 }
 
 #[streamable(message)]
 pub struct MempoolItemsRemoved {
+    #[chia(max_length = 30000)]
     removed_items: Vec<RemovedMempoolItem>,
 }
 

crates/chia-traits/src/streamable.rs

@@ -124,19 +124,28 @@ impl<T: Streamable> Streamable for Vec<T> {
     }
 
     fn parse<const TRUSTED: bool>(input: &mut Cursor<&[u8]>) -> Result<Self> {
-        let len = u32::parse::<TRUSTED>(input)?;
+        parse_vec_with_max_length::<T, TRUSTED>(input, u32::MAX as usize)
+    }
+}
 
-        let mut ret = if mem::size_of::<T>() == 0 {
-            Vec::<T>::new()
-        } else {
-            let limit = 2 * 1024 * 1024 / mem::size_of::<T>();
-            Vec::<T>::with_capacity(std::cmp::min(limit, len as usize))
-        };
-        for _ in 0..len {
-            ret.push(T::parse::<TRUSTED>(input)?);
-        }
-        Ok(ret)
+pub fn parse_vec_with_max_length<T: Streamable, const TRUSTED: bool>(
+    input: &mut Cursor<&[u8]>,
+    max_length: usize,
+) -> Result<Vec<T>> {
+    let len = u32::parse::<TRUSTED>(input)?;
+    if (len as usize) > max_length {
+        return Err(Error::SequenceTooLarge);
     }
+    let mut ret = if mem::size_of::<T>() == 0 {
+        Vec::<T>::new()
+    } else {
+        let limit = 2 * 1024 * 1024 / mem::size_of::<T>();
+        Vec::<T>::with_capacity(std::cmp::min(limit, len as usize))
+    };
+    for _ in 0..len {
+        ret.push(T::parse::<TRUSTED>(input)?);
+    }
+    Ok(ret)
 }
 
 impl Streamable for String {
@@ -784,3 +793,57 @@ fn test_stream_enum() {
     assert_eq!(stream::<TestEnum>(&TestEnum::B), &[1_u8]);
     assert_eq!(stream::<TestEnum>(&TestEnum::C), &[255_u8]);
 }
+
+#[cfg(test)]
+#[derive(Streamable, PartialEq, Debug)]
+struct TestBoundedVec {
+    #[chia(max_length = 3)]
+    items: Vec<u32>,
+}
+
+#[test]
+fn test_parse_bounded_vec() {
+    let buf: &[u8] = &[0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 3];
+    from_bytes::<TestBoundedVec>(
+        buf,
+        TestBoundedVec {
+            items: vec![1, 2, 3],
+        },
+    );
+}
+
+#[test]
+fn test_parse_bounded_vec_too_large() {
+    let buf: &[u8] = &[0, 0, 0, 4, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0, 4];
+    from_bytes_fail::<TestBoundedVec>(buf, Error::SequenceTooLarge);
+}
+
+#[cfg(test)]
+#[derive(Streamable, PartialEq, Debug)]
+struct TestBoundedOptionVec {
+    #[chia(max_length = 2)]
+    items: Option<Vec<u32>>,
+}
+
+#[test]
+fn test_parse_bounded_option_vec_none() {
+    let buf: &[u8] = &[0];
+    from_bytes::<TestBoundedOptionVec>(buf, TestBoundedOptionVec { items: None });
+}
+
+#[test]
+fn test_parse_bounded_option_vec_some() {
+    let buf: &[u8] = &[1, 0, 0, 0, 2, 0, 0, 0, 1, 0, 0, 0, 2];
+    from_bytes::<TestBoundedOptionVec>(
+        buf,
+        TestBoundedOptionVec {
+            items: Some(vec![1, 2]),
+        },
+    );
+}
+
+#[test]
+fn test_parse_bounded_option_vec_too_large() {
+    let buf: &[u8] = &[1, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 2, 0, 0, 0, 3];
+    from_bytes_fail::<TestBoundedOptionVec>(buf, Error::SequenceTooLarge);
+}