1use 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#[pin_project]
35#[derive(Debug)]
36pub struct UnauthedEthStream<S> {
37 #[pin]
38 inner: S,
39}
40
41impl<S> UnauthedEthStream<S> {
42 pub const fn new(inner: S) -> Self {
44 Self { inner }
45 }
46
47 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 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 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 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 let stream = EthStream::new(status.version, self.inner);
100
101 Ok((stream, their_status))
102 }
103}
104
105#[derive(Debug)]
107pub struct EthStreamInner<N> {
108 version: EthVersion,
110 max_message_size: usize,
112 reject_block_announcements: bool,
115 pending_responses: BTreeMap<u64, PendingResponse>,
122 _pd: std::marker::PhantomData<N>,
123}
124
125impl<N> EthStreamInner<N>
126where
127 N: NetworkPrimitives,
128{
129 pub const fn new(version: EthVersion) -> Self {
131 Self::with_max_message_size(version, MAX_MESSAGE_SIZE)
132 }
133
134 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 #[inline]
147 pub const fn version(&self) -> EthVersion {
148 self.version
149 }
150
151 pub const fn set_reject_block_announcements(&mut self, reject: bool) {
154 self.reject_block_announcements = reject;
155 }
156
157 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 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 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 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#[pin_project]
272#[derive(Debug)]
273pub struct EthStream<S, N = EthNetworkPrimitives> {
274 eth: EthStreamInner<N>,
276 #[pin]
277 inner: S,
278}
279
280impl<S, N: NetworkPrimitives> EthStream<S, N> {
281 #[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 #[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 #[inline]
300 pub const fn version(&self) -> EthVersion {
301 self.eth.version()
302 }
303
304 pub const fn set_reject_block_announcements(&mut self, reject: bool) {
307 self.eth.set_reject_block_announcements(reject);
308 }
309
310 #[inline]
312 pub const fn inner(&self) -> &S {
313 &self.inner
314 }
315
316 #[inline]
318 pub const fn inner_mut(&mut self) -> &mut S {
319 &mut self.inner
320 }
321
322 #[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 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 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 let mut this = self.project();
391 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#[derive(Clone, Copy, Debug)]
430struct PendingResponse {
431 message_id: EthMessageID,
433 max_items: usize,
435}
436
437impl PendingResponse {
438 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
473fn 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 EthMessageID::Receipts if version >= EthVersion::Eth70 => {
484 bool::decode(&mut payload)?;
485 }
486 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
497pub(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 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 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 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 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 let (_, their_status) = UnauthedEthStream::new(sink)
586 .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
587 .await
588 .unwrap();
589
590 assert_eq!(their_status, unified_status);
592
593 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 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 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 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 let (_, their_status) = UnauthedEthStream::new(sink)
636 .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
637 .await
638 .unwrap();
639
640 assert_eq!(their_status, unified_status);
642
643 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 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 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 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 let handshake_res = UnauthedEthStream::new(sink)
690 .handshake::<EthNetworkPrimitives>(unified_status, fork_filter)
691 .await;
692
693 assert!(matches!(
695 handshake_res,
696 Err(EthStreamError::EthHandshakeError(
697 EthHandshakeError::TotalDifficultyBitLenTooLarge { got: 165, maximum: 160 }
698 ))
699 ));
700
701 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 let (incoming, _) = listener.accept().await.unwrap();
721 let stream = PassthroughCodec::default().framed(incoming);
722 let mut stream = EthStream::new(EthVersion::Eth67, stream);
723
724 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 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 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 let message = stream.next().await.unwrap().unwrap();
761 assert_eq!(message, test_msg_clone);
762 });
763
764 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 handle.await.unwrap();
777 }
778
779 #[tokio::test(flavor = "multi_thread")]
780 async fn ethstream_over_p2p() {
781 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 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 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 let message = eth_stream.next().await.unwrap().unwrap();
833 assert_eq!(message, test_msg_clone);
834 });
835
836 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 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 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 tokio::time::sleep(Duration::from_secs(11)).await;
890 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 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 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!(
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 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 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 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 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 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 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 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}