@@ -177,16 +177,18 @@ where
177177mod tests {
178178 use super :: super :: ssz_snappy:: * ;
179179 use super :: * ;
180- use crate :: rpc:: methods:: StatusMessage ;
181180 use crate :: rpc:: protocol:: * ;
182- use snap:: write:: FrameEncoder ;
183- use ssz:: Encode ;
184- use std:: io:: Write ;
185- use types:: { Epoch , Hash256 , Slot } ;
181+
182+ use std:: sync:: Arc ;
183+ use types:: { ForkContext , Hash256 } ;
186184 use unsigned_varint:: codec:: Uvi ;
187185
188186 type Spec = types:: MainnetEthSpec ;
189187
188+ fn fork_context ( ) -> ForkContext {
189+ ForkContext :: new ( Hash256 :: zero ( ) , & Spec :: default_spec ( ) )
190+ }
191+
190192 #[ test]
191193 fn test_decode_status_message ( ) {
192194 let message = hex:: decode ( "0054ff060000734e615070590032000006e71e7b54989925efd6c9cbcb8ceb9b5f71216f5137282bf6a1e3b50f64e42d6c7fb347abe07eb0db8200000005029e2800" ) . unwrap ( ) ;
@@ -196,8 +198,9 @@ mod tests {
196198 let snappy_protocol_id =
197199 ProtocolId :: new ( Protocol :: Status , Version :: V1 , Encoding :: SSZSnappy ) ;
198200
201+ let fork_context = Arc :: new ( fork_context ( ) ) ;
199202 let mut snappy_outbound_codec =
200- SSZSnappyOutboundCodec :: < Spec > :: new ( snappy_protocol_id, 1_048_576 ) ;
203+ SSZSnappyOutboundCodec :: < Spec > :: new ( snappy_protocol_id, 1_048_576 , fork_context ) ;
201204
202205 // remove response code
203206 let mut snappy_buf = buf. clone ( ) ;
@@ -229,8 +232,10 @@ mod tests {
229232
230233 let snappy_protocol_id =
231234 ProtocolId :: new ( Protocol :: Status , Version :: V1 , Encoding :: SSZSnappy ) ;
235+
236+ let fork_context = Arc :: new ( fork_context ( ) ) ;
232237 let mut snappy_outbound_codec =
233- SSZSnappyOutboundCodec :: < Spec > :: new ( snappy_protocol_id, 1_048_576 ) ;
238+ SSZSnappyOutboundCodec :: < Spec > :: new ( snappy_protocol_id, 1_048_576 , fork_context ) ;
234239
235240 let snappy_decoded_message = snappy_outbound_codec. decode ( & mut dst) . unwrap_err ( ) ;
236241
@@ -256,80 +261,34 @@ mod tests {
256261 // Response limits
257262 let limit = protocol_id. rpc_response_limits :: < Spec > ( ) ;
258263 let mut max = encode_len ( limit. max + 1 ) ;
259- let mut codec = SSZSnappyOutboundCodec :: < Spec > :: new ( protocol_id. clone ( ) , 1_048_576 ) ;
264+ let fork_context = Arc :: new ( fork_context ( ) ) ;
265+ let mut codec = SSZSnappyOutboundCodec :: < Spec > :: new (
266+ protocol_id. clone ( ) ,
267+ 1_048_576 ,
268+ fork_context. clone ( ) ,
269+ ) ;
260270 assert_eq ! ( codec. decode( & mut max) . unwrap_err( ) , RPCError :: InvalidData ) ;
261271
262272 let mut min = encode_len ( limit. min - 1 ) ;
263- let mut codec = SSZSnappyOutboundCodec :: < Spec > :: new ( protocol_id. clone ( ) , 1_048_576 ) ;
273+ let mut codec = SSZSnappyOutboundCodec :: < Spec > :: new (
274+ protocol_id. clone ( ) ,
275+ 1_048_576 ,
276+ fork_context. clone ( ) ,
277+ ) ;
264278 assert_eq ! ( codec. decode( & mut min) . unwrap_err( ) , RPCError :: InvalidData ) ;
265279
266280 // Request limits
267281 let limit = protocol_id. rpc_request_limits ( ) ;
268282 let mut max = encode_len ( limit. max + 1 ) ;
269- let mut codec = SSZSnappyOutboundCodec :: < Spec > :: new ( protocol_id. clone ( ) , 1_048_576 ) ;
283+ let mut codec = SSZSnappyOutboundCodec :: < Spec > :: new (
284+ protocol_id. clone ( ) ,
285+ 1_048_576 ,
286+ fork_context. clone ( ) ,
287+ ) ;
270288 assert_eq ! ( codec. decode( & mut max) . unwrap_err( ) , RPCError :: InvalidData ) ;
271289
272290 let mut min = encode_len ( limit. min - 1 ) ;
273- let mut codec = SSZSnappyOutboundCodec :: < Spec > :: new ( protocol_id, 1_048_576 ) ;
291+ let mut codec = SSZSnappyOutboundCodec :: < Spec > :: new ( protocol_id, 1_048_576 , fork_context ) ;
274292 assert_eq ! ( codec. decode( & mut min) . unwrap_err( ) , RPCError :: InvalidData ) ;
275293 }
276-
277- #[ test]
278- fn test_decode_malicious_status_message ( ) {
279- // 10 byte snappy stream identifier
280- let stream_identifier: & ' static [ u8 ] = b"\xFF \x06 \x00 \x00 sNaPpY" ;
281-
282- assert_eq ! ( stream_identifier. len( ) , 10 ) ;
283-
284- // byte 0(0xFE) is padding chunk type identifier for snappy messages
285- // byte 1,2,3 are chunk length (little endian)
286- let malicious_padding: & ' static [ u8 ] = b"\xFE \x00 \x00 \x00 " ;
287-
288- // Status message is 84 bytes uncompressed. `max_compressed_len` is 32 + 84 + 84/6 = 130.
289- let status_message_bytes = StatusMessage {
290- fork_digest : [ 0 ; 4 ] ,
291- finalized_root : Hash256 :: from_low_u64_be ( 0 ) ,
292- finalized_epoch : Epoch :: new ( 1 ) ,
293- head_root : Hash256 :: from_low_u64_be ( 0 ) ,
294- head_slot : Slot :: new ( 1 ) ,
295- }
296- . as_ssz_bytes ( ) ;
297-
298- assert_eq ! ( status_message_bytes. len( ) , 84 ) ;
299- assert_eq ! ( snap:: raw:: max_compress_len( status_message_bytes. len( ) ) , 130 ) ;
300-
301- let mut uvi_codec: Uvi < usize > = Uvi :: default ( ) ;
302- let mut dst = BytesMut :: with_capacity ( 1024 ) ;
303-
304- // Insert length-prefix
305- uvi_codec
306- . encode ( status_message_bytes. len ( ) , & mut dst)
307- . unwrap ( ) ;
308-
309- // Insert snappy stream identifier
310- dst. extend_from_slice ( stream_identifier) ;
311-
312- // Insert malicious padding of 80 bytes.
313- for _ in 0 ..20 {
314- dst. extend_from_slice ( malicious_padding) ;
315- }
316-
317- // Insert payload (42 bytes compressed)
318- let mut writer = FrameEncoder :: new ( Vec :: new ( ) ) ;
319- writer. write_all ( & status_message_bytes) . unwrap ( ) ;
320- writer. flush ( ) . unwrap ( ) ;
321- assert_eq ! ( writer. get_ref( ) . len( ) , 42 ) ;
322- dst. extend_from_slice ( writer. get_ref ( ) ) ;
323-
324- // 10 (for stream identifier) + 80 + 42 = 132 > `max_compressed_len`. Hence, decoding should fail with `InvalidData`.
325-
326- let snappy_protocol_id =
327- ProtocolId :: new ( Protocol :: Status , Version :: V1 , Encoding :: SSZSnappy ) ;
328-
329- let mut snappy_outbound_codec =
330- SSZSnappyOutboundCodec :: < Spec > :: new ( snappy_protocol_id, 1_048_576 ) ;
331-
332- let snappy_decoded_message = snappy_outbound_codec. decode ( & mut dst) . unwrap_err ( ) ;
333- assert_eq ! ( snappy_decoded_message, RPCError :: InvalidData ) ;
334- }
335294}
0 commit comments