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, 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 futures::{ready, Sink, SinkExt};
17use pin_project::pin_project;
18use reth_eth_wire_types::{EthMessageID, NetworkPrimitives, RawCapabilityMessage};
19use reth_ethereum_forks::ForkFilter;
20use std::{
21    future::Future,
22    pin::Pin,
23    task::{Context, Poll},
24    time::Duration,
25};
26use tokio::time::timeout;
27use tokio_stream::Stream;
28use tracing::{debug, trace};
29
30/// An un-authenticated [`EthStream`]. This is consumed and returns a [`EthStream`] after the
31/// `Status` handshake is completed.
32#[pin_project]
33#[derive(Debug)]
34pub struct UnauthedEthStream<S> {
35    #[pin]
36    inner: S,
37}
38
39impl<S> UnauthedEthStream<S> {
40    /// Create a new `UnauthedEthStream` from a type `S` which implements `Stream` and `Sink`.
41    pub const fn new(inner: S) -> Self {
42        Self { inner }
43    }
44
45    /// Consumes the type and returns the wrapped stream
46    pub fn into_inner(self) -> S {
47        self.inner
48    }
49}
50
51impl<S, E> UnauthedEthStream<S>
52where
53    S: Stream<Item = Result<BytesMut, E>> + CanDisconnect<Bytes> + Send + Unpin,
54    EthStreamError: From<E> + From<<S as Sink<Bytes>>::Error>,
55{
56    /// Consumes the [`UnauthedEthStream`] and returns an [`EthStream`] after the `Status`
57    /// handshake is completed successfully. This also returns the `Status` message sent by the
58    /// remote peer.
59    ///
60    /// Caution: This expects that the [`UnifiedStatus`] has the proper eth version configured, with
61    /// ETH69 the initial status message changed.
62    pub async fn handshake<N: NetworkPrimitives>(
63        self,
64        status: UnifiedStatus,
65        fork_filter: ForkFilter,
66    ) -> Result<(EthStream<S, N>, UnifiedStatus), EthStreamError> {
67        self.handshake_with_timeout(status, fork_filter, HANDSHAKE_TIMEOUT).await
68    }
69
70    /// Wrapper around handshake which enforces a timeout.
71    pub async fn handshake_with_timeout<N: NetworkPrimitives>(
72        self,
73        status: UnifiedStatus,
74        fork_filter: ForkFilter,
75        timeout_limit: Duration,
76    ) -> Result<(EthStream<S, N>, UnifiedStatus), EthStreamError> {
77        timeout(timeout_limit, Self::handshake_without_timeout(self, status, fork_filter))
78            .await
79            .map_err(|_| EthStreamError::StreamTimeout)?
80    }
81
82    /// Handshake with no timeout
83    pub async fn handshake_without_timeout<N: NetworkPrimitives>(
84        mut self,
85        status: UnifiedStatus,
86        fork_filter: ForkFilter,
87    ) -> Result<(EthStream<S, N>, UnifiedStatus), EthStreamError> {
88        trace!(
89            status = %status.into_message(),
90            "sending eth status to peer"
91        );
92        let their_status =
93            EthereumEthHandshake(&mut self.inner).eth_handshake(status, fork_filter).await?;
94
95        // now we can create the `EthStream` because the peer has successfully completed
96        // the handshake
97        let stream = EthStream::new(status.version, self.inner);
98
99        Ok((stream, their_status))
100    }
101}
102
103/// Contains eth protocol specific logic for processing messages
104#[derive(Debug)]
105pub struct EthStreamInner<N> {
106    /// Negotiated eth version
107    version: EthVersion,
108    /// Maximum allowed ETH message size.
109    max_message_size: usize,
110    /// When true, `NewBlock` (0x07) and `NewBlockHashes` (0x01) messages are rejected before RLP
111    /// decoding to avoid any memory impact for non-PoW networks.
112    reject_block_announcements: bool,
113    _pd: std::marker::PhantomData<N>,
114}
115
116impl<N> EthStreamInner<N>
117where
118    N: NetworkPrimitives,
119{
120    /// Creates a new [`EthStreamInner`] with the given eth version
121    pub const fn new(version: EthVersion) -> Self {
122        Self::with_max_message_size(version, MAX_MESSAGE_SIZE)
123    }
124
125    /// Creates a new [`EthStreamInner`] with the given eth version and message size limit.
126    pub const fn with_max_message_size(version: EthVersion, max_message_size: usize) -> Self {
127        Self {
128            version,
129            max_message_size,
130            reject_block_announcements: false,
131            _pd: std::marker::PhantomData,
132        }
133    }
134
135    /// Returns the eth version
136    #[inline]
137    pub const fn version(&self) -> EthVersion {
138        self.version
139    }
140
141    /// Sets whether to reject block announcement messages (`NewBlock`, `NewBlockHashes`) before
142    /// RLP decoding.
143    pub const fn set_reject_block_announcements(&mut self, reject: bool) {
144        self.reject_block_announcements = reject;
145    }
146
147    /// Decodes incoming bytes into an [`EthMessage`].
148    pub fn decode_message(&self, bytes: BytesMut) -> Result<EthMessage<N>, EthStreamError> {
149        if bytes.len() > self.max_message_size {
150            return Err(EthStreamError::MessageTooBig(bytes.len()));
151        }
152
153        if self.reject_block_announcements &&
154            let Some(&id) = bytes.first() &&
155            (id == EthMessageID::NewBlock.to_u8() || id == EthMessageID::NewBlockHashes.to_u8())
156        {
157            return Err(EthStreamError::UnsupportedMessage { message_id: id });
158        }
159
160        let msg = match ProtocolMessage::decode_message_with_tx_memory_budget(
161            self.version,
162            &mut bytes.as_ref(),
163            self.max_message_size * TX_MEMORY_BUDGET_MULTIPLIER,
164        ) {
165            Ok(m) => m,
166            Err(err) => {
167                let msg = if bytes.len() > 50 {
168                    format!("{:02x?}...{:x?}", &bytes[..10], &bytes[bytes.len() - 10..])
169                } else {
170                    format!("{bytes:02x?}")
171                };
172                debug!(
173                    version=?self.version,
174                    %msg,
175                    "failed to decode protocol message"
176                );
177                return Err(EthStreamError::InvalidMessage(err));
178            }
179        };
180
181        if matches!(msg.message, EthMessage::Status(_)) {
182            return Err(EthStreamError::EthHandshakeError(EthHandshakeError::StatusNotInHandshake));
183        }
184
185        Ok(msg.message)
186    }
187
188    /// Encodes an [`EthMessage`] to bytes.
189    ///
190    /// Validates that Status messages are not sent after handshake, enforcing protocol rules.
191    pub fn encode_message(&self, item: EthMessage<N>) -> Result<Bytes, EthStreamError> {
192        if matches!(item, EthMessage::Status(_)) {
193            return Err(EthStreamError::EthHandshakeError(EthHandshakeError::StatusNotInHandshake));
194        }
195
196        Ok(Bytes::from(alloy_rlp::encode(ProtocolMessage::from(item))))
197    }
198}
199
200/// An `EthStream` wraps over any `Stream` that yields bytes and makes it
201/// compatible with eth-networking protocol messages, which get RLP encoded/decoded.
202#[pin_project]
203#[derive(Debug)]
204pub struct EthStream<S, N = EthNetworkPrimitives> {
205    /// Eth-specific logic
206    eth: EthStreamInner<N>,
207    #[pin]
208    inner: S,
209}
210
211impl<S, N: NetworkPrimitives> EthStream<S, N> {
212    /// Creates a new unauthed [`EthStream`] from a provided stream. You will need
213    /// to manually handshake a peer.
214    #[inline]
215    pub const fn new(version: EthVersion, inner: S) -> Self {
216        Self::with_max_message_size(version, inner, MAX_MESSAGE_SIZE)
217    }
218
219    /// Creates a new unauthed [`EthStream`] with a custom max message size.
220    #[inline]
221    pub const fn with_max_message_size(
222        version: EthVersion,
223        inner: S,
224        max_message_size: usize,
225    ) -> Self {
226        Self { eth: EthStreamInner::with_max_message_size(version, max_message_size), inner }
227    }
228
229    /// Returns the eth version.
230    #[inline]
231    pub const fn version(&self) -> EthVersion {
232        self.eth.version()
233    }
234
235    /// Sets whether to reject block announcement messages (`NewBlock`, `NewBlockHashes`) before
236    /// RLP decoding.
237    pub const fn set_reject_block_announcements(&mut self, reject: bool) {
238        self.eth.set_reject_block_announcements(reject);
239    }
240
241    /// Returns the underlying stream.
242    #[inline]
243    pub const fn inner(&self) -> &S {
244        &self.inner
245    }
246
247    /// Returns mutable access to the underlying stream.
248    #[inline]
249    pub const fn inner_mut(&mut self) -> &mut S {
250        &mut self.inner
251    }
252
253    /// Consumes this type and returns the wrapped stream.
254    #[inline]
255    pub fn into_inner(self) -> S {
256        self.inner
257    }
258}
259
260impl<S, E, N> EthStream<S, N>
261where
262    S: Sink<Bytes, Error = E> + Unpin,
263    EthStreamError: From<E>,
264    N: NetworkPrimitives,
265{
266    /// Same as [`Sink::start_send`] but accepts a [`EthBroadcastMessage`] instead.
267    pub fn start_send_broadcast(
268        &mut self,
269        item: EthBroadcastMessage<N>,
270    ) -> Result<(), EthStreamError> {
271        self.inner.start_send_unpin(item.encoded())?;
272        Ok(())
273    }
274
275    /// Sends a raw capability message directly over the stream
276    pub fn start_send_raw(&mut self, msg: RawCapabilityMessage) -> Result<(), EthStreamError> {
277        self.inner.start_send_unpin(msg.encoded())?;
278        Ok(())
279    }
280}
281
282impl<S, E, N> Stream for EthStream<S, N>
283where
284    S: Stream<Item = Result<BytesMut, E>> + Unpin,
285    EthStreamError: From<E>,
286    N: NetworkPrimitives,
287{
288    type Item = Result<EthMessage<N>, EthStreamError>;
289
290    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
291        let this = self.project();
292        let res = ready!(this.inner.poll_next(cx));
293
294        match res {
295            Some(Ok(bytes)) => Poll::Ready(Some(this.eth.decode_message(bytes))),
296            Some(Err(err)) => Poll::Ready(Some(Err(err.into()))),
297            None => Poll::Ready(None),
298        }
299    }
300}
301
302impl<S, N> Sink<EthMessage<N>> for EthStream<S, N>
303where
304    S: CanDisconnect<Bytes> + Unpin,
305    EthStreamError: From<<S as Sink<Bytes>>::Error>,
306    N: NetworkPrimitives,
307{
308    type Error = EthStreamError;
309
310    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
311        self.project().inner.poll_ready(cx).map_err(Into::into)
312    }
313
314    fn start_send(self: Pin<&mut Self>, item: EthMessage<N>) -> Result<(), Self::Error> {
315        if matches!(item, EthMessage::Status(_)) {
316            // Attempt to disconnect the peer for protocol breach when trying to send Status
317            // message after handshake is complete
318            let mut this = self.project();
319            // We can't await the disconnect future here since this is a synchronous method,
320            // but we can start the disconnect process. The actual disconnect will be handled
321            // asynchronously by the caller or the stream's poll methods.
322            let _disconnect_future = this.inner.disconnect(DisconnectReason::ProtocolBreach);
323            return Err(EthStreamError::EthHandshakeError(EthHandshakeError::StatusNotInHandshake))
324        }
325
326        self.project()
327            .inner
328            .start_send(Bytes::from(alloy_rlp::encode(ProtocolMessage::from(item))))?;
329
330        Ok(())
331    }
332
333    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
334        self.project().inner.poll_flush(cx).map_err(Into::into)
335    }
336
337    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
338        self.project().inner.poll_close(cx).map_err(Into::into)
339    }
340}
341
342impl<S, N> CanDisconnect<EthMessage<N>> for EthStream<S, N>
343where
344    S: CanDisconnect<Bytes> + Send,
345    EthStreamError: From<<S as Sink<Bytes>>::Error>,
346    N: NetworkPrimitives,
347{
348    fn disconnect(
349        &mut self,
350        reason: DisconnectReason,
351    ) -> Pin<Box<dyn Future<Output = Result<(), EthStreamError>> + Send + '_>> {
352        Box::pin(async move { self.inner.disconnect(reason).await.map_err(Into::into) })
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use super::UnauthedEthStream;
359    use crate::{
360        broadcast::BlockHashNumber,
361        errors::{EthHandshakeError, EthStreamError},
362        ethstream::RawCapabilityMessage,
363        hello::DEFAULT_TCP_PORT,
364        p2pstream::UnauthedP2PStream,
365        EthMessage, EthStream, EthVersion, HelloMessageWithProtocols, PassthroughCodec,
366        ProtocolVersion, Status, StatusMessage,
367    };
368    use alloy_chains::NamedChain;
369    use alloy_primitives::{bytes::Bytes, B256, U256};
370    use alloy_rlp::Decodable;
371    use futures::{SinkExt, StreamExt};
372    use reth_ecies::stream::ECIESStream;
373    use reth_eth_wire_types::{EthNetworkPrimitives, UnifiedStatus};
374    use reth_ethereum_forks::{ForkFilter, Head};
375    use reth_network_peers::pk2id;
376    use secp256k1::{SecretKey, SECP256K1};
377    use std::time::Duration;
378    use tokio::net::{TcpListener, TcpStream};
379    use tokio_util::codec::Decoder;
380
381    #[tokio::test]
382    async fn can_handshake() {
383        let genesis = B256::random();
384        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
385
386        let status = Status {
387            version: EthVersion::Eth67,
388            chain: NamedChain::Mainnet.into(),
389            total_difficulty: U256::ZERO,
390            blockhash: B256::random(),
391            genesis,
392            // Pass the current fork id.
393            forkid: fork_filter.current(),
394        };
395        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
396
397        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
398        let local_addr = listener.local_addr().unwrap();
399
400        let status_clone = unified_status;
401        let fork_filter_clone = fork_filter.clone();
402        let handle = tokio::spawn(async move {
403            // roughly based off of the design of tokio::net::TcpListener
404            let (incoming, _) = listener.accept().await.unwrap();
405            let stream = PassthroughCodec::default().framed(incoming);
406            let (_, their_status) = UnauthedEthStream::new(stream)
407                .handshake::<EthNetworkPrimitives>(status_clone, fork_filter_clone)
408                .await
409                .unwrap();
410
411            // just make sure it equals our status (our status is a clone of their status)
412            assert_eq!(their_status, status_clone);
413        });
414
415        let outgoing = TcpStream::connect(local_addr).await.unwrap();
416        let sink = PassthroughCodec::default().framed(outgoing);
417
418        // try to connect
419        let (_, their_status) = UnauthedEthStream::new(sink)
420            .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
421            .await
422            .unwrap();
423
424        // their status is a clone of our status, these should be equal
425        assert_eq!(their_status, unified_status);
426
427        // wait for it to finish
428        handle.await.unwrap();
429    }
430
431    #[tokio::test]
432    async fn pass_handshake_on_low_td_bitlen() {
433        let genesis = B256::random();
434        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
435
436        let status = Status {
437            version: EthVersion::Eth67,
438            chain: NamedChain::Mainnet.into(),
439            total_difficulty: U256::from(2).pow(U256::from(100)) - U256::from(1),
440            blockhash: B256::random(),
441            genesis,
442            // Pass the current fork id.
443            forkid: fork_filter.current(),
444        };
445        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
446
447        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
448        let local_addr = listener.local_addr().unwrap();
449
450        let status_clone = unified_status;
451        let fork_filter_clone = fork_filter.clone();
452        let handle = tokio::spawn(async move {
453            // roughly based off of the design of tokio::net::TcpListener
454            let (incoming, _) = listener.accept().await.unwrap();
455            let stream = PassthroughCodec::default().framed(incoming);
456            let (_, their_status) = UnauthedEthStream::new(stream)
457                .handshake::<EthNetworkPrimitives>(status_clone, fork_filter_clone)
458                .await
459                .unwrap();
460
461            // just make sure it equals our status, and that the handshake succeeded
462            assert_eq!(their_status, status_clone);
463        });
464
465        let outgoing = TcpStream::connect(local_addr).await.unwrap();
466        let sink = PassthroughCodec::default().framed(outgoing);
467
468        // try to connect
469        let (_, their_status) = UnauthedEthStream::new(sink)
470            .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
471            .await
472            .unwrap();
473
474        // their status is a clone of our status, these should be equal
475        assert_eq!(their_status, unified_status);
476
477        // await the other handshake
478        handle.await.unwrap();
479    }
480
481    #[tokio::test]
482    async fn fail_handshake_on_high_td_bitlen() {
483        let genesis = B256::random();
484        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
485
486        let status = Status {
487            version: EthVersion::Eth67,
488            chain: NamedChain::Mainnet.into(),
489            total_difficulty: U256::from(2).pow(U256::from(164)),
490            blockhash: B256::random(),
491            genesis,
492            // Pass the current fork id.
493            forkid: fork_filter.current(),
494        };
495        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
496
497        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
498        let local_addr = listener.local_addr().unwrap();
499
500        let status_clone = unified_status;
501        let fork_filter_clone = fork_filter.clone();
502        let handle = tokio::spawn(async move {
503            // roughly based off of the design of tokio::net::TcpListener
504            let (incoming, _) = listener.accept().await.unwrap();
505            let stream = PassthroughCodec::default().framed(incoming);
506            let handshake_res = UnauthedEthStream::new(stream)
507                .handshake::<EthNetworkPrimitives>(status_clone, fork_filter_clone)
508                .await;
509
510            // make sure the handshake fails due to td too high
511            assert!(matches!(
512                handshake_res,
513                Err(EthStreamError::EthHandshakeError(
514                    EthHandshakeError::TotalDifficultyBitLenTooLarge { got: 165, maximum: 160 }
515                ))
516            ));
517        });
518
519        let outgoing = TcpStream::connect(local_addr).await.unwrap();
520        let sink = PassthroughCodec::default().framed(outgoing);
521
522        // try to connect
523        let handshake_res = UnauthedEthStream::new(sink)
524            .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
525            .await;
526
527        // this handshake should also fail due to td too high
528        assert!(matches!(
529            handshake_res,
530            Err(EthStreamError::EthHandshakeError(
531                EthHandshakeError::TotalDifficultyBitLenTooLarge { got: 165, maximum: 160 }
532            ))
533        ));
534
535        // await the other handshake
536        handle.await.unwrap();
537    }
538
539    #[tokio::test]
540    async fn can_write_and_read_cleartext() {
541        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
542        let local_addr = listener.local_addr().unwrap();
543        let test_msg = EthMessage::<EthNetworkPrimitives>::NewBlockHashes(
544            vec![
545                BlockHashNumber { hash: B256::random(), number: 5 },
546                BlockHashNumber { hash: B256::random(), number: 6 },
547            ]
548            .into(),
549        );
550
551        let test_msg_clone = test_msg.clone();
552        let handle = tokio::spawn(async move {
553            // roughly based off of the design of tokio::net::TcpListener
554            let (incoming, _) = listener.accept().await.unwrap();
555            let stream = PassthroughCodec::default().framed(incoming);
556            let mut stream = EthStream::new(EthVersion::Eth67, stream);
557
558            // use the stream to get the next message
559            let message = stream.next().await.unwrap().unwrap();
560            assert_eq!(message, test_msg_clone);
561        });
562
563        let outgoing = TcpStream::connect(local_addr).await.unwrap();
564        let sink = PassthroughCodec::default().framed(outgoing);
565        let mut client_stream = EthStream::new(EthVersion::Eth67, sink);
566
567        client_stream.send(test_msg).await.unwrap();
568
569        // make sure the server receives the message and asserts before ending the test
570        handle.await.unwrap();
571    }
572
573    #[tokio::test]
574    async fn can_write_and_read_ecies() {
575        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
576        let local_addr = listener.local_addr().unwrap();
577        let server_key = SecretKey::new(&mut rand_08::thread_rng());
578        let test_msg = EthMessage::<EthNetworkPrimitives>::NewBlockHashes(
579            vec![
580                BlockHashNumber { hash: B256::random(), number: 5 },
581                BlockHashNumber { hash: B256::random(), number: 6 },
582            ]
583            .into(),
584        );
585
586        let test_msg_clone = test_msg.clone();
587        let handle = tokio::spawn(async move {
588            // roughly based off of the design of tokio::net::TcpListener
589            let (incoming, _) = listener.accept().await.unwrap();
590            let stream = ECIESStream::incoming(incoming, server_key).await.unwrap();
591            let mut stream = EthStream::new(EthVersion::Eth67, stream);
592
593            // use the stream to get the next message
594            let message = stream.next().await.unwrap().unwrap();
595            assert_eq!(message, test_msg_clone);
596        });
597
598        // create the server pubkey
599        let server_id = pk2id(&server_key.public_key(SECP256K1));
600
601        let client_key = SecretKey::new(&mut rand_08::thread_rng());
602
603        let outgoing = TcpStream::connect(local_addr).await.unwrap();
604        let outgoing = ECIESStream::connect(outgoing, client_key, server_id).await.unwrap();
605        let mut client_stream = EthStream::new(EthVersion::Eth67, outgoing);
606
607        client_stream.send(test_msg).await.unwrap();
608
609        // make sure the server receives the message and asserts before ending the test
610        handle.await.unwrap();
611    }
612
613    #[tokio::test(flavor = "multi_thread")]
614    async fn ethstream_over_p2p() {
615        // create a p2p stream and server, then confirm that the two are authed
616        // create tcpstream
617        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
618        let local_addr = listener.local_addr().unwrap();
619        let server_key = SecretKey::new(&mut rand_08::thread_rng());
620        let test_msg = EthMessage::<EthNetworkPrimitives>::NewBlockHashes(
621            vec![
622                BlockHashNumber { hash: B256::random(), number: 5 },
623                BlockHashNumber { hash: B256::random(), number: 6 },
624            ]
625            .into(),
626        );
627
628        let genesis = B256::random();
629        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
630
631        let status = Status {
632            version: EthVersion::Eth67,
633            chain: NamedChain::Mainnet.into(),
634            total_difficulty: U256::ZERO,
635            blockhash: B256::random(),
636            genesis,
637            // Pass the current fork id.
638            forkid: fork_filter.current(),
639        };
640        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
641
642        let status_copy = unified_status;
643        let fork_filter_clone = fork_filter.clone();
644        let test_msg_clone = test_msg.clone();
645        let handle = tokio::spawn(async move {
646            // roughly based off of the design of tokio::net::TcpListener
647            let (incoming, _) = listener.accept().await.unwrap();
648            let stream = ECIESStream::incoming(incoming, server_key).await.unwrap();
649
650            let server_hello = HelloMessageWithProtocols {
651                protocol_version: ProtocolVersion::V5,
652                client_version: "bitcoind/1.0.0".to_string(),
653                protocols: vec![EthVersion::Eth67.into()],
654                port: DEFAULT_TCP_PORT,
655                id: pk2id(&server_key.public_key(SECP256K1)),
656            };
657
658            let unauthed_stream = UnauthedP2PStream::new(stream);
659            let (p2p_stream, _) = unauthed_stream.handshake(server_hello).await.unwrap();
660            let (mut eth_stream, _) = UnauthedEthStream::new(p2p_stream)
661                .handshake(status_copy, fork_filter_clone)
662                .await
663                .unwrap();
664
665            // use the stream to get the next message
666            let message = eth_stream.next().await.unwrap().unwrap();
667            assert_eq!(message, test_msg_clone);
668        });
669
670        // create the server pubkey
671        let server_id = pk2id(&server_key.public_key(SECP256K1));
672
673        let client_key = SecretKey::new(&mut rand_08::thread_rng());
674
675        let outgoing = TcpStream::connect(local_addr).await.unwrap();
676        let sink = ECIESStream::connect(outgoing, client_key, server_id).await.unwrap();
677
678        let client_hello = HelloMessageWithProtocols {
679            protocol_version: ProtocolVersion::V5,
680            client_version: "bitcoind/1.0.0".to_string(),
681            protocols: vec![EthVersion::Eth67.into()],
682            port: DEFAULT_TCP_PORT,
683            id: pk2id(&client_key.public_key(SECP256K1)),
684        };
685
686        let unauthed_stream = UnauthedP2PStream::new(sink);
687        let (p2p_stream, _) = unauthed_stream.handshake(client_hello).await.unwrap();
688
689        let (mut client_stream, _) = UnauthedEthStream::new(p2p_stream)
690            .handshake(unified_status, fork_filter)
691            .await
692            .unwrap();
693
694        client_stream.send(test_msg).await.unwrap();
695
696        // make sure the server receives the message and asserts before ending the test
697        handle.await.unwrap();
698    }
699
700    #[tokio::test]
701    async fn handshake_should_timeout() {
702        let genesis = B256::random();
703        let fork_filter = ForkFilter::new(Head::default(), genesis, 0, Vec::new());
704
705        let status = Status {
706            version: EthVersion::Eth67,
707            chain: NamedChain::Mainnet.into(),
708            total_difficulty: U256::ZERO,
709            blockhash: B256::random(),
710            genesis,
711            // Pass the current fork id.
712            forkid: fork_filter.current(),
713        };
714        let unified_status = UnifiedStatus::from_message(StatusMessage::Legacy(status));
715
716        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
717        let local_addr = listener.local_addr().unwrap();
718
719        let status_clone = unified_status;
720        let fork_filter_clone = fork_filter.clone();
721        let _handle = tokio::spawn(async move {
722            // Delay accepting the connection for longer than the client's timeout period
723            tokio::time::sleep(Duration::from_secs(11)).await;
724            // roughly based off of the design of tokio::net::TcpListener
725            let (incoming, _) = listener.accept().await.unwrap();
726            let stream = PassthroughCodec::default().framed(incoming);
727            let (_, their_status) = UnauthedEthStream::new(stream)
728                .handshake::<EthNetworkPrimitives>(status_clone, fork_filter_clone)
729                .await
730                .unwrap();
731
732            // just make sure it equals our status (our status is a clone of their status)
733            assert_eq!(their_status, status_clone);
734        });
735
736        let outgoing = TcpStream::connect(local_addr).await.unwrap();
737        let sink = PassthroughCodec::default().framed(outgoing);
738
739        // try to connect
740        let handshake_result = UnauthedEthStream::new(sink)
741            .handshake_with_timeout::<EthNetworkPrimitives>(
742                unified_status,
743                fork_filter,
744                Duration::from_secs(1),
745            )
746            .await;
747
748        // Assert that a timeout error occurred
749        assert!(
750            matches!(handshake_result, Err(e) if e.to_string() == EthStreamError::StreamTimeout.to_string())
751        );
752    }
753
754    #[tokio::test]
755    async fn can_write_and_read_raw_capability() {
756        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
757        let local_addr = listener.local_addr().unwrap();
758
759        let test_msg = RawCapabilityMessage { id: 0x1234, payload: Bytes::from(vec![1, 2, 3, 4]) };
760
761        let test_msg_clone = test_msg.clone();
762        let handle = tokio::spawn(async move {
763            let (incoming, _) = listener.accept().await.unwrap();
764            let stream = PassthroughCodec::default().framed(incoming);
765            let mut stream = EthStream::<_, EthNetworkPrimitives>::new(EthVersion::Eth67, stream);
766
767            let bytes = stream.inner_mut().next().await.unwrap().unwrap();
768
769            // Create a cursor to track position while decoding
770            let mut id_bytes = &bytes[..];
771            let decoded_id = <usize as Decodable>::decode(&mut id_bytes).unwrap();
772            assert_eq!(decoded_id, test_msg_clone.id);
773
774            // Get remaining bytes after ID decoding
775            let remaining = id_bytes;
776            assert_eq!(remaining, &test_msg_clone.payload[..]);
777        });
778
779        let outgoing = TcpStream::connect(local_addr).await.unwrap();
780        let sink = PassthroughCodec::default().framed(outgoing);
781        let mut client_stream = EthStream::<_, EthNetworkPrimitives>::new(EthVersion::Eth67, sink);
782
783        client_stream.start_send_raw(test_msg).unwrap();
784        client_stream.inner_mut().flush().await.unwrap();
785
786        handle.await.unwrap();
787    }
788
789    #[tokio::test]
790    async fn status_message_after_handshake_triggers_disconnect() {
791        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
792        let local_addr = listener.local_addr().unwrap();
793
794        let handle = tokio::spawn(async move {
795            let (incoming, _) = listener.accept().await.unwrap();
796            let stream = PassthroughCodec::default().framed(incoming);
797            let mut stream = EthStream::<_, EthNetworkPrimitives>::new(EthVersion::Eth67, stream);
798
799            // Try to send a Status message after handshake - this should trigger disconnect
800            let status = Status {
801                version: EthVersion::Eth67,
802                chain: NamedChain::Mainnet.into(),
803                total_difficulty: U256::ZERO,
804                blockhash: B256::random(),
805                genesis: B256::random(),
806                forkid: ForkFilter::new(Head::default(), B256::random(), 0, Vec::new()).current(),
807            };
808            let status_message =
809                EthMessage::<EthNetworkPrimitives>::Status(StatusMessage::Legacy(status));
810
811            // This should return an error and trigger disconnect
812            let result = stream.send(status_message).await;
813            assert!(result.is_err());
814            assert!(matches!(
815                result.unwrap_err(),
816                EthStreamError::EthHandshakeError(EthHandshakeError::StatusNotInHandshake)
817            ));
818        });
819
820        let outgoing = TcpStream::connect(local_addr).await.unwrap();
821        let sink = PassthroughCodec::default().framed(outgoing);
822        let mut client_stream = EthStream::<_, EthNetworkPrimitives>::new(EthVersion::Eth67, sink);
823
824        // Send a valid message to keep the connection alive
825        let test_msg = EthMessage::<EthNetworkPrimitives>::NewBlockHashes(
826            vec![BlockHashNumber { hash: B256::random(), number: 5 }].into(),
827        );
828        client_stream.send(test_msg).await.unwrap();
829
830        handle.await.unwrap();
831    }
832}
reth_eth_wire/ethstream.rs

reth_eth_wire/
ethstream.rs
