diff --git a/crates/chia-protocol/src/full_node_protocol.rs b/crates/chia-protocol/src/full_node_protocol.rs index 246c9b95b..f4b74e1ca 100644 --- a/crates/chia-protocol/src/full_node_protocol.rs +++ b/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, } diff --git a/crates/chia-protocol/src/wallet_protocol.rs b/crates/chia-protocol/src/wallet_protocol.rs index 2919544a6..625fb85d1 100644 --- a/crates/chia-protocol/src/wallet_protocol.rs +++ b/crates/chia-protocol/src/wallet_protocol.rs @@ -72,6 +72,7 @@ pub struct RejectHeaderRequest { pub struct RequestRemovals { height: u32, header_hash: Bytes32, + #[chia(max_length = 30000)] coin_names: Option>, } @@ -79,7 +80,9 @@ pub struct RequestRemovals { pub struct RespondRemovals { height: u32, header_hash: Bytes32, + #[chia(max_length = 30000)] coins: Vec<(Bytes32, Option)>, + #[chia(max_length = 30000)] proofs: Option>, } @@ -93,6 +96,7 @@ pub struct RejectRemovalsRequest { pub struct RequestAdditions { height: u32, header_hash: Option, + #[chia(max_length = 30000)] puzzle_hashes: Option>, } @@ -100,7 +104,9 @@ pub struct RequestAdditions { pub struct RespondAdditions { height: u32, header_hash: Bytes32, + #[chia(max_length = 30000)] coins: Vec<(Bytes32, Vec)>, + #[chia(max_length = 30000)] proofs: Option)>>, } @@ -114,6 +120,7 @@ pub struct RejectAdditionsRequest { pub struct RespondBlockHeaders { start_height: u32, end_height: u32, + #[chia(max_length = 128)] header_blocks: Vec, } @@ -146,32 +153,39 @@ pub struct RejectHeaderBlocks { pub struct RespondHeaderBlocks { start_height: u32, end_height: u32, + #[chia(max_length = 64)] header_blocks: Vec, } #[streamable(message)] pub struct RegisterForPhUpdates { + #[chia(max_length = 1600000)] puzzle_hashes: Vec, min_height: u32, } #[streamable(message)] pub struct RespondToPhUpdates { + #[chia(max_length = 1600000)] puzzle_hashes: Vec, min_height: u32, + #[chia(max_length = 500000)] coin_states: Vec, } #[streamable(message)] pub struct RegisterForCoinUpdates { + #[chia(max_length = 1600000)] coin_ids: Vec, min_height: u32, } #[streamable(message)] pub struct RespondToCoinUpdates { + #[chia(max_length = 1600000)] coin_ids: Vec, min_height: u32, + #[chia(max_length = 500000)] coin_states: Vec, } @@ -180,6 +194,7 @@ pub struct CoinStateUpdate { height: u32, fork_height: u32, peak_hash: Bytes32, + #[chia(max_length = 30000)] items: Vec, } @@ -190,6 +205,7 @@ pub struct RequestChildren { #[streamable(message)] pub struct RespondChildren { + #[chia(max_length = 30000)] coin_states: Vec, } @@ -217,21 +233,25 @@ pub struct RespondFeeEstimates { #[streamable(message)] pub struct RequestRemovePuzzleSubscriptions { + #[chia(max_length = 1600000)] puzzle_hashes: Option>, } #[streamable(message)] pub struct RespondRemovePuzzleSubscriptions { + #[chia(max_length = 1600000)] puzzle_hashes: Vec, } #[streamable(message)] pub struct RequestRemoveCoinSubscriptions { + #[chia(max_length = 1600000)] coin_ids: Option>, } #[streamable(message)] pub struct RespondRemoveCoinSubscriptions { + #[chia(max_length = 1600000)] coin_ids: Vec, } @@ -245,6 +265,7 @@ pub struct CoinStateFilters { #[streamable(message)] pub struct RequestPuzzleState { + #[chia(max_length = 35000)] puzzle_hashes: Vec, previous_height: Option, header_hash: Bytes32, @@ -254,10 +275,12 @@ pub struct RequestPuzzleState { #[streamable(message)] pub struct RespondPuzzleState { + #[chia(max_length = 33000)] puzzle_hashes: Vec, height: u32, header_hash: Bytes32, is_finished: bool, + #[chia(max_length = 500000)] coin_states: Vec, } @@ -268,6 +291,7 @@ pub struct RejectPuzzleState { #[streamable(message)] pub struct RequestCoinState { + #[chia(max_length = 500000)] coin_ids: Vec, previous_height: Option, header_hash: Bytes32, @@ -276,7 +300,9 @@ pub struct RequestCoinState { #[streamable(message)] pub struct RespondCoinState { + #[chia(max_length = 500000)] coin_ids: Vec, + #[chia(max_length = 500000)] coin_states: Vec, } @@ -334,11 +360,13 @@ pub struct RemovedMempoolItem { #[streamable(message)] pub struct MempoolItemsAdded { + #[chia(max_length = 30000)] transaction_ids: Vec, } #[streamable(message)] pub struct MempoolItemsRemoved { + #[chia(max_length = 30000)] removed_items: Vec, } diff --git a/crates/chia-traits/src/streamable.rs b/crates/chia-traits/src/streamable.rs index 4d7670b45..57de9aa55 100644 --- a/crates/chia-traits/src/streamable.rs +++ b/crates/chia-traits/src/streamable.rs @@ -124,19 +124,28 @@ impl Streamable for Vec { } fn parse(input: &mut Cursor<&[u8]>) -> Result { - let len = u32::parse::(input)?; + parse_vec_with_max_length::(input, u32::MAX as usize) + } +} - let mut ret = if mem::size_of::() == 0 { - Vec::::new() - } else { - let limit = 2 * 1024 * 1024 / mem::size_of::(); - Vec::::with_capacity(std::cmp::min(limit, len as usize)) - }; - for _ in 0..len { - ret.push(T::parse::(input)?); - } - Ok(ret) +pub fn parse_vec_with_max_length( + input: &mut Cursor<&[u8]>, + max_length: usize, +) -> Result> { + let len = u32::parse::(input)?; + if (len as usize) > max_length { + return Err(Error::SequenceTooLarge); } + let mut ret = if mem::size_of::() == 0 { + Vec::::new() + } else { + let limit = 2 * 1024 * 1024 / mem::size_of::(); + Vec::::with_capacity(std::cmp::min(limit, len as usize)) + }; + for _ in 0..len { + ret.push(T::parse::(input)?); + } + Ok(ret) } impl Streamable for String { @@ -784,3 +793,57 @@ fn test_stream_enum() { assert_eq!(stream::(&TestEnum::B), &[1_u8]); assert_eq!(stream::(&TestEnum::C), &[255_u8]); } + +#[cfg(test)] +#[derive(Streamable, PartialEq, Debug)] +struct TestBoundedVec { + #[chia(max_length = 3)] + items: Vec, +} + +#[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::( + 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::(buf, Error::SequenceTooLarge); +} + +#[cfg(test)] +#[derive(Streamable, PartialEq, Debug)] +struct TestBoundedOptionVec { + #[chia(max_length = 2)] + items: Option>, +} + +#[test] +fn test_parse_bounded_option_vec_none() { + let buf: &[u8] = &[0]; + from_bytes::(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::( + 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::(buf, Error::SequenceTooLarge); +} diff --git a/crates/chia_streamable_macro/src/lib.rs b/crates/chia_streamable_macro/src/lib.rs index 15d48106c..8bde37de8 100644 --- a/crates/chia_streamable_macro/src/lib.rs +++ b/crates/chia_streamable_macro/src/lib.rs @@ -6,8 +6,8 @@ use proc_macro2::{Ident, Span, TokenStream as TokenStream2}; use quote::quote; use syn::token::Pub; use syn::{ - Data, DeriveInput, Expr, Fields, FieldsNamed, FieldsUnnamed, Index, Lit, Type, Visibility, - parse_macro_input, + Attribute, Data, DeriveInput, Expr, Fields, FieldsNamed, FieldsUnnamed, GenericArgument, Index, + Lit, PathArguments, Type, Visibility, parse_macro_input, }; #[proc_macro_attribute] @@ -152,7 +152,71 @@ pub fn streamable(attr: TokenStream, item: TokenStream) -> TokenStream { .into() } -#[proc_macro_derive(Streamable)] +fn get_max_length(attrs: &[Attribute]) -> Option { + for attr in attrs { + if attr.path().is_ident("chia") { + let mut max_length = None; + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("max_length") { + let value = meta.value()?; + let lit: syn::LitInt = value.parse()?; + max_length = Some( + lit.base10_parse::() + .expect("max_length must be a valid usize"), + ); + Ok(()) + } else { + Err(meta.error("unsupported chia attribute")) + } + }) + .expect("failed to parse chia attribute"); + return max_length; + } + } + None +} + +fn extract_inner_type<'a>(ty: &'a Type, name: &str) -> Option<&'a Type> { + let Type::Path(type_path) = ty else { + return None; + }; + let segment = type_path.path.segments.last()?; + if segment.ident != name { + return None; + } + let PathArguments::AngleBracketed(args) = &segment.arguments else { + return None; + }; + match args.args.first() { + Some(GenericArgument::Type(inner)) => Some(inner), + _ => None, + } +} + +fn generate_bounded_parse(ty: &Type, max_len: usize, crate_name: &TokenStream2) -> TokenStream2 { + if let Some(inner_type) = extract_inner_type(ty, "Vec") { + quote! { + #crate_name::parse_vec_with_max_length::<#inner_type, TRUSTED>(input, #max_len)? + } + } else if let Some(inner_option_type) = extract_inner_type(ty, "Option") { + let vec_inner_type = extract_inner_type(inner_option_type, "Vec") + .expect("max_length attribute is only supported on Vec or Option> fields"); + quote! { + { + let val = #crate_name::read_bytes(input, 1)?[0]; + match val { + 0 => None, + 1 => Some(#crate_name::parse_vec_with_max_length::<#vec_inner_type, TRUSTED>(input, #max_len)?), + _ => return Err(#crate_name::chia_error::Error::InvalidOptional), + } + } + } + } else { + panic!("max_length attribute is only supported on Vec or Option> fields"); + } +} + +#[proc_macro_derive(Streamable, attributes(chia))] pub fn chia_streamable_macro(input: TokenStream) -> TokenStream { let found_crate = crate_name("chia-traits").expect("chia-traits is present in `Cargo.toml`"); @@ -169,6 +233,7 @@ pub fn chia_streamable_macro(input: TokenStream) -> TokenStream { let mut fnames = Vec::::new(); let mut findices = Vec::::new(); let mut ftypes = Vec::::new(); + let mut max_lengths = Vec::>::new(); match data { Data::Enum(e) => { let mut names = Vec::::new(); @@ -216,6 +281,7 @@ pub fn chia_streamable_macro(input: TokenStream) -> TokenStream { for (index, f) in unnamed.iter().enumerate() { findices.push(Index::from(index)); ftypes.push(f.ty.clone()); + max_lengths.push(get_max_length(&f.attrs)); } } Fields::Unit => {} @@ -223,11 +289,24 @@ pub fn chia_streamable_macro(input: TokenStream) -> TokenStream { for f in &named { fnames.push(f.ident.as_ref().unwrap().clone()); ftypes.push(f.ty.clone()); + max_lengths.push(get_max_length(&f.attrs)); } } }, } + let parse_exprs: Vec = ftypes + .iter() + .zip(max_lengths.iter()) + .map(|(ty, max_len)| { + if let Some(max_len) = max_len { + generate_bounded_parse(ty, *max_len, &crate_name) + } else { + quote!(<#ty as #crate_name::Streamable>::parse::(input)?) + } + }) + .collect(); + if !fnames.is_empty() { let ret = quote! { impl #crate_name::Streamable for #ident { @@ -239,7 +318,7 @@ pub fn chia_streamable_macro(input: TokenStream) -> TokenStream { Ok(()) } fn parse(input: &mut std::io::Cursor<&[u8]>) -> #crate_name::chia_error::Result { - Ok(Self { #( #fnames: <#ftypes as #crate_name::Streamable>::parse::(input)?, )* }) + Ok(Self { #( #fnames: #parse_exprs, )* }) } } }; @@ -255,7 +334,7 @@ pub fn chia_streamable_macro(input: TokenStream) -> TokenStream { Ok(()) } fn parse(input: &mut std::io::Cursor<&[u8]>) -> #crate_name::chia_error::Result { - Ok(Self( #( <#ftypes as #crate_name::Streamable>::parse::(input)?, )* )) + Ok(Self( #( #parse_exprs, )* )) } } };