reth_eth_wire/
multiplex.rs

1//! Rlpx protocol multiplexer and satellite stream
2//!
3//! A Satellite is a Stream that primarily drives a single `RLPx` subprotocol but can also handle
4//! additional subprotocols.
5//!
6//! Most of other subprotocols are "dependent satellite" protocols of "eth" and not a fully standalone protocol, for example "snap", See also [snap protocol](https://github.com/ethereum/devp2p/blob/298d7a77c3bf833641579ecbbb5b13f0311eeeea/caps/snap.md?plain=1#L71)
7//! Hence it is expected that the primary protocol is "eth" and the additional protocols are
8//! "dependent satellite" protocols.
9
10use std::{
11    collections::VecDeque,
12    fmt,
13    future::Future,
14    io,
15    pin::{pin, Pin},
16    sync::Arc,
17    task::{ready, Context, Poll},
18};
19
20use crate::{
21    capability::{SharedCapabilities, SharedCapability, UnsupportedCapabilityError},
22    errors::{EthStreamError, P2PStreamError},
23    handshake::EthRlpxHandshake,
24    p2pstream::DisconnectP2P,
25    CanDisconnect, Capability, DisconnectReason, EthStream, P2PStream, UnifiedStatus,
26    HANDSHAKE_TIMEOUT,
27};
28use bytes::{Bytes, BytesMut};
29use futures::{Sink, SinkExt, Stream, StreamExt, TryStream, TryStreamExt};
30use reth_eth_wire_types::NetworkPrimitives;
31use reth_ethereum_forks::ForkFilter;
32use tokio::sync::{mpsc, mpsc::UnboundedSender};
33use tokio_stream::wrappers::UnboundedReceiverStream;
34
35/// A Stream and Sink type that wraps a raw rlpx stream [`P2PStream`] and handles message ID
36/// multiplexing.
37#[derive(Debug)]
38pub struct RlpxProtocolMultiplexer<St> {
39    inner: MultiplexInner<St>,
40}
41
42impl<St> RlpxProtocolMultiplexer<St> {
43    /// Wraps the raw p2p stream
44    pub fn new(conn: P2PStream<St>) -> Self {
45        Self {
46            inner: MultiplexInner {
47                conn,
48                protocols: Default::default(),
49                out_buffer: Default::default(),
50            },
51        }
52    }
53
54    /// Installs a new protocol on top of the raw p2p stream.
55    ///
56    /// This accepts a closure that receives a [`ProtocolConnection`] that will yield messages for
57    /// the given capability.
58    pub fn install_protocol<F, Proto>(
59        &mut self,
60        cap: &Capability,
61        f: F,
62    ) -> Result<(), UnsupportedCapabilityError>
63    where
64        F: FnOnce(ProtocolConnection) -> Proto,
65        Proto: Stream<Item = BytesMut> + Send + 'static,
66    {
67        self.inner.install_protocol(cap, f)
68    }
69
70    /// Returns the [`SharedCapabilities`] of the underlying raw p2p stream
71    pub const fn shared_capabilities(&self) -> &SharedCapabilities {
72        self.inner.shared_capabilities()
73    }
74
75    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
76    pub fn into_satellite_stream<F, Primary>(
77        self,
78        cap: &Capability,
79        primary: F,
80    ) -> Result<RlpxSatelliteStream<St, Primary>, P2PStreamError>
81    where
82        F: FnOnce(ProtocolProxy) -> Primary,
83    {
84        let Ok(shared_cap) = self.shared_capabilities().ensure_matching_capability(cap).cloned()
85        else {
86            return Err(P2PStreamError::CapabilityNotShared)
87        };
88
89        let (to_primary, from_wire) = mpsc::unbounded_channel();
90        let (to_wire, from_primary) = mpsc::unbounded_channel();
91        let proxy = ProtocolProxy {
92            shared_cap: shared_cap.clone(),
93            from_wire: UnboundedReceiverStream::new(from_wire),
94            to_wire,
95        };
96
97        let st = primary(proxy);
98        Ok(RlpxSatelliteStream {
99            inner: self.inner,
100            primary: PrimaryProtocol {
101                to_primary,
102                from_primary: UnboundedReceiverStream::new(from_primary),
103                st,
104                shared_cap,
105            },
106            next_outbound: 0,
107        })
108    }
109
110    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
111    ///
112    /// Returns an error if the primary protocol is not supported by the remote or the handshake
113    /// failed.
114    pub async fn into_satellite_stream_with_handshake<F, Fut, Err, Primary>(
115        self,
116        cap: &Capability,
117        handshake: F,
118    ) -> Result<RlpxSatelliteStream<St, Primary>, Err>
119    where
120        F: FnOnce(ProtocolProxy) -> Fut,
121        Fut: Future<Output = Result<Primary, Err>>,
122        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
123        P2PStreamError: Into<Err>,
124    {
125        self.into_satellite_stream_with_tuple_handshake(cap, async move |proxy| {
126            let st = handshake(proxy).await?;
127            Ok((st, ()))
128        })
129        .await
130        .map(|(st, _)| st)
131    }
132
133    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
134    ///
135    /// Returns an error if the primary protocol is not supported by the remote or the handshake
136    /// failed.
137    ///
138    /// This accepts a closure that does a handshake with the remote peer and returns a tuple of the
139    /// primary stream and extra data.
140    ///
141    /// See also [`UnauthedEthStream::handshake`](crate::UnauthedEthStream)
142    pub async fn into_satellite_stream_with_tuple_handshake<F, Fut, Err, Primary, Extra>(
143        mut self,
144        cap: &Capability,
145        handshake: F,
146    ) -> Result<(RlpxSatelliteStream<St, Primary>, Extra), Err>
147    where
148        F: FnOnce(ProtocolProxy) -> Fut,
149        Fut: Future<Output = Result<(Primary, Extra), Err>>,
150        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
151        P2PStreamError: Into<Err>,
152    {
153        let Ok(shared_cap) = self.shared_capabilities().ensure_matching_capability(cap).cloned()
154        else {
155            return Err(P2PStreamError::CapabilityNotShared.into())
156        };
157
158        let (to_primary, from_wire) = mpsc::unbounded_channel();
159        let (to_wire, mut from_primary) = mpsc::unbounded_channel();
160        let proxy = ProtocolProxy {
161            shared_cap: shared_cap.clone(),
162            from_wire: UnboundedReceiverStream::new(from_wire),
163            to_wire,
164        };
165
166        let f = handshake(proxy);
167        let mut f = pin!(f);
168
169        // this polls the connection and the primary stream concurrently until the handshake is
170        // complete
171        loop {
172            tokio::select! {
173                biased;
174                Some(Ok(msg)) = self.inner.conn.next() => {
175                    // Ensure the message belongs to the primary protocol
176                    let Some(offset) = msg.first().copied()
177                    else {
178                        return Err(P2PStreamError::EmptyProtocolMessage.into())
179                    };
180                    if let Some(cap) = self.shared_capabilities().find_by_relative_offset(offset).cloned() {
181                            if cap == shared_cap {
182                                // delegate to primary
183                                let _ = to_primary.send(msg);
184                            } else {
185                                // delegate to satellite
186                                self.inner.delegate_message(&cap, msg);
187                            }
188                        } else {
189                           return Err(P2PStreamError::UnknownReservedMessageId(offset).into())
190                        }
191                }
192                Some(msg) = from_primary.recv() => {
193                    self.inner.conn.send(msg).await.map_err(Into::into)?;
194                }
195                // Poll all subprotocols for new messages
196                msg = ProtocolsPoller::new(&mut self.inner.protocols) => {
197                     self.inner.conn.send(msg.map_err(Into::into)?).await.map_err(Into::into)?;
198                }
199                res = &mut f => {
200                    let (st, extra) = res?;
201                    return Ok((
202                        RlpxSatelliteStream {
203                            inner: self.inner,
204                            primary: PrimaryProtocol {
205                                to_primary,
206                                from_primary: UnboundedReceiverStream::new(from_primary),
207                                st,
208                                shared_cap,
209                            },
210                            next_outbound: 0,
211                        },
212                        extra,
213                    ))
214                }
215            }
216        }
217    }
218
219    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with eth protocol as the given
220    /// primary protocol and the handshake implementation.
221    pub async fn into_eth_satellite_stream<N: NetworkPrimitives>(
222        self,
223        status: UnifiedStatus,
224        fork_filter: ForkFilter,
225        handshake: Arc<dyn EthRlpxHandshake>,
226        eth_max_message_size: usize,
227    ) -> Result<(RlpxSatelliteStream<St, EthStream<ProtocolProxy, N>>, UnifiedStatus), EthStreamError>
228    where
229        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
230    {
231        let eth_cap = self.inner.conn.shared_capabilities().eth_version()?;
232        self.into_satellite_stream_with_tuple_handshake(
233            &Capability::eth(eth_cap),
234            async move |proxy| {
235                let handshake = handshake.clone();
236                let mut unauth = UnauthProxy { inner: proxy };
237                let their_status = handshake
238                    .handshake(&mut unauth, status, fork_filter, HANDSHAKE_TIMEOUT)
239                    .await?;
240                let eth_stream = EthStream::with_max_message_size(
241                    eth_cap,
242                    unauth.into_inner(),
243                    eth_max_message_size,
244                );
245                Ok((eth_stream, their_status))
246            },
247        )
248        .await
249    }
250}
251
252#[derive(Debug)]
253struct MultiplexInner<St> {
254    /// The raw p2p stream
255    conn: P2PStream<St>,
256    /// All the subprotocols that are multiplexed on top of the raw p2p stream
257    protocols: VecDeque<ProtocolStream>,
258    /// Buffer for outgoing messages on the wire.
259    out_buffer: OutBuffer,
260}
261
262impl<St> MultiplexInner<St> {
263    const fn shared_capabilities(&self) -> &SharedCapabilities {
264        self.conn.shared_capabilities()
265    }
266
267    /// Delegates a message to the matching protocol.
268    fn delegate_message(&self, cap: &SharedCapability, msg: BytesMut) -> bool {
269        for proto in &self.protocols {
270            if proto.shared_cap == *cap {
271                proto.send_raw(msg);
272                return true
273            }
274        }
275        false
276    }
277
278    fn install_protocol<F, Proto>(
279        &mut self,
280        cap: &Capability,
281        f: F,
282    ) -> Result<(), UnsupportedCapabilityError>
283    where
284        F: FnOnce(ProtocolConnection) -> Proto,
285        Proto: Stream<Item = BytesMut> + Send + 'static,
286    {
287        let shared_cap =
288            self.conn.shared_capabilities().ensure_matching_capability(cap).cloned()?;
289        let (to_satellite, rx) = mpsc::unbounded_channel();
290        let proto_conn = ProtocolConnection { from_wire: UnboundedReceiverStream::new(rx) };
291        let st = f(proto_conn);
292        let st = ProtocolStream { shared_cap, to_satellite, satellite_st: Box::pin(st) };
293        self.protocols.push_back(st);
294        Ok(())
295    }
296}
297
298/// Represents a protocol in the multiplexer that is used as the primary protocol.
299#[derive(Debug)]
300struct PrimaryProtocol<Primary> {
301    /// Channel to send messages to the primary protocol.
302    to_primary: UnboundedSender<BytesMut>,
303    /// Receiver for messages from the primary protocol.
304    from_primary: UnboundedReceiverStream<Bytes>,
305    /// Shared capability of the primary protocol.
306    shared_cap: SharedCapability,
307    /// The primary stream.
308    st: Primary,
309}
310
311/// A Stream and Sink type that acts as a wrapper around a primary `RLPx` subprotocol (e.g. "eth")
312///
313/// Only emits and sends _non-empty_ messages
314#[derive(Debug)]
315pub struct ProtocolProxy {
316    shared_cap: SharedCapability,
317    /// Receives _non-empty_ messages from the wire
318    from_wire: UnboundedReceiverStream<BytesMut>,
319    /// Sends _non-empty_ messages from the wire
320    to_wire: UnboundedSender<Bytes>,
321}
322
323impl ProtocolProxy {
324    /// Sends a _non-empty_ message on the wire.
325    fn try_send(&self, msg: Bytes) -> Result<(), io::Error> {
326        if msg.is_empty() {
327            // message must not be empty
328            return Err(io::ErrorKind::InvalidInput.into())
329        }
330        self.to_wire.send(self.mask_msg_id(msg)?).map_err(|_| io::ErrorKind::BrokenPipe.into())
331    }
332
333    /// Masks the message ID of a message to be sent on the wire.
334    #[inline]
335    fn mask_msg_id(&self, msg: Bytes) -> Result<Bytes, io::Error> {
336        if msg.is_empty() {
337            // message must not be empty
338            return Err(io::ErrorKind::InvalidInput.into())
339        }
340
341        let offset = self.shared_cap.relative_message_id_offset();
342        if offset == 0 {
343            return Ok(msg);
344        }
345
346        let mut masked: BytesMut = msg.into();
347        masked[0] = masked[0].checked_add(offset).ok_or(io::ErrorKind::InvalidInput)?;
348        Ok(masked.freeze())
349    }
350
351    /// Unmasks the message ID of a message received from the wire.
352    #[inline]
353    fn unmask_id(&self, mut msg: BytesMut) -> Result<BytesMut, io::Error> {
354        if msg.is_empty() {
355            // message must not be empty
356            return Err(io::ErrorKind::InvalidInput.into())
357        }
358        msg[0] = msg[0]
359            .checked_sub(self.shared_cap.relative_message_id_offset())
360            .ok_or(io::ErrorKind::InvalidInput)?;
361        Ok(msg)
362    }
363}
364
365impl Stream for ProtocolProxy {
366    type Item = Result<BytesMut, io::Error>;
367
368    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
369        let msg = ready!(self.from_wire.poll_next_unpin(cx));
370        Poll::Ready(msg.map(|msg| self.get_mut().unmask_id(msg)))
371    }
372}
373
374impl Sink<Bytes> for ProtocolProxy {
375    type Error = io::Error;
376
377    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
378        Poll::Ready(Ok(()))
379    }
380
381    fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
382        self.get_mut().try_send(item)
383    }
384
385    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
386        Poll::Ready(Ok(()))
387    }
388
389    fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
390        Poll::Ready(Ok(()))
391    }
392}
393
394impl CanDisconnect<Bytes> for ProtocolProxy {
395    fn disconnect(
396        &mut self,
397        _reason: DisconnectReason,
398    ) -> Pin<Box<dyn Future<Output = Result<(), <Self as Sink<Bytes>>::Error>> + Send + '_>> {
399        Box::pin(async move { Ok(()) })
400    }
401}
402
403/// Adapter so the injected `EthRlpxHandshake` can run over a multiplexed `ProtocolProxy`
404/// using the same error type expectations (`P2PStreamError`).
405#[derive(Debug)]
406struct UnauthProxy {
407    inner: ProtocolProxy,
408}
409
410impl UnauthProxy {
411    fn into_inner(self) -> ProtocolProxy {
412        self.inner
413    }
414}
415
416impl Stream for UnauthProxy {
417    type Item = Result<BytesMut, P2PStreamError>;
418
419    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
420        self.inner.poll_next_unpin(cx).map(|opt| opt.map(|res| res.map_err(P2PStreamError::from)))
421    }
422}
423
424impl Sink<Bytes> for UnauthProxy {
425    type Error = P2PStreamError;
426
427    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
428        self.inner.poll_ready_unpin(cx).map_err(P2PStreamError::from)
429    }
430
431    fn start_send(mut self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
432        self.inner.start_send_unpin(item).map_err(P2PStreamError::from)
433    }
434
435    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
436        self.inner.poll_flush_unpin(cx).map_err(P2PStreamError::from)
437    }
438
439    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
440        self.inner.poll_close_unpin(cx).map_err(P2PStreamError::from)
441    }
442}
443
444impl CanDisconnect<Bytes> for UnauthProxy {
445    fn disconnect(
446        &mut self,
447        reason: DisconnectReason,
448    ) -> Pin<Box<dyn Future<Output = Result<(), <Self as Sink<Bytes>>::Error>> + Send + '_>> {
449        let fut = self.inner.disconnect(reason);
450        Box::pin(async move { fut.await.map_err(P2PStreamError::from) })
451    }
452}
453
454/// A connection channel to receive _`non_empty`_ messages for the negotiated protocol.
455///
456/// This is a [Stream] that returns raw bytes of the received messages for this protocol.
457#[derive(Debug)]
458pub struct ProtocolConnection {
459    from_wire: UnboundedReceiverStream<BytesMut>,
460}
461
462impl Stream for ProtocolConnection {
463    type Item = BytesMut;
464
465    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
466        self.from_wire.poll_next_unpin(cx)
467    }
468}
469
470/// A Stream and Sink type that acts as a wrapper around a primary `RLPx` subprotocol (e.g. "eth")
471/// [`EthStream`] and can also handle additional subprotocols.
472#[derive(Debug)]
473pub struct RlpxSatelliteStream<St, Primary> {
474    inner: MultiplexInner<St>,
475    primary: PrimaryProtocol<Primary>,
476    /// Round-robin cursor for the next outbound producer to poll.
477    next_outbound: usize,
478}
479
480impl<St, Primary> RlpxSatelliteStream<St, Primary> {
481    /// Installs a new protocol on top of the raw p2p stream.
482    ///
483    /// This accepts a closure that receives a [`ProtocolConnection`] that will yield messages for
484    /// the given capability.
485    pub fn install_protocol<F, Proto>(
486        &mut self,
487        cap: &Capability,
488        f: F,
489    ) -> Result<(), UnsupportedCapabilityError>
490    where
491        F: FnOnce(ProtocolConnection) -> Proto,
492        Proto: Stream<Item = BytesMut> + Send + 'static,
493    {
494        self.inner.install_protocol(cap, f)
495    }
496
497    /// Returns the primary protocol.
498    #[inline]
499    pub const fn primary(&self) -> &Primary {
500        &self.primary.st
501    }
502
503    /// Returns mutable access to the primary protocol.
504    #[inline]
505    pub const fn primary_mut(&mut self) -> &mut Primary {
506        &mut self.primary.st
507    }
508
509    /// Returns the underlying [`P2PStream`].
510    #[inline]
511    pub const fn inner(&self) -> &P2PStream<St> {
512        &self.inner.conn
513    }
514
515    /// Returns mutable access to the underlying [`P2PStream`].
516    #[inline]
517    pub const fn inner_mut(&mut self) -> &mut P2PStream<St> {
518        &mut self.inner.conn
519    }
520
521    /// Consumes this type and returns the wrapped [`P2PStream`].
522    #[inline]
523    pub fn into_inner(self) -> P2PStream<St> {
524        self.inner.conn
525    }
526
527    /// Polls primary and satellite outbound producers round-robin until the `OutBuffer` is full or
528    /// every producer is pending.
529    ///
530    /// The cursor advances after each producer poll, so a ready producer cannot drain repeatedly
531    /// before later producers get a turn.
532    fn poll_outbound_producers(&mut self, cx: &mut Context<'_>) -> Result<ProducerPoll, io::Error> {
533        let producers = self.inner.protocols.len() + 1;
534        let mut pending = 0;
535
536        while pending < producers {
537            if self.inner.out_buffer.is_full() {
538                return Ok(ProducerPoll::Full)
539            }
540
541            if self.next_outbound >= producers {
542                self.next_outbound = 0;
543            }
544
545            let producer = self.next_outbound;
546            self.next_outbound = (self.next_outbound + 1) % producers;
547
548            let msg = if producer == 0 {
549                match self.primary.from_primary.poll_next_unpin(cx) {
550                    Poll::Ready(Some(msg)) => msg,
551                    Poll::Ready(None) => return Ok(ProducerPoll::Closed),
552                    Poll::Pending => {
553                        pending += 1;
554                        continue
555                    }
556                }
557            } else {
558                let proto = self
559                    .inner
560                    .protocols
561                    .get_mut(producer - 1)
562                    .expect("outbound producer index checked against protocol count");
563                match proto.poll_next_unpin(cx) {
564                    Poll::Ready(Some(msg)) => msg?,
565                    Poll::Ready(None) => return Ok(ProducerPoll::Closed),
566                    Poll::Pending => {
567                        pending += 1;
568                        continue
569                    }
570                }
571            };
572
573            pending = 0;
574            self.inner.out_buffer.push_back(msg);
575        }
576
577        Ok(ProducerPoll::Pending)
578    }
579}
580
581impl<St, Primary, PrimaryErr> Stream for RlpxSatelliteStream<St, Primary>
582where
583    St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
584    Primary: TryStream<Error = PrimaryErr> + Unpin,
585    P2PStreamError: Into<PrimaryErr>,
586{
587    type Item = Result<Primary::Ok, Primary::Error>;
588
589    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
590        let this = self.get_mut();
591
592        loop {
593            // first drain the primary stream
594            if let Poll::Ready(Some(msg)) = this.primary.st.try_poll_next_unpin(cx) {
595                return Poll::Ready(Some(msg))
596            }
597
598            let mut conn_ready = true;
599            loop {
600                match this.inner.conn.poll_ready_unpin(cx) {
601                    Poll::Ready(Ok(())) => {
602                        if let Some(msg) = this.inner.out_buffer.pop_front() {
603                            if let Err(err) = this.inner.conn.start_send_unpin(msg) {
604                                return Poll::Ready(Some(Err(err.into())))
605                            }
606                        } else {
607                            break
608                        }
609                    }
610                    Poll::Ready(Err(err)) => {
611                        if let Err(disconnect_err) =
612                            this.inner.conn.start_disconnect(DisconnectReason::DisconnectRequested)
613                        {
614                            return Poll::Ready(Some(Err(disconnect_err.into())))
615                        }
616                        return Poll::Ready(Some(Err(err.into())))
617                    }
618                    Poll::Pending => {
619                        conn_ready = false;
620                        break
621                    }
622                }
623            }
624            // The connection only buffers frames on `start_send`; `poll_flush` performs the
625            // actual writes and flushes the transport once for the batch handed to it above.
626            // This also resumes a flush that returned pending on an earlier pass; a no-op if
627            // nothing is buffered.
628            match this.inner.conn.poll_flush_unpin(cx) {
629                Poll::Ready(Ok(())) => {}
630                Poll::Ready(Err(err)) => return Poll::Ready(Some(Err(err.into()))),
631                Poll::Pending => {
632                    conn_ready = false;
633                }
634            }
635
636            match this.poll_outbound_producers(cx) {
637                Ok(ProducerPoll::Pending | ProducerPoll::Full) => {}
638                Ok(ProducerPoll::Closed) => return Poll::Ready(None),
639                Err(err) => return Poll::Ready(Some(Err(P2PStreamError::Io(err).into()))),
640            }
641
642            let mut delegated = false;
643            loop {
644                // pull messages from connection
645                match this.inner.conn.poll_next_unpin(cx) {
646                    Poll::Ready(Some(Ok(msg))) => {
647                        delegated = true;
648                        let Some(offset) = msg.first().copied() else {
649                            return Poll::Ready(Some(Err(
650                                P2PStreamError::EmptyProtocolMessage.into()
651                            )))
652                        };
653                        // delegate the multiplexed message to the correct protocol
654                        if let Some(cap) =
655                            this.inner.conn.shared_capabilities().find_by_relative_offset(offset)
656                        {
657                            if cap == &this.primary.shared_cap {
658                                // delegate to primary
659                                let _ = this.primary.to_primary.send(msg);
660                            } else {
661                                // delegate to installed satellite if any
662                                for proto in &this.inner.protocols {
663                                    if proto.shared_cap == *cap {
664                                        proto.send_raw(msg);
665                                        break
666                                    }
667                                }
668                            }
669                        } else {
670                            return Poll::Ready(Some(Err(P2PStreamError::UnknownReservedMessageId(
671                                offset,
672                            )
673                            .into())))
674                        }
675                    }
676                    Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err.into()))),
677                    Poll::Ready(None) => {
678                        // connection closed
679                        return Poll::Ready(None)
680                    }
681                    Poll::Pending => break,
682                }
683            }
684
685            if !conn_ready || (!delegated && this.inner.out_buffer.is_empty()) {
686                return Poll::Pending
687            }
688        }
689    }
690}
691
692impl<St, Primary, T> Sink<T> for RlpxSatelliteStream<St, Primary>
693where
694    St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
695    Primary: Sink<T> + Unpin,
696    P2PStreamError: Into<<Primary as Sink<T>>::Error>,
697{
698    type Error = <Primary as Sink<T>>::Error;
699
700    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
701        let this = self.get_mut();
702        if let Err(err) = ready!(this.inner.conn.poll_ready_unpin(cx)) {
703            return Poll::Ready(Err(err.into()))
704        }
705        if let Err(err) = ready!(this.primary.st.poll_ready_unpin(cx)) {
706            return Poll::Ready(Err(err))
707        }
708        Poll::Ready(Ok(()))
709    }
710
711    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
712        self.get_mut().primary.st.start_send_unpin(item)
713    }
714
715    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
716        self.get_mut().inner.conn.poll_flush_unpin(cx).map_err(Into::into)
717    }
718
719    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
720        self.get_mut().inner.conn.poll_close_unpin(cx).map_err(Into::into)
721    }
722}
723
724/// Wraps a `RLPx` subprotocol and handles message ID multiplexing.
725struct ProtocolStream {
726    shared_cap: SharedCapability,
727    /// the channel shared with the satellite stream
728    to_satellite: UnboundedSender<BytesMut>,
729    satellite_st: Pin<Box<dyn Stream<Item = BytesMut> + Send>>,
730}
731
732impl ProtocolStream {
733    /// Masks the message ID of a message to be sent on the wire.
734    #[inline]
735    fn mask_msg_id(&self, mut msg: BytesMut) -> Result<Bytes, io::Error> {
736        if msg.is_empty() {
737            // message must not be empty
738            return Err(io::ErrorKind::InvalidInput.into())
739        }
740        msg[0] = msg[0]
741            .checked_add(self.shared_cap.relative_message_id_offset())
742            .ok_or(io::ErrorKind::InvalidInput)?;
743        Ok(msg.freeze())
744    }
745
746    /// Unmasks the message ID of a message received from the wire.
747    #[inline]
748    fn unmask_id(&self, mut msg: BytesMut) -> Result<BytesMut, io::Error> {
749        if msg.is_empty() {
750            // message must not be empty
751            return Err(io::ErrorKind::InvalidInput.into())
752        }
753        msg[0] = msg[0]
754            .checked_sub(self.shared_cap.relative_message_id_offset())
755            .ok_or(io::ErrorKind::InvalidInput)?;
756        Ok(msg)
757    }
758
759    /// Sends the message to the satellite stream.
760    fn send_raw(&self, msg: BytesMut) {
761        let _ = self.unmask_id(msg).map(|msg| self.to_satellite.send(msg));
762    }
763}
764
765impl Stream for ProtocolStream {
766    type Item = Result<Bytes, io::Error>;
767
768    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
769        let this = self.get_mut();
770        let msg = ready!(this.satellite_st.as_mut().poll_next(cx));
771        Poll::Ready(msg.map(|msg| this.mask_msg_id(msg)))
772    }
773}
774
775impl fmt::Debug for ProtocolStream {
776    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
777        f.debug_struct("ProtocolStream").field("cap", &self.shared_cap).finish_non_exhaustive()
778    }
779}
780
781/// Helper to poll multiple protocol streams in a `tokio::select`! branch
782struct ProtocolsPoller<'a> {
783    protocols: &'a mut VecDeque<ProtocolStream>,
784}
785
786impl<'a> ProtocolsPoller<'a> {
787    const fn new(protocols: &'a mut VecDeque<ProtocolStream>) -> Self {
788        Self { protocols }
789    }
790}
791
792impl<'a> Future for ProtocolsPoller<'a> {
793    type Output = Result<Bytes, P2PStreamError>;
794
795    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
796        let protocols = self.protocols.len();
797        for _ in 0..protocols {
798            let mut proto = self.protocols.pop_front().expect("protocol count checked");
799            match proto.poll_next_unpin(cx) {
800                Poll::Ready(Some(Err(err))) => {
801                    self.protocols.push_back(proto);
802                    return Poll::Ready(Err(P2PStreamError::from(err)))
803                }
804                Poll::Ready(Some(Ok(msg))) => {
805                    // Got a message, put protocol back and return the message
806                    self.protocols.push_back(proto);
807                    return Poll::Ready(Ok(msg));
808                }
809                _ => {
810                    // push it back because we still want to complete the handshake first
811                    self.protocols.push_back(proto);
812                }
813            }
814        }
815
816        // All protocols processed, nothing ready
817        Poll::Pending
818    }
819}
820
821/// Soft cap for per-connection outbound `RLPx` messages waiting in the multiplexer.
822///
823/// The cap is soft because the next message size is only known after polling a protocol stream.
824/// The buffer may exceed this by at most one message before producer polling is paused.
825///
826/// The lower [`P2PStream`] sink admits two outbound messages and rejects uncompressed payloads
827/// above 16 MiB, so 32 MiB mirrors the largest payload volume the lower p2p layer is already
828/// prepared to buffer.
829const MAX_MUX_OUT_BUFFER_BYTES: usize = 32 * 1024 * 1024;
830
831#[derive(Debug)]
832struct OutBuffer {
833    messages: VecDeque<Bytes>,
834    bytes: usize,
835    max_bytes: usize,
836}
837
838impl Default for OutBuffer {
839    fn default() -> Self {
840        Self { messages: Default::default(), bytes: 0, max_bytes: MAX_MUX_OUT_BUFFER_BYTES }
841    }
842}
843
844impl OutBuffer {
845    fn push_back(&mut self, msg: Bytes) {
846        self.bytes += msg.len();
847        self.messages.push_back(msg);
848    }
849
850    fn pop_front(&mut self) -> Option<Bytes> {
851        let msg = self.messages.pop_front()?;
852        self.bytes -= msg.len();
853        Some(msg)
854    }
855
856    fn is_empty(&self) -> bool {
857        self.messages.is_empty()
858    }
859
860    const fn is_full(&self) -> bool {
861        self.bytes >= self.max_bytes
862    }
863}
864
865/// Result of polling outbound producers into the mux buffer.
866#[derive(Clone, Copy, Debug, Eq, PartialEq)]
867enum ProducerPoll {
868    /// All outbound producers are pending.
869    Pending,
870    /// The mux buffer reached its soft cap.
871    Full,
872    /// An outbound producer closed.
873    Closed,
874}
875
876#[cfg(test)]
877mod tests {
878    use super::*;
879    use crate::{
880        handshake::EthHandshake,
881        message::MAX_MESSAGE_SIZE,
882        protocol::Protocol,
883        test_utils::{
884            connect_passthrough, eth_handshake, eth_hello,
885            proto::{test_hello, TestProtoMessage},
886        },
887        UnauthedEthStream, UnauthedP2PStream,
888    };
889    use futures::{stream, task::noop_waker_ref};
890    use reth_eth_wire_types::EthNetworkPrimitives;
891    use std::task::Poll;
892    use tokio::{net::TcpListener, sync::oneshot};
893    use tokio_util::codec::Decoder;
894
895    #[derive(Debug)]
896    struct PendingPrimary {
897        _proxy: ProtocolProxy,
898    }
899
900    impl Stream for PendingPrimary {
901        type Item = Result<(), P2PStreamError>;
902
903        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
904            Poll::Pending
905        }
906    }
907
908    #[derive(Debug)]
909    struct StalledTransport;
910
911    impl Stream for StalledTransport {
912        type Item = io::Result<BytesMut>;
913
914        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
915            Poll::Pending
916        }
917    }
918
919    impl Sink<Bytes> for StalledTransport {
920        type Error = io::Error;
921
922        fn poll_ready(
923            self: Pin<&mut Self>,
924            _cx: &mut Context<'_>,
925        ) -> Poll<Result<(), Self::Error>> {
926            Poll::Pending
927        }
928
929        fn start_send(self: Pin<&mut Self>, _item: Bytes) -> Result<(), Self::Error> {
930            Ok(())
931        }
932
933        fn poll_flush(
934            self: Pin<&mut Self>,
935            _cx: &mut Context<'_>,
936        ) -> Poll<Result<(), Self::Error>> {
937            Poll::Pending
938        }
939
940        fn poll_close(
941            self: Pin<&mut Self>,
942            _cx: &mut Context<'_>,
943        ) -> Poll<Result<(), Self::Error>> {
944            Poll::Pending
945        }
946    }
947
948    #[tokio::test]
949    async fn satellite_mux_stops_polling_protocols_when_out_buffer_is_full() {
950        let (hello, _) = test_hello();
951        let shared_capabilities =
952            SharedCapabilities::try_new(hello.protocols.clone(), hello.message().capabilities)
953                .unwrap();
954        let conn = P2PStream::new(StalledTransport, shared_capabilities);
955        let eth = conn.shared_capabilities().eth().unwrap().clone();
956
957        let mut st = RlpxProtocolMultiplexer::new(conn)
958            .into_satellite_stream(eth.capability().as_ref(), |proxy| PendingPrimary {
959                _proxy: proxy,
960            })
961            .unwrap();
962        const MESSAGE_COUNT: usize = 4096;
963        const MESSAGE_BYTES: usize = 1024;
964        st.inner.out_buffer.max_bytes = 4 * MESSAGE_BYTES + 1;
965        st.install_protocol(&TestProtoMessage::capability(), |_conn| {
966            stream::iter((0..MESSAGE_COUNT).map(|_| {
967                let mut msg = BytesMut::zeroed(MESSAGE_BYTES);
968                msg[0] = TestProtoMessage::ping().message_type as u8;
969                msg
970            }))
971        })
972        .unwrap();
973
974        let mut cx = Context::from_waker(noop_waker_ref());
975        assert!(Pin::new(&mut st).poll_next(&mut cx).is_pending());
976
977        assert!(st.inner.out_buffer.bytes > st.inner.out_buffer.max_bytes);
978        assert!(st.inner.out_buffer.bytes <= st.inner.out_buffer.max_bytes + MESSAGE_BYTES);
979        assert!(st.inner.out_buffer.messages.len() < MESSAGE_COUNT);
980    }
981
982    #[tokio::test]
983    async fn satellite_mux_round_robins_ready_protocols_when_out_buffer_fills() {
984        let (mut hello, _) = eth_hello();
985        let cap_a = Capability::new_static("aaa", 1);
986        let cap_b = Capability::new_static("bbb", 1);
987        hello.protocols.push(Protocol::new(cap_a.clone(), 1));
988        hello.protocols.push(Protocol::new(cap_b.clone(), 1));
989
990        let shared_capabilities =
991            SharedCapabilities::try_new(hello.protocols.clone(), hello.message().capabilities)
992                .unwrap();
993        let conn = P2PStream::new(StalledTransport, shared_capabilities);
994        let eth = conn.shared_capabilities().eth().unwrap().clone();
995        let cap_a_offset =
996            conn.shared_capabilities().find(&cap_a).unwrap().relative_message_id_offset();
997        let cap_b_offset =
998            conn.shared_capabilities().find(&cap_b).unwrap().relative_message_id_offset();
999
1000        let mut st = RlpxProtocolMultiplexer::new(conn)
1001            .into_satellite_stream(eth.capability().as_ref(), |proxy| PendingPrimary {
1002                _proxy: proxy,
1003            })
1004            .unwrap();
1005        st.inner.out_buffer.max_bytes = 5;
1006        st.install_protocol(&cap_a, |_conn| {
1007            stream::iter((0..16).map(|_| BytesMut::from(&[0, b'a'][..])))
1008        })
1009        .unwrap();
1010        st.install_protocol(&cap_b, |_conn| {
1011            stream::iter((0..16).map(|_| BytesMut::from(&[0, b'b'][..])))
1012        })
1013        .unwrap();
1014
1015        let mut cx = Context::from_waker(noop_waker_ref());
1016        assert!(Pin::new(&mut st).poll_next(&mut cx).is_pending());
1017
1018        let message_ids =
1019            st.inner.out_buffer.messages.iter().take(2).map(|msg| msg[0]).collect::<Vec<_>>();
1020        assert_eq!(message_ids.len(), 2);
1021        assert_ne!(message_ids[0], message_ids[1]);
1022        assert!(message_ids.contains(&cap_a_offset));
1023        assert!(message_ids.contains(&cap_b_offset));
1024    }
1025
1026    #[tokio::test]
1027    async fn eth_satellite() {
1028        reth_tracing::init_test_tracing();
1029        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1030        let local_addr = listener.local_addr().unwrap();
1031        let (status, fork_filter) = eth_handshake();
1032        let other_status = status;
1033        let other_fork_filter = fork_filter.clone();
1034        let _handle = tokio::spawn(async move {
1035            let (incoming, _) = listener.accept().await.unwrap();
1036            let stream = crate::PassthroughCodec::default().framed(incoming);
1037            let (server_hello, _) = eth_hello();
1038            let (p2p_stream, _) =
1039                UnauthedP2PStream::new(stream).handshake(server_hello).await.unwrap();
1040
1041            let (_eth_stream, _) = UnauthedEthStream::new(p2p_stream)
1042                .handshake::<EthNetworkPrimitives>(other_status, other_fork_filter)
1043                .await
1044                .unwrap();
1045
1046            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
1047        });
1048
1049        let conn = connect_passthrough(local_addr, eth_hello().0).await;
1050        let eth = conn.shared_capabilities().eth().unwrap().clone();
1051
1052        let multiplexer = RlpxProtocolMultiplexer::new(conn);
1053        let _satellite = multiplexer
1054            .into_satellite_stream_with_handshake(eth.capability().as_ref(), async move |proxy| {
1055                UnauthedEthStream::new(proxy)
1056                    .handshake::<EthNetworkPrimitives>(status, fork_filter)
1057                    .await
1058            })
1059            .await
1060            .unwrap();
1061    }
1062
1063    /// A test that install a satellite stream eth+test protocol and sends messages between them.
1064    #[tokio::test(flavor = "multi_thread")]
1065    async fn eth_test_protocol_satellite() {
1066        reth_tracing::init_test_tracing();
1067        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1068        let local_addr = listener.local_addr().unwrap();
1069        let (status, fork_filter) = eth_handshake();
1070        let other_status = status;
1071        let other_fork_filter = fork_filter.clone();
1072        let _handle = tokio::spawn(async move {
1073            let (incoming, _) = listener.accept().await.unwrap();
1074            let stream = crate::PassthroughCodec::default().framed(incoming);
1075            let (server_hello, _) = test_hello();
1076            let (conn, _) = UnauthedP2PStream::new(stream).handshake(server_hello).await.unwrap();
1077
1078            let (mut st, _their_status) = RlpxProtocolMultiplexer::new(conn)
1079                .into_eth_satellite_stream::<EthNetworkPrimitives>(
1080                    other_status,
1081                    other_fork_filter,
1082                    Arc::new(EthHandshake::default()),
1083                    MAX_MESSAGE_SIZE,
1084                )
1085                .await
1086                .unwrap();
1087
1088            st.install_protocol(&TestProtoMessage::capability(), |mut conn| {
1089                async_stream::stream! {
1090                    yield TestProtoMessage::ping().encoded();
1091                    let msg = conn.next().await.unwrap();
1092                    let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
1093                    assert_eq!(msg, TestProtoMessage::pong());
1094
1095                    yield TestProtoMessage::message("hello").encoded();
1096                    let msg = conn.next().await.unwrap();
1097                    let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
1098                    assert_eq!(msg, TestProtoMessage::message("good bye!"));
1099
1100                    yield TestProtoMessage::message("good bye!").encoded();
1101
1102                    futures::future::pending::<()>().await;
1103                    unreachable!()
1104                }
1105            })
1106            .unwrap();
1107
1108            loop {
1109                let _ = st.next().await;
1110            }
1111        });
1112
1113        let conn = connect_passthrough(local_addr, test_hello().0).await;
1114        let (mut st, _their_status) = RlpxProtocolMultiplexer::new(conn)
1115            .into_eth_satellite_stream::<EthNetworkPrimitives>(
1116                status,
1117                fork_filter,
1118                Arc::new(EthHandshake::default()),
1119                MAX_MESSAGE_SIZE,
1120            )
1121            .await
1122            .unwrap();
1123
1124        let (tx, mut rx) = oneshot::channel();
1125
1126        st.install_protocol(&TestProtoMessage::capability(), |mut conn| {
1127            async_stream::stream! {
1128                let msg = conn.next().await.unwrap();
1129                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
1130                assert_eq!(msg, TestProtoMessage::ping());
1131
1132                yield TestProtoMessage::pong().encoded();
1133
1134                let msg = conn.next().await.unwrap();
1135                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
1136                assert_eq!(msg, TestProtoMessage::message("hello"));
1137
1138                yield TestProtoMessage::message("good bye!").encoded();
1139
1140                let msg = conn.next().await.unwrap();
1141                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
1142                assert_eq!(msg, TestProtoMessage::message("good bye!"));
1143
1144                tx.send(()).unwrap();
1145
1146                futures::future::pending::<()>().await;
1147                unreachable!()
1148            }
1149        })
1150        .unwrap();
1151
1152        loop {
1153            tokio::select! {
1154                _ = &mut rx => {
1155                    break
1156                }
1157               _ = st.next() => {
1158                }
1159            }
1160        }
1161    }
1162}
reth_eth_wire/multiplex.rs

reth_eth_wire/
multiplex.rs
