Skip to main content

reth_eth_wire/
eth_snap.rs

1//! A dedicated `eth` + `snap/2` stream.
2//!
3//! [`EthSnapStream`] carries `eth` as the primary protocol and `snap/2` (EIP-8189) as a typed
4//! side-channel over a single `RLPx` connection. It owns the raw [`P2PStream`] directly and reuses
5//! the transport-free [`EthStreamInner`] codec for `eth` framing, demultiplexing the two protocols
6//! by capability message-id offset.
7
8use 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/// A dedicated stream that carries `eth` as the primary protocol and `snap/2` (EIP-8189) as a typed
37/// side-channel on the same `RLPx` connection.
38///
39/// Both protocols are exposed through the [`Stream`]/[`Sink`] impls, which yield and accept
40/// [`EthSnapMessage`] (an `eth` message or a `snap/2` message). A single poll surfaces whichever
41/// protocol the next inbound frame belongs to, so the owner drives one stream rather than two.
42#[derive(Debug)]
43pub struct EthSnapStream<St, N: NetworkPrimitives = EthNetworkPrimitives> {
44    /// The raw `RLPx` stream carrying both `eth` and `snap/2` frames.
45    conn: P2PStream<St>,
46    /// Transport-free `eth` codec, reused from [`EthStream`](crate::EthStream).
47    eth: EthStreamInner<N>,
48    /// Relative message-id offset of the `snap/2` capability in the combined message space. Ids at
49    /// or above it are snap, rebased to snap-relative and validated by
50    /// [`SnapProtocolMessage::decode_versioned`].
51    snap_offset: u8,
52    /// Responses the remote owes for `snap/2` requests sent on this stream, keyed by request id.
53    ///
54    /// Like the `eth` codec's pending requests, a snap response is only decoded if it answers one
55    /// of these, and the entry is removed when it does, so the map is bounded by the snap requests
56    /// the owner keeps in flight.
57    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    /// Performs the `eth` status handshake over a connection that negotiated both `eth` and
66    /// `snap/2`, returning the established [`EthSnapStream`] and the remote's status.
67    ///
68    /// The handshake runs directly on the underlying [`P2PStream`]; a `snap/2` message arriving
69    /// before the status exchange completes is a protocol violation and surfaces as a handshake
70    /// error.
71    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    /// Returns the negotiated `eth` version.
91    #[inline]
92    pub const fn version(&self) -> EthVersion {
93        self.eth.version()
94    }
95
96    /// Sets whether to reject block announcement messages before RLP decoding.
97    #[inline]
98    pub const fn set_reject_block_announcements(&mut self, reject: bool) {
99        self.eth.set_reject_block_announcements(reject);
100    }
101
102    /// Returns a reference to the underlying [`P2PStream`].
103    #[inline]
104    pub const fn inner(&self) -> &P2PStream<St> {
105        &self.conn
106    }
107
108    /// Returns mutable access to the underlying [`P2PStream`].
109    #[inline]
110    pub const fn inner_mut(&mut self) -> &mut P2PStream<St> {
111        &mut self.conn
112    }
113
114    /// Consumes this type and returns the underlying [`P2PStream`].
115    #[inline]
116    pub fn into_inner(self) -> P2PStream<St> {
117        self.conn
118    }
119
120    /// Consumes the pending request a framed `snap/2` response answers, without decoding the
121    /// response's lists.
122    ///
123    /// Like the `eth` codec, this rejects a response before decoding allocates per list entry if
124    /// nobody asked for it, if it holds more entries than requested, or if its range proof holds
125    /// more than [`MAX_RANGE_PROOF_NODES`] nodes. Messages that are not responses pass through.
126    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        // Every response is a list that starts with the request id.
140        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        // The request sizes the list right after the request id; range responses follow it with
152        // their proof.
153        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    /// Queues an [`EthBroadcastMessage`] to be sent on the wire.
172    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    /// Sends a raw capability message over the connection.
180    ///
181    /// Raw messages are not recorded as requests, so the response to an `eth` or `snap/2` request
182    /// sent this way is rejected as unsolicited. Send requests as [`EthSnapMessage`]s instead.
183    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/// A message carried by an [`EthSnapStream`]: either an `eth` message or a `snap/2` message.
189#[derive(Debug)]
190pub enum EthSnapMessage<N: NetworkPrimitives = EthNetworkPrimitives> {
191    /// An `eth` protocol message.
192    Eth(EthMessage<N>),
193    /// A `snap/2` (EIP-8189) protocol message.
194    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        // `eth` occupies ids below the snap offset and is decoded by the shared codec. Ids at or
217        // above it are snap: rebased to snap-relative (`0x00..`) and validated by
218        // `decode_versioned`, which rejects ids that are out of range or invalid for `snap/2`.
219        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                // Reclaim the freshly-encoded buffer as mutable without copying the payload, then
254                // rebase the snap-relative id into the combined message space.
255                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/// The response the remote owes for a `snap/2` request sent on an [`EthSnapStream`].
274#[derive(Clone, Copy, Debug)]
275struct PendingSnapResponse {
276    /// Message id the response must carry.
277    message_id: SnapMessageId,
278    /// Number of entries the request sized the response's leading list to.
279    max_items: usize,
280}
281
282impl PendingSnapResponse {
283    /// Returns the request id and the response owed for `msg`, or `None` if it is a response.
284    const fn for_request(msg: &SnapProtocolMessage) -> Option<(u64, Self)> {
285        let (message_id, max_items) = match msg {
286            // Accounts are bounded by the requested response bytes, not by a count.
287            SnapProtocolMessage::GetAccountRange(_) => (SnapMessageId::AccountRange, usize::MAX),
288            // One slot list per requested account.
289            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
305/// Maximum number of nodes in a `snap/2` range proof.
306///
307/// The proof covers two boundary paths, each holding at most one hashed node per nibble of a
308/// 32-byte key plus its leaf, since inline nodes are never separate proof entries. The range proof
309/// verifier in `reth-trie-common` rejects longer proofs as well.
310const MAX_RANGE_PROOF_NODES: usize = 2 * (2 * B256::len_bytes() + 1);
311
312/// Resolves the `snap/2` message-id offset from the negotiated capabilities, accepting only a
313/// connection that shares exactly `eth` at relative offset 0 immediately followed by `snap/2`.
314///
315/// The stream forwards every combined id below the snap offset to the eth codec and treats every
316/// id at or above it as snap, so any other shared capability before `eth`, between `eth` and
317/// `snap/2`, or after `snap/2` would be routed incorrectly. Such layouts belong on the
318/// general-purpose satellite multiplexer and are rejected here.
319fn 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/// Rebases a snap-relative message id to the combined message space.
336#[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    /// Connects two peers that negotiate `eth` + `snap/2` and returns their streams.
369    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    /// Sends `msg` from `from` and returns the error `to` reports for it.
406    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    /// Builds shared capabilities from matching local protocols and peer capabilities.
434    fn shared_caps(local: Vec<Protocol>, peer: Vec<Capability>) -> SharedCapabilities {
435        SharedCapabilities::try_new(local, peer).unwrap()
436    }
437
438    /// Builds a hello advertising `eth` + `snap/2`.
439    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        // "aaa" sorts before "eth", so it takes relative offset 0 and eth no longer starts at 0.
464        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        // "les" sorts between "eth" and "snap", leaving a gap so snap no longer directly follows
475        // eth.
476        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        // "zzz" sorts after "snap"; eth+snap still line up, but its frames would be routed
487        // incorrectly as snap, so the layout must be rejected (it belongs on the satellite
488        // multiplexer).
489        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    /// End-to-end: two peers negotiate `eth` + `snap/2`, then a `GetBlockAccessLists` request and
505    /// its `BlockAccessLists` response round-trip over live [`EthSnapStream`]s.
506    #[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        // Server: accept, negotiate, and answer one GetBlockAccessLists echoing the request id.
516        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        // Client: connect, negotiate, send the request, and await the correlated response.
544        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    /// End-to-end: message ids removed by snap/2 are rejected after capability negotiation rather
579    /// than being decoded as trie-node requests or responses.
580    #[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}