1use crate::{
9 capability::SharedCapabilities,
10 errors::{EthStreamError, P2PStreamError},
11 ethstream::list_exceeds_items,
12 handshake::EthRlpxHandshake,
13 message::EthBroadcastMessage,
14 Capability, EthMessage, EthNetworkPrimitives, EthStreamInner, EthVersion, NetworkPrimitives,
15 P2PStream, UnifiedStatus, HANDSHAKE_TIMEOUT,
16};
17use alloy_primitives::{
18 bytes::{Bytes, BytesMut},
19 B256,
20};
21use alloy_rlp::{Decodable, Header};
22use futures::{Sink, SinkExt, Stream, StreamExt};
23use reth_eth_wire_types::{
24 snap::{SnapMessageId, SnapProtocolError, SnapProtocolMessage, SnapVersion},
25 RawCapabilityMessage,
26};
27use reth_ethereum_forks::ForkFilter;
28use std::{
29 collections::BTreeMap,
30 io,
31 pin::Pin,
32 sync::Arc,
33 task::{ready, Context, Poll},
34};
35
36#[derive(Debug)]
43pub struct EthSnapStream<St, N: NetworkPrimitives = EthNetworkPrimitives> {
44 conn: P2PStream<St>,
46 eth: EthStreamInner<N>,
48 snap_offset: u8,
52 pending_snap_responses: BTreeMap<u64, PendingSnapResponse>,
58}
59
60impl<St, N> EthSnapStream<St, N>
61where
62 St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin + Send + Sync,
63 N: NetworkPrimitives,
64{
65 pub async fn handshake(
72 mut conn: P2PStream<St>,
73 status: UnifiedStatus,
74 fork_filter: ForkFilter,
75 handshake: Arc<dyn EthRlpxHandshake>,
76 eth_max_message_size: usize,
77 ) -> Result<(Self, UnifiedStatus), EthStreamError> {
78 let eth_version = conn.shared_capabilities().eth_version()?;
79 let snap_offset = eth_snap_layout(conn.shared_capabilities())?;
80
81 let their_status =
82 handshake.handshake(&mut conn, status, fork_filter, HANDSHAKE_TIMEOUT).await?;
83
84 let eth = EthStreamInner::with_max_message_size(eth_version, eth_max_message_size);
85 Ok((Self { conn, eth, snap_offset, pending_snap_responses: BTreeMap::new() }, their_status))
86 }
87}
88
89impl<St, N: NetworkPrimitives> EthSnapStream<St, N> {
90 #[inline]
92 pub const fn version(&self) -> EthVersion {
93 self.eth.version()
94 }
95
96 #[inline]
98 pub const fn set_reject_block_announcements(&mut self, reject: bool) {
99 self.eth.set_reject_block_announcements(reject);
100 }
101
102 #[inline]
104 pub const fn inner(&self) -> &P2PStream<St> {
105 &self.conn
106 }
107
108 #[inline]
110 pub const fn inner_mut(&mut self) -> &mut P2PStream<St> {
111 &mut self.conn
112 }
113
114 #[inline]
116 pub fn into_inner(self) -> P2PStream<St> {
117 self.conn
118 }
119
120 fn take_pending_snap_response(&mut self, bytes: &[u8]) -> Result<(), SnapProtocolError> {
127 let Some((&id, mut buf)) = bytes.split_first() else { return Ok(()) };
128 let Some(message_id) = [
129 SnapMessageId::AccountRange,
130 SnapMessageId::StorageRanges,
131 SnapMessageId::ByteCodes,
132 SnapMessageId::BlockAccessLists,
133 ]
134 .into_iter()
135 .find(|response| *response as u8 == id) else {
136 return Ok(())
137 };
138
139 let mut payload = Header::decode_bytes(&mut buf, true)?;
141 let request_id = u64::decode(&mut payload)?;
142
143 let Some(pending) = self
144 .pending_snap_responses
145 .remove(&request_id)
146 .filter(|pending| pending.message_id == message_id)
147 else {
148 return Err(SnapProtocolError::UnsolicitedResponse { message_id, request_id })
149 };
150
151 let max_items = if list_exceeds_items(&mut payload, pending.max_items)? {
154 pending.max_items
155 } else if matches!(message_id, SnapMessageId::AccountRange | SnapMessageId::StorageRanges) &&
156 list_exceeds_items(&mut payload, MAX_RANGE_PROOF_NODES)?
157 {
158 MAX_RANGE_PROOF_NODES
159 } else {
160 return Ok(())
161 };
162 Err(SnapProtocolError::ResponseExceedsRequest { message_id, request_id, max_items })
163 }
164}
165
166impl<St, N> EthSnapStream<St, N>
167where
168 St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
169 N: NetworkPrimitives,
170{
171 pub fn start_send_broadcast(
173 &mut self,
174 item: EthBroadcastMessage<N>,
175 ) -> Result<(), EthStreamError> {
176 self.conn.start_send_unpin(item.encoded()).map_err(Into::into)
177 }
178
179 pub fn start_send_raw(&mut self, msg: RawCapabilityMessage) -> Result<(), EthStreamError> {
184 self.conn.start_send_unpin(msg.encoded()).map_err(Into::into)
185 }
186}
187
188#[derive(Debug)]
190pub enum EthSnapMessage<N: NetworkPrimitives = EthNetworkPrimitives> {
191 Eth(EthMessage<N>),
193 Snap(SnapProtocolMessage),
195}
196
197impl<St, N> Stream for EthSnapStream<St, N>
198where
199 St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
200 N: NetworkPrimitives,
201{
202 type Item = Result<EthSnapMessage<N>, EthStreamError>;
203
204 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
205 let this = self.get_mut();
206 let mut bytes = match ready!(this.conn.poll_next_unpin(cx)) {
207 Some(Ok(bytes)) => bytes,
208 Some(Err(err)) => return Poll::Ready(Some(Err(err.into()))),
209 None => return Poll::Ready(None),
210 };
211
212 let Some(&id) = bytes.first() else {
213 return Poll::Ready(Some(Err(P2PStreamError::EmptyProtocolMessage.into())))
214 };
215
216 let Some(snap_id) = id.checked_sub(this.snap_offset) else {
220 return Poll::Ready(Some(this.eth.decode_message(bytes).map(EthSnapMessage::Eth)))
221 };
222 bytes[0] = snap_id;
223 if let Err(err) = this.take_pending_snap_response(&bytes) {
224 return Poll::Ready(Some(Err(err.into())))
225 }
226 Poll::Ready(Some(
227 SnapProtocolMessage::decode_versioned(SnapVersion::V2, &bytes)
228 .map(EthSnapMessage::Snap)
229 .map_err(Into::into),
230 ))
231 }
232}
233
234impl<St, N> Sink<EthSnapMessage<N>> for EthSnapStream<St, N>
235where
236 St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
237 N: NetworkPrimitives,
238{
239 type Error = EthStreamError;
240
241 fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
242 self.get_mut().conn.poll_ready_unpin(cx).map_err(Into::into)
243 }
244
245 fn start_send(self: Pin<&mut Self>, item: EthSnapMessage<N>) -> Result<(), Self::Error> {
246 let this = self.get_mut();
247 let bytes = match item {
248 EthSnapMessage::Eth(msg) => this.eth.encode_message(msg)?,
249 EthSnapMessage::Snap(msg) => {
250 if let Some((request_id, response)) = PendingSnapResponse::for_request(&msg) {
251 this.pending_snap_responses.insert(request_id, response);
252 }
253 let mut buf =
256 msg.encode().0.try_into_mut().unwrap_or_else(|b| BytesMut::from(b.as_ref()));
257 mask_snap(&mut buf, this.snap_offset)?;
258 buf.freeze()
259 }
260 };
261 this.conn.start_send_unpin(bytes).map_err(Into::into)
262 }
263
264 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
265 self.get_mut().conn.poll_flush_unpin(cx).map_err(Into::into)
266 }
267
268 fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
269 self.get_mut().conn.poll_close_unpin(cx).map_err(Into::into)
270 }
271}
272
273#[derive(Clone, Copy, Debug)]
275struct PendingSnapResponse {
276 message_id: SnapMessageId,
278 max_items: usize,
280}
281
282impl PendingSnapResponse {
283 const fn for_request(msg: &SnapProtocolMessage) -> Option<(u64, Self)> {
285 let (message_id, max_items) = match msg {
286 SnapProtocolMessage::GetAccountRange(_) => (SnapMessageId::AccountRange, usize::MAX),
288 SnapProtocolMessage::GetStorageRanges(req) => {
290 (SnapMessageId::StorageRanges, req.account_hashes.len())
291 }
292 SnapProtocolMessage::GetByteCodes(req) => (SnapMessageId::ByteCodes, req.hashes.len()),
293 SnapProtocolMessage::GetBlockAccessLists(req) => {
294 (SnapMessageId::BlockAccessLists, req.block_hashes.len())
295 }
296 SnapProtocolMessage::AccountRange(_) |
297 SnapProtocolMessage::StorageRanges(_) |
298 SnapProtocolMessage::ByteCodes(_) |
299 SnapProtocolMessage::BlockAccessLists(_) => return None,
300 };
301 Some((msg.request_id(), Self { message_id, max_items }))
302 }
303}
304
305const MAX_RANGE_PROOF_NODES: usize = 2 * (2 * B256::len_bytes() + 1);
311
312fn eth_snap_layout(caps: &SharedCapabilities) -> Result<u8, EthStreamError> {
320 if !caps.is_exact_eth_snap_v2() {
321 return Err(P2PStreamError::CapabilityNotShared.into())
322 }
323 let snap = caps
324 .ensure_matching_capability(&Capability::snap_2())
325 .map_err(|_| EthStreamError::from(P2PStreamError::CapabilityNotShared))?;
326 let snap_offset = snap.relative_message_id_offset();
327
328 let eth = caps.eth()?;
329 if eth.relative_message_id_offset() != 0 || eth.num_messages() != snap_offset {
330 return Err(P2PStreamError::CapabilityNotShared.into())
331 }
332 Ok(snap_offset)
333}
334
335#[inline]
337fn mask_snap(bytes: &mut BytesMut, snap_offset: u8) -> Result<(), io::Error> {
338 let id = bytes.first().ok_or_else(|| io::Error::from(io::ErrorKind::InvalidInput))?;
339 bytes[0] =
340 id.checked_add(snap_offset).ok_or_else(|| io::Error::from(io::ErrorKind::InvalidInput))?;
341 Ok(())
342}
343
344#[cfg(test)]
345mod tests {
346 use super::*;
347 use crate::{
348 handshake::EthHandshake,
349 message::MAX_MESSAGE_SIZE,
350 protocol::Protocol,
351 test_utils::{connect_passthrough, eth_handshake, eth_hello},
352 UnauthedP2PStream,
353 };
354 use reth_eth_wire_types::{
355 snap::{
356 AccountRangeMessage, BlockAccessListsMessage, ByteCodesMessage, GetAccountRangeMessage,
357 GetBlockAccessListsMessage, GetByteCodesMessage, GetStorageRangesMessage,
358 StorageRangesMessage,
359 },
360 BlockAccessLists, EthVersion,
361 };
362 use tokio::net::{TcpListener, TcpStream};
363 use tokio_util::codec::{Decoder, Framed};
364
365 type TestEthSnapStream =
366 EthSnapStream<Framed<TcpStream, crate::PassthroughCodec>, EthNetworkPrimitives>;
367
368 async fn eth_snap_pair() -> (TestEthSnapStream, TestEthSnapStream) {
370 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
371 let local_addr = listener.local_addr().unwrap();
372 let (status, fork_filter) = eth_handshake();
373 let server_fork_filter = fork_filter.clone();
374
375 let server = tokio::spawn(async move {
376 let (incoming, _) = listener.accept().await.unwrap();
377 let stream = crate::PassthroughCodec::default().framed(incoming);
378 let (conn, _) =
379 UnauthedP2PStream::new(stream).handshake(eth_snap_hello()).await.unwrap();
380 EthSnapStream::handshake(
381 conn,
382 status,
383 server_fork_filter,
384 Arc::new(EthHandshake::default()),
385 MAX_MESSAGE_SIZE,
386 )
387 .await
388 .unwrap()
389 .0
390 });
391
392 let conn = connect_passthrough(local_addr, eth_snap_hello()).await;
393 let (client, _) = EthSnapStream::handshake(
394 conn,
395 status,
396 fork_filter,
397 Arc::new(EthHandshake::default()),
398 MAX_MESSAGE_SIZE,
399 )
400 .await
401 .unwrap();
402 (client, server.await.unwrap())
403 }
404
405 async fn send_and_expect_error(
407 from: &mut TestEthSnapStream,
408 to: &mut TestEthSnapStream,
409 msg: SnapProtocolMessage,
410 ) -> SnapProtocolError {
411 from.send(EthSnapMessage::Snap(msg)).await.unwrap();
412 match to.next().await.unwrap() {
413 Err(EthStreamError::InvalidSnapMessage(err)) => err,
414 other => panic!("expected a snap error, got {other:?}"),
415 }
416 }
417
418 fn get_byte_codes(request_id: u64, hashes: usize) -> SnapProtocolMessage {
419 SnapProtocolMessage::GetByteCodes(GetByteCodesMessage {
420 request_id,
421 hashes: vec![B256::ZERO; hashes],
422 response_bytes: u64::MAX,
423 })
424 }
425
426 fn byte_codes(request_id: u64, codes: usize) -> SnapProtocolMessage {
427 SnapProtocolMessage::ByteCodes(ByteCodesMessage {
428 request_id,
429 codes: vec![Bytes::new().into(); codes],
430 })
431 }
432
433 fn shared_caps(local: Vec<Protocol>, peer: Vec<Capability>) -> SharedCapabilities {
435 SharedCapabilities::try_new(local, peer).unwrap()
436 }
437
438 fn eth_snap_hello() -> crate::HelloMessageWithProtocols {
440 let mut hello = eth_hello().0;
441 hello.try_add_protocol(Protocol::snap_2()).unwrap();
442 hello
443 }
444
445 #[test]
446 fn eth_snap_layout_accepts_eth_then_snap() {
447 let caps = shared_caps(
448 vec![EthVersion::Eth68.into(), Protocol::snap_2()],
449 vec![EthVersion::Eth68.into(), Capability::snap_2()],
450 );
451 let offset = eth_snap_layout(&caps).unwrap();
452 assert_eq!(offset, caps.eth().unwrap().num_messages());
453 }
454
455 #[test]
456 fn eth_snap_layout_rejects_missing_snap() {
457 let caps = shared_caps(vec![EthVersion::Eth68.into()], vec![EthVersion::Eth68.into()]);
458 assert!(eth_snap_layout(&caps).is_err());
459 }
460
461 #[test]
462 fn eth_snap_layout_rejects_capability_before_eth() {
463 let cap = Capability::new_static("aaa", 1);
465 let caps = shared_caps(
466 vec![Protocol::new(cap.clone(), 5), EthVersion::Eth68.into(), Protocol::snap_2()],
467 vec![cap, EthVersion::Eth68.into(), Capability::snap_2()],
468 );
469 assert!(eth_snap_layout(&caps).is_err());
470 }
471
472 #[test]
473 fn eth_snap_layout_rejects_capability_between_eth_and_snap() {
474 let cap = Capability::new_static("les", 1);
477 let caps = shared_caps(
478 vec![EthVersion::Eth68.into(), Protocol::new(cap.clone(), 5), Protocol::snap_2()],
479 vec![EthVersion::Eth68.into(), cap, Capability::snap_2()],
480 );
481 assert!(eth_snap_layout(&caps).is_err());
482 }
483
484 #[test]
485 fn eth_snap_layout_rejects_capability_after_snap() {
486 let cap = Capability::new_static("zzz", 1);
490 let caps = shared_caps(
491 vec![EthVersion::Eth68.into(), Protocol::snap_2(), Protocol::new(cap.clone(), 5)],
492 vec![EthVersion::Eth68.into(), Capability::snap_2(), cap],
493 );
494 assert!(eth_snap_layout(&caps).is_err());
495 }
496
497 #[test]
498 fn mask_snap_rebases_into_combined_space() {
499 let mut bytes = BytesMut::from(&[8u8, 0xcc][..]);
500 mask_snap(&mut bytes, 17).unwrap();
501 assert_eq!(bytes[0], 17 + 8);
502 }
503
504 #[tokio::test(flavor = "multi_thread")]
507 async fn snap_request_response_round_trips_over_the_wire() {
508 reth_tracing::init_test_tracing();
509 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
510 let local_addr = listener.local_addr().unwrap();
511 let (status, fork_filter) = eth_handshake();
512 let server_status = status;
513 let server_fork_filter = fork_filter.clone();
514
515 let server = tokio::spawn(async move {
517 let (incoming, _) = listener.accept().await.unwrap();
518 let stream = crate::PassthroughCodec::default().framed(incoming);
519 let server_hello = eth_snap_hello();
520 let (conn, _) = UnauthedP2PStream::new(stream).handshake(server_hello).await.unwrap();
521
522 let (mut stream, _) = EthSnapStream::<_, EthNetworkPrimitives>::handshake(
523 conn,
524 server_status,
525 server_fork_filter,
526 Arc::new(EthHandshake::default()),
527 MAX_MESSAGE_SIZE,
528 )
529 .await
530 .unwrap();
531
532 while let Some(Ok(msg)) = stream.next().await {
533 if let EthSnapMessage::Snap(SnapProtocolMessage::GetBlockAccessLists(req)) = msg {
534 let response = SnapProtocolMessage::BlockAccessLists(BlockAccessListsMessage {
535 request_id: req.request_id,
536 block_access_lists: reth_eth_wire_types::BlockAccessLists(vec![None]),
537 });
538 stream.send(EthSnapMessage::Snap(response)).await.unwrap();
539 }
540 }
541 });
542
543 let conn = connect_passthrough(local_addr, eth_snap_hello()).await;
545 let (mut stream, _) = EthSnapStream::<_, EthNetworkPrimitives>::handshake(
546 conn,
547 status,
548 fork_filter,
549 Arc::new(EthHandshake::default()),
550 MAX_MESSAGE_SIZE,
551 )
552 .await
553 .unwrap();
554
555 stream
556 .send(EthSnapMessage::Snap(SnapProtocolMessage::GetBlockAccessLists(
557 GetBlockAccessListsMessage {
558 request_id: 7,
559 block_hashes: vec![B256::ZERO],
560 response_bytes: u64::MAX,
561 },
562 )))
563 .await
564 .unwrap();
565
566 let response = loop {
567 if let EthSnapMessage::Snap(SnapProtocolMessage::BlockAccessLists(resp)) =
568 stream.next().await.unwrap().unwrap()
569 {
570 break resp;
571 }
572 };
573 assert_eq!(response.request_id, 7);
574
575 server.abort();
576 }
577
578 #[tokio::test(flavor = "multi_thread")]
581 async fn snap_two_rejects_removed_trie_node_messages_over_the_wire() {
582 reth_tracing::init_test_tracing();
583
584 for removed_snap_id in [0x06, 0x07] {
585 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
586 let local_addr = listener.local_addr().unwrap();
587 let (status, fork_filter) = eth_handshake();
588 let server_status = status;
589 let server_fork_filter = fork_filter.clone();
590
591 let server = tokio::spawn(async move {
592 let (incoming, _) = listener.accept().await.unwrap();
593 let stream = crate::PassthroughCodec::default().framed(incoming);
594 let (conn, _) =
595 UnauthedP2PStream::new(stream).handshake(eth_snap_hello()).await.unwrap();
596 let (mut stream, _) = EthSnapStream::<_, EthNetworkPrimitives>::handshake(
597 conn,
598 server_status,
599 server_fork_filter,
600 Arc::new(EthHandshake::default()),
601 MAX_MESSAGE_SIZE,
602 )
603 .await
604 .unwrap();
605
606 stream.next().await.expect("peer should send a frame").unwrap_err()
607 });
608
609 let conn = connect_passthrough(local_addr, eth_snap_hello()).await;
610 let (mut stream, _) = EthSnapStream::<_, EthNetworkPrimitives>::handshake(
611 conn,
612 status,
613 fork_filter,
614 Arc::new(EthHandshake::default()),
615 MAX_MESSAGE_SIZE,
616 )
617 .await
618 .unwrap();
619
620 let combined_id = stream.snap_offset as usize + removed_snap_id;
621 stream
622 .start_send_raw(RawCapabilityMessage::new(
623 combined_id,
624 Bytes::from_static(&[alloy_rlp::EMPTY_LIST_CODE]),
625 ))
626 .unwrap();
627 stream.flush().await.unwrap();
628
629 let error = server.await.unwrap();
630 assert!(
631 error.to_string().contains("invalid for snap"),
632 "unexpected error for removed snap id {removed_snap_id:#x}: {error}"
633 );
634 }
635 }
636
637 #[tokio::test(flavor = "multi_thread")]
638 async fn unsolicited_snap_response_is_rejected() {
639 let (mut client, mut server) = eth_snap_pair().await;
640
641 let err = send_and_expect_error(&mut client, &mut server, byte_codes(7, 1)).await;
642 assert!(matches!(
643 err,
644 SnapProtocolError::UnsolicitedResponse {
645 message_id: SnapMessageId::ByteCodes,
646 request_id: 7
647 }
648 ));
649 }
650
651 #[tokio::test(flavor = "multi_thread")]
652 async fn snap_response_to_sent_request_is_decoded_once() {
653 let (mut client, mut server) = eth_snap_pair().await;
654 server.send(EthSnapMessage::Snap(get_byte_codes(7, 2))).await.unwrap();
655 assert!(matches!(
656 client.next().await.unwrap().unwrap(),
657 EthSnapMessage::Snap(SnapProtocolMessage::GetByteCodes(_))
658 ));
659
660 client.send(EthSnapMessage::Snap(byte_codes(7, 2))).await.unwrap();
661 assert!(matches!(
662 server.next().await.unwrap().unwrap(),
663 EthSnapMessage::Snap(SnapProtocolMessage::ByteCodes(_))
664 ));
665
666 let err = send_and_expect_error(&mut client, &mut server, byte_codes(7, 2)).await;
667 assert!(matches!(
668 err,
669 SnapProtocolError::UnsolicitedResponse {
670 message_id: SnapMessageId::ByteCodes,
671 request_id: 7
672 }
673 ));
674 }
675
676 #[tokio::test(flavor = "multi_thread")]
677 async fn snap_response_of_other_type_is_rejected() {
678 let (mut client, mut server) = eth_snap_pair().await;
679 server.send(EthSnapMessage::Snap(get_byte_codes(7, 1))).await.unwrap();
680 client.next().await.unwrap().unwrap();
681
682 let response = SnapProtocolMessage::BlockAccessLists(BlockAccessListsMessage {
683 request_id: 7,
684 block_access_lists: BlockAccessLists(Vec::new()),
685 });
686 let err = send_and_expect_error(&mut client, &mut server, response).await;
687 assert!(matches!(
688 err,
689 SnapProtocolError::UnsolicitedResponse {
690 message_id: SnapMessageId::BlockAccessLists,
691 request_id: 7
692 }
693 ));
694 }
695
696 #[tokio::test(flavor = "multi_thread")]
697 async fn snap_response_with_more_entries_than_requested_is_rejected() {
698 let (mut client, mut server) = eth_snap_pair().await;
699 server.send(EthSnapMessage::Snap(get_byte_codes(7, 2))).await.unwrap();
700 client.next().await.unwrap().unwrap();
701
702 let err = send_and_expect_error(&mut client, &mut server, byte_codes(7, 3)).await;
703 assert!(matches!(
704 err,
705 SnapProtocolError::ResponseExceedsRequest {
706 message_id: SnapMessageId::ByteCodes,
707 request_id: 7,
708 max_items: 2
709 }
710 ));
711 }
712
713 #[tokio::test(flavor = "multi_thread")]
714 async fn snap_range_proof_is_bounded() {
715 let (mut client, mut server) = eth_snap_pair().await;
716 let request = |request_id| {
717 SnapProtocolMessage::GetAccountRange(GetAccountRangeMessage {
718 request_id,
719 root_hash: B256::ZERO,
720 starting_hash: B256::ZERO,
721 limit_hash: B256::repeat_byte(0xff),
722 response_bytes: u64::MAX,
723 })
724 };
725 let response = |request_id, proof_nodes| {
726 SnapProtocolMessage::AccountRange(AccountRangeMessage {
727 request_id,
728 accounts: Vec::new(),
729 proof: vec![Bytes::new().into(); proof_nodes],
730 })
731 };
732
733 server.send(EthSnapMessage::Snap(request(1))).await.unwrap();
734 client.next().await.unwrap().unwrap();
735 client.send(EthSnapMessage::Snap(response(1, MAX_RANGE_PROOF_NODES))).await.unwrap();
736 assert!(matches!(
737 server.next().await.unwrap().unwrap(),
738 EthSnapMessage::Snap(SnapProtocolMessage::AccountRange(_))
739 ));
740
741 server.send(EthSnapMessage::Snap(request(2))).await.unwrap();
742 client.next().await.unwrap().unwrap();
743 let err =
744 send_and_expect_error(&mut client, &mut server, response(2, MAX_RANGE_PROOF_NODES + 1))
745 .await;
746 assert!(matches!(
747 err,
748 SnapProtocolError::ResponseExceedsRequest {
749 message_id: SnapMessageId::AccountRange,
750 request_id: 2,
751 max_items: MAX_RANGE_PROOF_NODES
752 }
753 ));
754 }
755
756 #[tokio::test(flavor = "multi_thread")]
757 async fn snap_storage_ranges_hold_one_slot_list_per_account() {
758 let (mut client, mut server) = eth_snap_pair().await;
759 server
760 .send(EthSnapMessage::Snap(SnapProtocolMessage::GetStorageRanges(
761 GetStorageRangesMessage {
762 request_id: 7,
763 root_hash: B256::ZERO,
764 account_hashes: vec![B256::ZERO],
765 starting_hash: B256::ZERO.into(),
766 limit_hash: B256::repeat_byte(0xff).into(),
767 response_bytes: u64::MAX,
768 },
769 )))
770 .await
771 .unwrap();
772 client.next().await.unwrap().unwrap();
773
774 let response = SnapProtocolMessage::StorageRanges(StorageRangesMessage {
775 request_id: 7,
776 slots: vec![Vec::new(); 2],
777 proof: Vec::new(),
778 });
779 let err = send_and_expect_error(&mut client, &mut server, response).await;
780 assert!(matches!(
781 err,
782 SnapProtocolError::ResponseExceedsRequest {
783 message_id: SnapMessageId::StorageRanges,
784 request_id: 7,
785 max_items: 1
786 }
787 ));
788 }
789}