diff --git a/rust/relay/src/proptest.rs b/rust/relay/src/proptest.rs index d79c8271c..92dfd13b7 100644 --- a/rust/relay/src/proptest.rs +++ b/rust/relay/src/proptest.rs @@ -1,6 +1,5 @@ use crate::Binding; use proptest::arbitrary::any; -use proptest::collection::vec; use proptest::strategy::Just; use proptest::strategy::Strategy; use proptest::string::string_regex; @@ -30,11 +29,16 @@ pub fn channel_number() -> impl Strategy { (ChannelNumber::MIN..ChannelNumber::MAX).prop_map(|n| ChannelNumber::new(n).unwrap()) } -pub fn channel_payload() -> impl Strategy, usize)> { - vec(any::(), 0..(u16::MAX as usize)).prop_flat_map(|vec| { - let len = vec.len(); - (Just(vec), 0..len) - }) +pub fn channel_payload() -> impl Strategy, u16)> { + any::>() + .prop_filter("payload does not fit into u16", |vec| { + vec.len() <= u16::MAX as usize + }) + .prop_map(|vec| { + let len = vec.len(); + + (vec, len as u16) + }) } pub fn username_salt() -> impl Strategy { diff --git a/rust/relay/src/server/channel_data.rs b/rust/relay/src/server/channel_data.rs index 7aeca1b6a..d1781ab3a 100644 --- a/rust/relay/src/server/channel_data.rs +++ b/rust/relay/src/server/channel_data.rs @@ -97,12 +97,12 @@ mod tests { #[test_strategy::proptest] fn channel_data_decoding( #[strategy(crate::proptest::channel_number())] channel: ChannelNumber, - #[strategy(crate::proptest::channel_payload())] payload: (Vec, usize), + #[strategy(crate::proptest::channel_payload())] payload: (Vec, u16), ) { - let encoded = to_bytes(channel.value(), payload.1 as u16, &payload.0); + let encoded = to_bytes(channel.value(), payload.1, &payload.0); let parsed = ChannelData::parse(&encoded).unwrap(); assert_eq!(channel.value(), parsed.channel); - assert_eq!(&payload.0[..payload.1], parsed.data) + assert_eq!(&payload.0[..(payload.1 as usize)], parsed.data) } }