Skip to main content

reth_eth_wire/
ethstream.rs

1//! Ethereum protocol stream implementations.
2//!
3//! Provides stream types for the Ethereum wire protocol.
4//! It separates protocol logic [`EthStreamInner`] from transport concerns [`EthStream`].
5//! Handles handshaking, message processing, and RLP serialization.
6
7use crate::{
8    errors::{EthHandshakeError, EthStreamError},
9    handshake::EthereumEthHandshake,
10    message::{EthBroadcastMessage, MessageError, MAX_MESSAGE_SIZE, TX_MEMORY_BUDGET_MULTIPLIER},
11    p2pstream::HANDSHAKE_TIMEOUT,
12    CanDisconnect, DisconnectReason, EthMessage, EthNetworkPrimitives, EthVersion, ProtocolMessage,
13    UnifiedStatus,
14};
15use alloy_primitives::bytes::{Bytes, BytesMut};
16use alloy_rlp::{Decodable, Header};
17use futures::{ready, Sink, SinkExt};
18use pin_project::pin_project;
19use reth_eth_wire_types::{EthMessageID, NetworkPrimitives, RawCapabilityMessage};
20use reth_ethereum_forks::ForkFilter;
21use std::{
22    collections::BTreeMap,
23    future::Future,
24    pin::Pin,
25    task::{Context, Poll},
26    time::Duration,
27};
28use tokio::time::timeout;
29use tokio_stream::Stream;
30use tracing::{debug, trace};
31
32/// An un-authenticated [`EthStream`]. This is consumed and returns a [`EthStream`] after the
33/// `Status` handshake is completed.
34#[pin_project]
35#[derive(Debug)]
36pub struct UnauthedEthStream<S> {
37    #[pin]
38    inner: S,
39}
40
41impl<S> UnauthedEthStream<S> {
42    /// Create a new `UnauthedEthStream` from a type `S` which implements `Stream` and `Sink`.
43    pub const fn new(inner: S) -> Self {
44        Self { inner }
45    }
46
47    /// Consumes the type and returns the wrapped stream
48    pub fn into_inner(self) -> S {
49        self.inner
50    }
51}
52
53impl<S, E> UnauthedEthStream<S>
54where
55    S: Stream<Item = Result<BytesMut, E>> + CanDisconnect<Bytes> + Send,
56    EthStreamError: From<E> + From<<S as Sink<Bytes>>::Error>,
57{
58    /// Consumes the [`UnauthedEthStream`] and returns an [`EthStream`] after the `Status`
59    /// handshake is completed successfully. This also returns the `Status` message sent by the
60    /// remote peer.
61    ///
62    /// Caution: This expects that the [`UnifiedStatus`] has the proper eth version configured, with
63    /// ETH69 the initial status message changed.
64    pub async fn handshake<N: NetworkPrimitives>(
65        self,
66        status: UnifiedStatus,
67        fork_filter: ForkFilter,
68    ) -> Result<(EthStream<S, N>, UnifiedStatus), EthStreamError> {
69        self.handshake_with_timeout(status, fork_filter, HANDSHAKE_TIMEOUT).await
70    }
71
72    /// Wrapper around handshake which enforces a timeout.
73    pub async fn handshake_with_timeout<N: NetworkPrimitives>(
74        self,
75        status: UnifiedStatus,
76        fork_filter: ForkFilter,
77        timeout_limit: Duration,
78    ) -> Result<(EthStream<S, N>, UnifiedStatus), EthStreamError> {
79        timeout(timeout_limit, Self::handshake_without_timeout(self, status, fork_filter))
80            .await
81            .map_err(|_| EthStreamError::StreamTimeout)?
82    }
83
84    /// Handshake with no timeout
85    pub async fn handshake_without_timeout<N: NetworkPrimitives>(
86        mut self,
87        status: UnifiedStatus,
88        fork_filter: ForkFilter,
89    ) -> Result<(EthStream<S, N>, UnifiedStatus), EthStreamError> {
90        trace!(
91            status = %status.into_message(),
92            "sending eth status to peer"
93        );
94        let their_status =
95            EthereumEthHandshake(&mut self.inner).eth_handshake(status, fork_filter).await?;
96
97        // now we can create the `EthStream` because the peer has successfully completed
98        // the handshake
99        let stream = EthStream::new(status.version, self.inner);
100
101        Ok((stream, their_status))
102    }
103}
104
105/// Contains eth protocol specific logic for processing messages
106#[derive(Debug)]
107pub struct EthStreamInner<N> {
108    /// Negotiated eth version
109    version: EthVersion,
110    /// Maximum allowed ETH message size.
111    max_message_size: usize,
112    /// When true, `NewBlock` (0x07) and `NewBlockHashes` (0x01) messages are rejected before RLP
113    /// decoding to avoid any memory impact for non-PoW networks.
114    reject_block_announcements: bool,
115    /// Responses the remote owes for requests encoded by [`Self::encode_message`], keyed by
116    /// request id.
117    ///
118    /// [`Self::decode_message`] only decodes a response that answers one of these requests, and
119    /// removes the entry when it does. The map is therefore bounded by the requests the owner
120    /// keeps in flight.
121    pending_responses: BTreeMap<u64, PendingResponse>,
122    _pd: std::marker::PhantomData<N>,
123}
124
125impl<N> EthStreamInner<N>
126where
127    N: NetworkPrimitives,
128{
129    /// Creates a new [`EthStreamInner`] with the given eth version
130    pub const fn new(version: EthVersion) -> Self {
131        Self::with_max_message_size(version, MAX_MESSAGE_SIZE)
132    }
133
134    /// Creates a new [`EthStreamInner`] with the given eth version and message size limit.
135    pub const fn with_max_message_size(version: EthVersion, max_message_size: usize) -> Self {
136        Self {
137            version,
138            max_message_size,
139            reject_block_announcements: false,
140            pending_responses: BTreeMap::new(),
141            _pd: std::marker::PhantomData,
142        }
143    }
144
145    /// Returns the eth version
146    #[inline]
147    pub const fn version(&self) -> EthVersion {
148        self.version
149    }
150
151    /// Sets whether to reject block announcement messages (`NewBlock`, `NewBlockHashes`) before
152    /// RLP decoding.
153    pub const fn set_reject_block_announcements(&mut self, reject: bool) {
154        self.reject_block_announcements = reject;
155    }
156
157    /// Decodes incoming bytes into an [`EthMessage`].
158    ///
159    /// A response is only decoded if it answers a request sent through [`Self::encode_message`]
160    /// and carries at most as many items as that request asked for, see
161    /// [`EthStreamError::UnsolicitedResponse`] and [`EthStreamError::ResponseExceedsRequest`].
162    pub fn decode_message(&mut self, bytes: BytesMut) -> Result<EthMessage<N>, EthStreamError> {
163        if bytes.len() > self.max_message_size {
164            return Err(EthStreamError::MessageTooBig(bytes.len()));
165        }
166
167        if self.reject_block_announcements &&
168            let Some(&id) = bytes.first() &&
169            (id == EthMessageID::NewBlock.to_u8() || id == EthMessageID::NewBlockHashes.to_u8())
170        {
171            return Err(EthStreamError::UnsupportedMessage { message_id: id });
172        }
173
174        self.take_pending_response(&bytes)?;
175
176        let msg = match ProtocolMessage::decode_message_with_tx_memory_budget(
177            self.version,
178            &mut bytes.as_ref(),
179            self.max_message_size * TX_MEMORY_BUDGET_MULTIPLIER,
180        ) {
181            Ok(m) => m,
182            Err(err) => {
183                let msg = if bytes.len() > 50 {
184                    format!("{:02x?}...{:x?}", &bytes[..10], &bytes[bytes.len() - 10..])
185                } else {
186                    format!("{bytes:02x?}")
187                };
188                debug!(
189                    version=?self.version,
190                    %msg,
191                    "failed to decode protocol message"
192                );
193                return Err(EthStreamError::InvalidMessage(err));
194            }
195        };
196
197        if matches!(msg.message, EthMessage::Status(_)) {
198            return Err(EthStreamError::EthHandshakeError(EthHandshakeError::StatusNotInHandshake));
199        }
200
201        Ok(msg.message)
202    }
203
204    /// Encodes an [`EthMessage`] to bytes.
205    ///
206    /// Validates that Status messages are not sent after handshake, enforcing protocol rules.
207    /// Requests are recorded so that [`Self::decode_message`] accepts their responses.
208    pub fn encode_message(&mut self, item: EthMessage<N>) -> Result<Bytes, EthStreamError> {
209        if matches!(item, EthMessage::Status(_)) {
210            return Err(EthStreamError::EthHandshakeError(EthHandshakeError::StatusNotInHandshake));
211        }
212
213        if let Some((request_id, response)) = PendingResponse::for_request(&item) {
214            self.pending_responses.insert(request_id, response);
215        }
216
217        Ok(Bytes::from(alloy_rlp::encode(ProtocolMessage::from(item))))
218    }
219
220    /// Consumes the pending request an encoded response answers, without decoding the response's
221    /// items.
222    ///
223    /// Decoding allocates per item, and an item can be a single RLP byte, so a few hundred KiB of
224    /// compressed frame can decode into hundreds of MiB. Responses nobody asked for, and responses
225    /// with more items than requested, are rejected before any of that happens. Messages that are
226    /// not responses pass through.
227    fn take_pending_response(&mut self, bytes: &[u8]) -> Result<(), EthStreamError> {
228        let mut buf = bytes;
229        let message_id = EthMessageID::decode(&mut buf).map_err(MessageError::from)?;
230        if !matches!(
231            message_id,
232            EthMessageID::BlockHeaders |
233                EthMessageID::BlockBodies |
234                EthMessageID::PooledTransactions |
235                EthMessageID::NodeData |
236                EthMessageID::Receipts |
237                EthMessageID::BlockAccessLists |
238                EthMessageID::Cells
239        ) {
240            return Ok(())
241        }
242
243        // Every response is a list that starts with the request id.
244        let mut payload = Header::decode_bytes(&mut buf, true).map_err(MessageError::from)?;
245        let request_id = u64::decode(&mut payload).map_err(MessageError::from)?;
246
247        let Some(pending) = self
248            .pending_responses
249            .remove(&request_id)
250            .filter(|pending| pending.message_id == message_id)
251        else {
252            return Err(EthStreamError::UnsolicitedResponse { message_id, request_id })
253        };
254
255        if response_exceeds_items(self.version, message_id, payload, pending.max_items)
256            .map_err(MessageError::from)?
257        {
258            return Err(EthStreamError::ResponseExceedsRequest {
259                message_id,
260                request_id,
261                max_items: pending.max_items,
262            })
263        }
264
265        Ok(())
266    }
267}
268
269/// An `EthStream` wraps over any `Stream` that yields bytes and makes it
270/// compatible with eth-networking protocol messages, which get RLP encoded/decoded.
271#[pin_project]
272#[derive(Debug)]
273pub struct EthStream<S, N = EthNetworkPrimitives> {
274    /// Eth-specific logic
275    eth: EthStreamInner<N>,
276    #[pin]
277    inner: S,
278}
279
280impl<S, N: NetworkPrimitives> EthStream<S, N> {
281    /// Creates a new unauthed [`EthStream`] from a provided stream. You will need
282    /// to manually handshake a peer.
283    #[inline]
284    pub const fn new(version: EthVersion, inner: S) -> Self {
285        Self::with_max_message_size(version, inner, MAX_MESSAGE_SIZE)
286    }
287
288    /// Creates a new unauthed [`EthStream`] with a custom max message size.
289    #[inline]
290    pub const fn with_max_message_size(
291        version: EthVersion,
292        inner: S,
293        max_message_size: usize,
294    ) -> Self {
295        Self { eth: EthStreamInner::with_max_message_size(version, max_message_size), inner }
296    }
297
298    /// Returns the eth version.
299    #[inline]
300    pub const fn version(&self) -> EthVersion {
301        self.eth.version()
302    }
303
304    /// Sets whether to reject block announcement messages (`NewBlock`, `NewBlockHashes`) before
305    /// RLP decoding.
306    pub const fn set_reject_block_announcements(&mut self, reject: bool) {
307        self.eth.set_reject_block_announcements(reject);
308    }
309
310    /// Returns the underlying stream.
311    #[inline]
312    pub const fn inner(&self) -> &S {
313        &self.inner
314    }
315
316    /// Returns mutable access to the underlying stream.
317    #[inline]
318    pub const fn inner_mut(&mut self) -> &mut S {
319        &mut self.inner
320    }
321
322    /// Consumes this type and returns the wrapped stream.
323    #[inline]
324    pub fn into_inner(self) -> S {
325        self.inner
326    }
327}
328
329impl<S, E, N> EthStream<S, N>
330where
331    S: Sink<Bytes, Error = E> + Unpin,
332    EthStreamError: From<E>,
333    N: NetworkPrimitives,
334{
335    /// Same as [`Sink::start_send`] but accepts a [`EthBroadcastMessage`] instead.
336    pub fn start_send_broadcast(
337        &mut self,
338        item: EthBroadcastMessage<N>,
339    ) -> Result<(), EthStreamError> {
340        self.inner.start_send_unpin(item.encoded())?;
341        Ok(())
342    }
343
344    /// Sends a raw capability message directly over the stream
345    ///
346    /// Raw messages are not recorded as requests, so the response to an `eth` request sent this
347    /// way is rejected as unsolicited. Send requests as [`EthMessage`]s instead.
348    pub fn start_send_raw(&mut self, msg: RawCapabilityMessage) -> Result<(), EthStreamError> {
349        self.inner.start_send_unpin(msg.encoded())?;
350        Ok(())
351    }
352}
353
354impl<S, E, N> Stream for EthStream<S, N>
355where
356    S: Stream<Item = Result<BytesMut, E>> + Unpin,
357    EthStreamError: From<E>,
358    N: NetworkPrimitives,
359{
360    type Item = Result<EthMessage<N>, EthStreamError>;
361
362    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
363        let this = self.project();
364        let res = ready!(this.inner.poll_next(cx));
365
366        match res {
367            Some(Ok(bytes)) => Poll::Ready(Some(this.eth.decode_message(bytes))),
368            Some(Err(err)) => Poll::Ready(Some(Err(err.into()))),
369            None => Poll::Ready(None),
370        }
371    }
372}
373
374impl<S, N> Sink<EthMessage<N>> for EthStream<S, N>
375where
376    S: CanDisconnect<Bytes>,
377    EthStreamError: From<<S as Sink<Bytes>>::Error>,
378    N: NetworkPrimitives,
379{
380    type Error = EthStreamError;
381
382    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
383        self.project().inner.poll_ready(cx).map_err(Into::into)
384    }
385
386    fn start_send(self: Pin<&mut Self>, item: EthMessage<N>) -> Result<(), Self::Error> {
387        if matches!(item, EthMessage::Status(_)) {
388            // Attempt to disconnect the peer for protocol breach when trying to send Status
389            // message after handshake is complete
390            let mut this = self.project();
391            // We can't await the disconnect future here since this is a synchronous method,
392            // but we can start the disconnect process. The actual disconnect will be handled
393            // asynchronously by the caller or the stream's poll methods.
394            let _disconnect_future = this.inner.disconnect(DisconnectReason::ProtocolBreach);
395            return Err(EthStreamError::EthHandshakeError(EthHandshakeError::StatusNotInHandshake))
396        }
397
398        let this = self.project();
399        let bytes = this.eth.encode_message(item)?;
400        this.inner.start_send(bytes)?;
401
402        Ok(())
403    }
404
405    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
406        self.project().inner.poll_flush(cx).map_err(Into::into)
407    }
408
409    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
410        self.project().inner.poll_close(cx).map_err(Into::into)
411    }
412}
413
414impl<S, N> CanDisconnect<EthMessage<N>> for EthStream<S, N>
415where
416    S: CanDisconnect<Bytes> + Send,
417    EthStreamError: From<<S as Sink<Bytes>>::Error>,
418    N: NetworkPrimitives,
419{
420    fn disconnect(
421        &mut self,
422        reason: DisconnectReason,
423    ) -> Pin<Box<dyn Future<Output = Result<(), EthStreamError>> + Send + '_>> {
424        Box::pin(async move { self.inner.disconnect(reason).await.map_err(Into::into) })
425    }
426}
427
428/// The response the remote owes for a request sent on an [`EthStreamInner`].
429#[derive(Clone, Copy, Debug)]
430struct PendingResponse {
431    /// Message id the response must carry.
432    message_id: EthMessageID,
433    /// Number of items the request asked for, which the response must not exceed.
434    max_items: usize,
435}
436
437impl PendingResponse {
438    /// Returns the request id and the response owed for `msg`, or `None` if it is not a request.
439    fn for_request<N: NetworkPrimitives>(msg: &EthMessage<N>) -> Option<(u64, Self)> {
440        let (request_id, message_id, max_items) = match msg {
441            EthMessage::GetBlockHeaders(req) => (
442                req.request_id,
443                EthMessageID::BlockHeaders,
444                usize::try_from(req.message.limit).unwrap_or(usize::MAX),
445            ),
446            EthMessage::GetBlockBodies(req) => {
447                (req.request_id, EthMessageID::BlockBodies, req.message.0.len())
448            }
449            EthMessage::GetPooledTransactions(req) => {
450                (req.request_id, EthMessageID::PooledTransactions, req.message.0.len())
451            }
452            EthMessage::GetNodeData(req) => {
453                (req.request_id, EthMessageID::NodeData, req.message.0.len())
454            }
455            EthMessage::GetReceipts(req) => {
456                (req.request_id, EthMessageID::Receipts, req.message.0.len())
457            }
458            EthMessage::GetReceipts70(req) => {
459                (req.request_id, EthMessageID::Receipts, req.message.block_hashes.len())
460            }
461            EthMessage::GetBlockAccessLists(req) => {
462                (req.request_id, EthMessageID::BlockAccessLists, req.message.0.len())
463            }
464            EthMessage::GetCells(req) => {
465                (req.request_id, EthMessageID::Cells, req.message.hashes.len())
466            }
467            _ => return None,
468        };
469        Some((request_id, Self { message_id, max_items }))
470    }
471}
472
473/// Returns whether the response payload following the request id holds more than `max_items`
474/// items.
475fn response_exceeds_items(
476    version: EthVersion,
477    message_id: EthMessageID,
478    mut payload: &[u8],
479    max_items: usize,
480) -> alloy_rlp::Result<bool> {
481    match message_id {
482        // eth/70 receipts carry the `lastBlockIncomplete` flag before the receipt lists.
483        EthMessageID::Receipts if version >= EthVersion::Eth70 => {
484            bool::decode(&mut payload)?;
485        }
486        // Cells wraps the per-transaction hashes and cell lists in a struct.
487        EthMessageID::Cells => {
488            let mut cells = Header::decode_bytes(&mut payload, true)?;
489            return Ok(list_exceeds_items(&mut cells, max_items)? ||
490                list_exceeds_items(&mut cells, max_items)?)
491        }
492        _ => {}
493    }
494    list_exceeds_items(&mut payload, max_items)
495}
496
497/// Decodes the header of the RLP list at the front of `buf` and returns whether the list holds
498/// more than `max_items` items, without decoding them.
499///
500/// Advances `buf` past the list.
501pub(crate) fn list_exceeds_items(buf: &mut &[u8], max_items: usize) -> alloy_rlp::Result<bool> {
502    let mut items = Header::decode_bytes(buf, true)?;
503    let mut count = 0;
504    while !items.is_empty() {
505        if count == max_items {
506            return Ok(true)
507        }
508        let header = Header::decode(&mut items)?;
509        // `Header::decode` checked that the payload is in bounds.
510        items = &items[header.payload_length..];
511        count += 1;
512    }
513    Ok(false)
514}
515
516#[cfg(test)]
517mod tests {
518    use super::{EthStreamInner, UnauthedEthStream};
519    use crate::{
520        broadcast::BlockHashNumber,
521        errors::{EthHandshakeError, EthStreamError},
522        ethstream::RawCapabilityMessage,
523        hello::DEFAULT_TCP_PORT,
524        message::RequestPair,
525        p2pstream::UnauthedP2PStream,
526        BlockAccessLists, BlockBodies, BlockHeaders, Cells, EthMessage, EthMessageID, EthStream,
527        EthVersion, GetBlockAccessLists, GetBlockBodies, GetCells, GetReceipts70,
528        HelloMessageWithProtocols, PassthroughCodec, ProtocolMessage, ProtocolVersion, Receipts70,
529        Status, StatusMessage,
530    };
531    use alloy_chains::NamedChain;
532    use alloy_primitives::{
533        bytes::{Bytes, BytesMut},
534        B128, B256, U256,
535    };
536    use alloy_rlp::Decodable;
537    use futures::{SinkExt, StreamExt};
538    use reth_ecies::stream::ECIESStream;
539    use reth_eth_wire_types::{EthNetworkPrimitives, UnifiedStatus};
540    use reth_ethereum_forks::{ForkFilter, Head};
541    use reth_network_peers::pk2id;
542    use secp256k1::SecretKey;
543    use std::time::Duration;
544    use tokio::net::{TcpListener, TcpStream};
545    use tokio_util::codec::Decoder;
546
547    #[tokio::test]
548    async fn can_handshake() {
549        let genesis = B256::random();
550        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
551
552        let status = Status {
553            version: EthVersion::Eth67,
554            chain: NamedChain::Mainnet.into(),
555            total_difficulty: U256::ZERO,
556            blockhash: B256::random(),
557            genesis,
558            // Pass the current fork id.
559            forkid: fork_filter.current(),
560        };
561        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
562
563        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
564        let local_addr = listener.local_addr().unwrap();
565
566        let status_clone = unified_status;
567        let fork_filter_clone = fork_filter.clone();
568        let handle = tokio::spawn(async move {
569            // roughly based off of the design of tokio::net::TcpListener
570            let (incoming, _) = listener.accept().await.unwrap();
571            let stream = PassthroughCodec::default().framed(incoming);
572            let (_, their_status) = UnauthedEthStream::new(stream)
573                .handshake::<EthNetworkPrimitives>(status_clone, fork_filter_clone)
574                .await
575                .unwrap();
576
577            // just make sure it equals our status (our status is a clone of their status)
578            assert_eq!(their_status, status_clone);
579        });
580
581        let outgoing = TcpStream::connect(local_addr).await.unwrap();
582        let sink = PassthroughCodec::default().framed(outgoing);
583
584        // try to connect
585        let (_, their_status) = UnauthedEthStream::new(sink)
586            .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
587            .await
588            .unwrap();
589
590        // their status is a clone of our status, these should be equal
591        assert_eq!(their_status, unified_status);
592
593        // wait for it to finish
594        handle.await.unwrap();
595    }
596
597    #[tokio::test]
598    async fn pass_handshake_on_low_td_bitlen() {
599        let genesis = B256::random();
600        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
601
602        let status = Status {
603            version: EthVersion::Eth67,
604            chain: NamedChain::Mainnet.into(),
605            total_difficulty: U256::from(2).pow(U256::from(100)) - U256::ONE,
606            blockhash: B256::random(),
607            genesis,
608            // Pass the current fork id.
609            forkid: fork_filter.current(),
610        };
611        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
612
613        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
614        let local_addr = listener.local_addr().unwrap();
615
616        let status_clone = unified_status;
617        let fork_filter_clone = fork_filter.clone();
618        let handle = tokio::spawn(async move {
619            // roughly based off of the design of tokio::net::TcpListener
620            let (incoming, _) = listener.accept().await.unwrap();
621            let stream = PassthroughCodec::default().framed(incoming);
622            let (_, their_status) = UnauthedEthStream::new(stream)
623                .handshake::<EthNetworkPrimitives>(status_clone, fork_filter_clone)
624                .await
625                .unwrap();
626
627            // just make sure it equals our status, and that the handshake succeeded
628            assert_eq!(their_status, status_clone);
629        });
630
631        let outgoing = TcpStream::connect(local_addr).await.unwrap();
632        let sink = PassthroughCodec::default().framed(outgoing);
633
634        // try to connect
635        let (_, their_status) = UnauthedEthStream::new(sink)
636            .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
637            .await
638            .unwrap();
639
640        // their status is a clone of our status, these should be equal
641        assert_eq!(their_status, unified_status);
642
643        // await the other handshake
644        handle.await.unwrap();
645    }
646
647    #[tokio::test]
648    async fn fail_handshake_on_high_td_bitlen() {
649        let genesis = B256::random();
650        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
651
652        let status = Status {
653            version: EthVersion::Eth67,
654            chain: NamedChain::Mainnet.into(),
655            total_difficulty: U256::from(2).pow(U256::from(164)),
656            blockhash: B256::random(),
657            genesis,
658            // Pass the current fork id.
659            forkid: fork_filter.current(),
660        };
661        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
662
663        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
664        let local_addr = listener.local_addr().unwrap();
665
666        let status_clone = unified_status;
667        let fork_filter_clone = fork_filter.clone();
668        let handle = tokio::spawn(async move {
669            // roughly based off of the design of tokio::net::TcpListener
670            let (incoming, _) = listener.accept().await.unwrap();
671            let stream = PassthroughCodec::default().framed(incoming);
672            let handshake_res = UnauthedEthStream::new(stream)
673                .handshake::<EthNetworkPrimitives>(status_clone, fork_filter_clone)
674                .await;
675
676            // make sure the handshake fails due to td too high
677            assert!(matches!(
678                handshake_res,
679                Err(EthStreamError::EthHandshakeError(
680                    EthHandshakeError::TotalDifficultyBitLenTooLarge { got: 165, maximum: 160 }
681                ))
682            ));
683        });
684
685        let outgoing = TcpStream::connect(local_addr).await.unwrap();
686        let sink = PassthroughCodec::default().framed(outgoing);
687
688        // try to connect
689        let handshake_res = UnauthedEthStream::new(sink)
690            .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
691            .await;
692
693        // this handshake should also fail due to td too high
694        assert!(matches!(
695            handshake_res,
696            Err(EthStreamError::EthHandshakeError(
697                EthHandshakeError::TotalDifficultyBitLenTooLarge { got: 165, maximum: 160 }
698            ))
699        ));
700
701        // await the other handshake
702        handle.await.unwrap();
703    }
704
705    #[tokio::test]
706    async fn can_write_and_read_cleartext() {
707        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
708        let local_addr = listener.local_addr().unwrap();
709        let test_msg = EthMessage::<EthNetworkPrimitives>::NewBlockHashes(
710            vec![
711                BlockHashNumber { hash: B256::random(), number: 5 },
712                BlockHashNumber { hash: B256::random(), number: 6 },
713            ]
714            .into(),
715        );
716
717        let test_msg_clone = test_msg.clone();
718        let handle = tokio::spawn(async move {
719            // roughly based off of the design of tokio::net::TcpListener
720            let (incoming, _) = listener.accept().await.unwrap();
721            let stream = PassthroughCodec::default().framed(incoming);
722            let mut stream = EthStream::new(EthVersion::Eth67, stream);
723
724            // use the stream to get the next message
725            let message = stream.next().await.unwrap().unwrap();
726            assert_eq!(message, test_msg_clone);
727        });
728
729        let outgoing = TcpStream::connect(local_addr).await.unwrap();
730        let sink = PassthroughCodec::default().framed(outgoing);
731        let mut client_stream = EthStream::new(EthVersion::Eth67, sink);
732
733        client_stream.send(test_msg).await.unwrap();
734
735        // make sure the server receives the message and asserts before ending the test
736        handle.await.unwrap();
737    }
738
739    #[tokio::test]
740    async fn can_write_and_read_ecies() {
741        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
742        let local_addr = listener.local_addr().unwrap();
743        let server_key = SecretKey::new(&mut rand::rng());
744        let test_msg = EthMessage::<EthNetworkPrimitives>::NewBlockHashes(
745            vec![
746                BlockHashNumber { hash: B256::random(), number: 5 },
747                BlockHashNumber { hash: B256::random(), number: 6 },
748            ]
749            .into(),
750        );
751
752        let test_msg_clone = test_msg.clone();
753        let handle = tokio::spawn(async move {
754            // roughly based off of the design of tokio::net::TcpListener
755            let (incoming, _) = listener.accept().await.unwrap();
756            let stream = ECIESStream::incoming(incoming, server_key).await.unwrap();
757            let mut stream = EthStream::new(EthVersion::Eth67, stream);
758
759            // use the stream to get the next message
760            let message = stream.next().await.unwrap().unwrap();
761            assert_eq!(message, test_msg_clone);
762        });
763
764        // create the server pubkey
765        let server_id = pk2id(&server_key.public_key());
766
767        let client_key = SecretKey::new(&mut rand::rng());
768
769        let outgoing = TcpStream::connect(local_addr).await.unwrap();
770        let outgoing = ECIESStream::connect(outgoing, client_key, server_id).await.unwrap();
771        let mut client_stream = EthStream::new(EthVersion::Eth67, outgoing);
772
773        client_stream.send(test_msg).await.unwrap();
774
775        // make sure the server receives the message and asserts before ending the test
776        handle.await.unwrap();
777    }
778
779    #[tokio::test(flavor = "multi_thread")]
780    async fn ethstream_over_p2p() {
781        // create a p2p stream and server, then confirm that the two are authed
782        // create tcpstream
783        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
784        let local_addr = listener.local_addr().unwrap();
785        let server_key = SecretKey::new(&mut rand::rng());
786        let test_msg = EthMessage::<EthNetworkPrimitives>::NewBlockHashes(
787            vec![
788                BlockHashNumber { hash: B256::random(), number: 5 },
789                BlockHashNumber { hash: B256::random(), number: 6 },
790            ]
791            .into(),
792        );
793
794        let genesis = B256::random();
795        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
796
797        let status = Status {
798            version: EthVersion::Eth67,
799            chain: NamedChain::Mainnet.into(),
800            total_difficulty: U256::ZERO,
801            blockhash: B256::random(),
802            genesis,
803            // Pass the current fork id.
804            forkid: fork_filter.current(),
805        };
806        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
807
808        let status_copy = unified_status;
809        let fork_filter_clone = fork_filter.clone();
810        let test_msg_clone = test_msg.clone();
811        let handle = tokio::spawn(async move {
812            // roughly based off of the design of tokio::net::TcpListener
813            let (incoming, _) = listener.accept().await.unwrap();
814            let stream = ECIESStream::incoming(incoming, server_key).await.unwrap();
815
816            let server_hello = HelloMessageWithProtocols {
817                protocol_version: ProtocolVersion::V5,
818                client_version: "bitcoind/1.0.0".to_string(),
819                protocols: vec![EthVersion::Eth67.into()],
820                port: DEFAULT_TCP_PORT,
821                id: pk2id(&server_key.public_key()),
822            };
823
824            let unauthed_stream = UnauthedP2PStream::new(stream);
825            let (p2p_stream, _) = unauthed_stream.handshake(server_hello).await.unwrap();
826            let (mut eth_stream, _) = UnauthedEthStream::new(p2p_stream)
827                .handshake(status_copy, fork_filter_clone)
828                .await
829                .unwrap();
830
831            // use the stream to get the next message
832            let message = eth_stream.next().await.unwrap().unwrap();
833            assert_eq!(message, test_msg_clone);
834        });
835
836        // create the server pubkey
837        let server_id = pk2id(&server_key.public_key());
838
839        let client_key = SecretKey::new(&mut rand::rng());
840
841        let outgoing = TcpStream::connect(local_addr).await.unwrap();
842        let sink = ECIESStream::connect(outgoing, client_key, server_id).await.unwrap();
843
844        let client_hello = HelloMessageWithProtocols {
845            protocol_version: ProtocolVersion::V5,
846            client_version: "bitcoind/1.0.0".to_string(),
847            protocols: vec![EthVersion::Eth67.into()],
848            port: DEFAULT_TCP_PORT,
849            id: pk2id(&client_key.public_key()),
850        };
851
852        let unauthed_stream = UnauthedP2PStream::new(sink);
853        let (p2p_stream, _) = unauthed_stream.handshake(client_hello).await.unwrap();
854
855        let (mut client_stream, _) = UnauthedEthStream::new(p2p_stream)
856            .handshake(unified_status, fork_filter)
857            .await
858            .unwrap();
859
860        client_stream.send(test_msg).await.unwrap();
861
862        // make sure the server receives the message and asserts before ending the test
863        handle.await.unwrap();
864    }
865
866    #[tokio::test]
867    async fn handshake_should_timeout() {
868        let genesis = B256::random();
869        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
870
871        let status = Status {
872            version: EthVersion::Eth67,
873            chain: NamedChain::Mainnet.into(),
874            total_difficulty: U256::ZERO,
875            blockhash: B256::random(),
876            genesis,
877            // Pass the current fork id.
878            forkid: fork_filter.current(),
879        };
880        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
881
882        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
883        let local_addr = listener.local_addr().unwrap();
884
885        let status_clone = unified_status;
886        let fork_filter_clone = fork_filter.clone();
887        let _handle = tokio::spawn(async move {
888            // Delay accepting the connection for longer than the client's timeout period
889            tokio::time::sleep(Duration::from_secs(11)).await;
890            // roughly based off of the design of tokio::net::TcpListener
891            let (incoming, _) = listener.accept().await.unwrap();
892            let stream = PassthroughCodec::default().framed(incoming);
893            let (_, their_status) = UnauthedEthStream::new(stream)
894                .handshake::<EthNetworkPrimitives>(status_clone, fork_filter_clone)
895                .await
896                .unwrap();
897
898            // just make sure it equals our status (our status is a clone of their status)
899            assert_eq!(their_status, status_clone);
900        });
901
902        let outgoing = TcpStream::connect(local_addr).await.unwrap();
903        let sink = PassthroughCodec::default().framed(outgoing);
904
905        // try to connect
906        let handshake_result = UnauthedEthStream::new(sink)
907            .handshake_with_timeout::<EthNetworkPrimitives>(
908                unified_status,
909                fork_filter,
910                Duration::from_secs(1),
911            )
912            .await;
913
914        // Assert that a timeout error occurred
915        assert!(
916            matches!(handshake_result, Err(e) if e.to_string() == EthStreamError::StreamTimeout.to_string())
917        );
918    }
919
920    #[tokio::test]
921    async fn can_write_and_read_raw_capability() {
922        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
923        let local_addr = listener.local_addr().unwrap();
924
925        let test_msg = RawCapabilityMessage { id: 0x1234, payload: Bytes::from(vec![1, 2, 3, 4]) };
926
927        let test_msg_clone = test_msg.clone();
928        let handle = tokio::spawn(async move {
929            let (incoming, _) = listener.accept().await.unwrap();
930            let stream = PassthroughCodec::default().framed(incoming);
931            let mut stream = EthStream::<_, EthNetworkPrimitives>::new(EthVersion::Eth67, stream);
932
933            let bytes = stream.inner_mut().next().await.unwrap().unwrap();
934
935            // Create a cursor to track position while decoding
936            let mut id_bytes = &bytes[..];
937            let decoded_id = <usize as Decodable>::decode(&mut id_bytes).unwrap();
938            assert_eq!(decoded_id, test_msg_clone.id);
939
940            // Get remaining bytes after ID decoding
941            let remaining = id_bytes;
942            assert_eq!(remaining, &test_msg_clone.payload[..]);
943        });
944
945        let outgoing = TcpStream::connect(local_addr).await.unwrap();
946        let sink = PassthroughCodec::default().framed(outgoing);
947        let mut client_stream = EthStream::<_, EthNetworkPrimitives>::new(EthVersion::Eth67, sink);
948
949        client_stream.start_send_raw(test_msg).unwrap();
950        client_stream.inner_mut().flush().await.unwrap();
951
952        handle.await.unwrap();
953    }
954
955    #[tokio::test]
956    async fn status_message_after_handshake_triggers_disconnect() {
957        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
958        let local_addr = listener.local_addr().unwrap();
959
960        let handle = tokio::spawn(async move {
961            let (incoming, _) = listener.accept().await.unwrap();
962            let stream = PassthroughCodec::default().framed(incoming);
963            let mut stream = EthStream::<_, EthNetworkPrimitives>::new(EthVersion::Eth67, stream);
964
965            // Try to send a Status message after handshake - this should trigger disconnect
966            let status = Status {
967                version: EthVersion::Eth67,
968                chain: NamedChain::Mainnet.into(),
969                total_difficulty: U256::ZERO,
970                blockhash: B256::random(),
971                genesis: B256::random(),
972                forkid: ForkFilter::new(Head::default(), B256::random(), 0, Vec::new()).current(),
973            };
974            let status_message =
975                EthMessage::<EthNetworkPrimitives>::Status(StatusMessage::Legacy(status));
976
977            // This should return an error and trigger disconnect
978            let result = stream.send(status_message).await;
979            assert!(result.is_err());
980            assert!(matches!(
981                result.unwrap_err(),
982                EthStreamError::EthHandshakeError(EthHandshakeError::StatusNotInHandshake)
983            ));
984        });
985
986        let outgoing = TcpStream::connect(local_addr).await.unwrap();
987        let sink = PassthroughCodec::default().framed(outgoing);
988        let mut client_stream = EthStream::<_, EthNetworkPrimitives>::new(EthVersion::Eth67, sink);
989
990        // Send a valid message to keep the connection alive
991        let test_msg = EthMessage::<EthNetworkPrimitives>::NewBlockHashes(
992            vec![BlockHashNumber { hash: B256::random(), number: 5 }].into(),
993        );
994        client_stream.send(test_msg).await.unwrap();
995
996        handle.await.unwrap();
997    }
998
999    /// Encodes `msg` with its message id, as [`EthStreamInner::decode_message`] receives it.
1000    fn encode(msg: EthMessage<EthNetworkPrimitives>) -> BytesMut {
1001        alloy_rlp::encode(ProtocolMessage::from(msg)).as_slice().into()
1002    }
1003
1004    fn get_block_access_lists(request_id: u64, hashes: usize) -> EthMessage<EthNetworkPrimitives> {
1005        EthMessage::GetBlockAccessLists(RequestPair {
1006            request_id,
1007            message: GetBlockAccessLists(vec![B256::ZERO; hashes]),
1008        })
1009    }
1010
1011    fn block_access_lists(request_id: u64, entries: usize) -> EthMessage<EthNetworkPrimitives> {
1012        EthMessage::BlockAccessLists(RequestPair {
1013            request_id,
1014            message: BlockAccessLists(vec![None; entries]),
1015        })
1016    }
1017
1018    #[test]
1019    fn unsolicited_response_is_rejected() {
1020        let mut eth = EthStreamInner::<EthNetworkPrimitives>::new(EthVersion::Eth71);
1021
1022        let err = eth.decode_message(encode(block_access_lists(7, 1))).unwrap_err();
1023        assert!(matches!(
1024            err,
1025            EthStreamError::UnsolicitedResponse {
1026                message_id: EthMessageID::BlockAccessLists,
1027                request_id: 7
1028            }
1029        ));
1030        assert!(err.is_protocol_breach());
1031    }
1032
1033    #[test]
1034    fn response_to_sent_request_is_decoded_once() {
1035        let mut eth = EthStreamInner::<EthNetworkPrimitives>::new(EthVersion::Eth71);
1036        eth.encode_message(get_block_access_lists(7, 2)).unwrap();
1037
1038        let response = block_access_lists(7, 2);
1039        assert_eq!(eth.decode_message(encode(response.clone())).unwrap(), response);
1040        assert!(matches!(
1041            eth.decode_message(encode(response)).unwrap_err(),
1042            EthStreamError::UnsolicitedResponse {
1043                message_id: EthMessageID::BlockAccessLists,
1044                request_id: 7
1045            }
1046        ));
1047    }
1048
1049    #[test]
1050    fn response_with_more_items_than_requested_is_rejected() {
1051        let mut eth = EthStreamInner::<EthNetworkPrimitives>::new(EthVersion::Eth71);
1052        eth.encode_message(get_block_access_lists(7, 2)).unwrap();
1053
1054        let err = eth.decode_message(encode(block_access_lists(7, 3))).unwrap_err();
1055        assert!(matches!(
1056            err,
1057            EthStreamError::ResponseExceedsRequest {
1058                message_id: EthMessageID::BlockAccessLists,
1059                request_id: 7,
1060                max_items: 2
1061            }
1062        ));
1063        assert!(err.is_protocol_breach());
1064    }
1065
1066    #[test]
1067    fn response_of_other_type_is_rejected() {
1068        let mut eth = EthStreamInner::<EthNetworkPrimitives>::new(EthVersion::Eth71);
1069        eth.encode_message(EthMessage::GetBlockBodies(RequestPair {
1070            request_id: 7,
1071            message: GetBlockBodies(vec![B256::ZERO]),
1072        }))
1073        .unwrap();
1074
1075        let response = EthMessage::BlockHeaders(RequestPair {
1076            request_id: 7,
1077            message: BlockHeaders(Vec::new()),
1078        });
1079        assert!(matches!(
1080            eth.decode_message(encode(response)).unwrap_err(),
1081            EthStreamError::UnsolicitedResponse {
1082                message_id: EthMessageID::BlockHeaders,
1083                request_id: 7
1084            }
1085        ));
1086
1087        // The request was consumed, so its proper response is now unsolicited as well.
1088        let response = EthMessage::BlockBodies(RequestPair {
1089            request_id: 7,
1090            message: BlockBodies(Vec::new()),
1091        });
1092        assert!(matches!(
1093            eth.decode_message(encode(response)).unwrap_err(),
1094            EthStreamError::UnsolicitedResponse { message_id: EthMessageID::BlockBodies, .. }
1095        ));
1096    }
1097
1098    #[test]
1099    fn receipts70_response_items_are_receipt_lists() {
1100        let mut eth = EthStreamInner::<EthNetworkPrimitives>::new(EthVersion::Eth70);
1101        let request = |request_id| {
1102            EthMessage::GetReceipts70(RequestPair {
1103                request_id,
1104                message: GetReceipts70 {
1105                    first_block_receipt_index: 0,
1106                    block_hashes: vec![B256::ZERO],
1107                },
1108            })
1109        };
1110        let response = |request_id, lists| {
1111            EthMessage::Receipts70(RequestPair {
1112                request_id,
1113                message: Receipts70 {
1114                    last_block_incomplete: true,
1115                    receipts: vec![Vec::new(); lists],
1116                },
1117            })
1118        };
1119
1120        eth.encode_message(request(1)).unwrap();
1121        assert_eq!(eth.decode_message(encode(response(1, 1))).unwrap(), response(1, 1));
1122
1123        eth.encode_message(request(2)).unwrap();
1124        assert!(matches!(
1125            eth.decode_message(encode(response(2, 2))).unwrap_err(),
1126            EthStreamError::ResponseExceedsRequest {
1127                message_id: EthMessageID::Receipts,
1128                request_id: 2,
1129                max_items: 1
1130            }
1131        ));
1132    }
1133
1134    #[test]
1135    fn cells_response_bounds_hashes_and_cells() {
1136        let mut eth = EthStreamInner::<EthNetworkPrimitives>::new(EthVersion::Eth72);
1137        let request = |request_id| {
1138            EthMessage::GetCells(RequestPair {
1139                request_id,
1140                message: GetCells { hashes: vec![B256::ZERO], cell_mask: B128::ZERO },
1141            })
1142        };
1143        let response = |request_id, hashes, cells| {
1144            EthMessage::Cells(RequestPair {
1145                request_id,
1146                message: Cells {
1147                    hashes: vec![B256::ZERO; hashes],
1148                    cells: vec![Vec::new(); cells],
1149                    cell_mask: B128::ZERO,
1150                },
1151            })
1152        };
1153
1154        eth.encode_message(request(1)).unwrap();
1155        assert_eq!(eth.decode_message(encode(response(1, 1, 1))).unwrap(), response(1, 1, 1));
1156
1157        for (request_id, hashes, cells) in [(2, 2, 1), (3, 1, 2)] {
1158            eth.encode_message(request(request_id)).unwrap();
1159            assert!(matches!(
1160                eth.decode_message(encode(response(request_id, hashes, cells))).unwrap_err(),
1161                EthStreamError::ResponseExceedsRequest {
1162                    message_id: EthMessageID::Cells,
1163                    max_items: 1,
1164                    ..
1165                }
1166            ));
1167        }
1168    }
1169}