reth_eth_wire/

multiplex.rs

//! Rlpx protocol multiplexer and satellite stream
//!
//! A Satellite is a Stream that primarily drives a single `RLPx` subprotocol but can also handle
//! additional subprotocols.
//!
//! 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)
//! Hence it is expected that the primary protocol is "eth" and the additional protocols are
//! "dependent satellite" protocols.

use std::{
    collections::VecDeque,
    fmt,
    future::Future,
    io,
    pin::{pin, Pin},
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    task::{ready, Context, Poll},
};

use crate::{
    capability::{SharedCapabilities, SharedCapability, UnsupportedCapabilityError},
    errors::{EthStreamError, P2PStreamError},
    handshake::EthRlpxHandshake,
    p2pstream::DisconnectP2P,
    protocol::ProtocolIngressLimits,
    CanDisconnect, Capability, DisconnectReason, EthStream, P2PStream, UnifiedStatus,
    HANDSHAKE_TIMEOUT,
};
use bytes::{Bytes, BytesMut};
use futures::{Sink, SinkExt, Stream, StreamExt, TryStream, TryStreamExt};
use reth_eth_wire_types::NetworkPrimitives;
use reth_ethereum_forks::ForkFilter;
use reth_metrics::metrics::counter;
use tokio::sync::{mpsc, mpsc::UnboundedSender};
use tokio_stream::wrappers::{ReceiverStream, UnboundedReceiverStream};

/// A Stream and Sink type that wraps a raw rlpx stream [`P2PStream`] and handles message ID
/// multiplexing.
#[derive(Debug)]
pub struct RlpxProtocolMultiplexer<St> {
    inner: MultiplexInner<St>,
}

impl<St> RlpxProtocolMultiplexer<St> {
    /// Wraps the raw p2p stream
    pub fn new(conn: P2PStream<St>) -> Self {
        Self {
            inner: MultiplexInner {
                conn,
                protocols: Default::default(),
                out_buffer: Default::default(),
                inbound_budget: Arc::new(InboundBudget::new(MAX_MUX_IN_BUFFER_BYTES)),
            },
        }
    }

    /// Installs a new protocol on top of the raw p2p stream.
    ///
    /// This accepts a closure that receives a [`ProtocolConnection`] that will yield messages for
    /// the given capability.
    pub fn install_protocol<F, Proto>(
        &mut self,
        cap: &Capability,
        f: F,
    ) -> Result<(), UnsupportedCapabilityError>
    where
        F: FnOnce(ProtocolConnection) -> Proto,
        Proto: Stream<Item = BytesMut> + Send + 'static,
    {
        self.inner.install_protocol(cap, ProtocolIngressLimits::default(), f)
    }

    /// Installs a new protocol with local limits for inbound messages.
    pub fn install_protocol_with_limits<F, Proto>(
        &mut self,
        cap: &Capability,
        limits: ProtocolIngressLimits,
        f: F,
    ) -> Result<(), UnsupportedCapabilityError>
    where
        F: FnOnce(ProtocolConnection) -> Proto,
        Proto: Stream<Item = BytesMut> + Send + 'static,
    {
        self.inner.install_protocol(cap, limits, f)
    }

    /// Returns the [`SharedCapabilities`] of the underlying raw p2p stream
    pub const fn shared_capabilities(&self) -> &SharedCapabilities {
        self.inner.shared_capabilities()
    }

    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
    pub fn into_satellite_stream<F, Primary>(
        self,
        cap: &Capability,
        primary: F,
    ) -> Result<RlpxSatelliteStream<St, Primary>, P2PStreamError>
    where
        F: FnOnce(ProtocolProxy) -> Primary,
    {
        let Ok(shared_cap) = self.shared_capabilities().ensure_matching_capability(cap).cloned()
        else {
            return Err(P2PStreamError::CapabilityNotShared)
        };

        let (to_primary, from_wire) =
            self.inner.inbound_channel(&shared_cap, ProtocolIngressLimits::default());
        let (to_wire, from_primary) = mpsc::unbounded_channel();
        let proxy = ProtocolProxy { shared_cap: shared_cap.clone(), from_wire, to_wire };

        let st = primary(proxy);
        Ok(RlpxSatelliteStream {
            inner: self.inner,
            primary: PrimaryProtocol {
                to_primary,
                from_primary: UnboundedReceiverStream::new(from_primary),
                st,
                shared_cap,
            },
            next_outbound: 0,
            producer_polls_since_inbound: 0,
        })
    }

    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
    ///
    /// Returns an error if the primary protocol is not supported by the remote or the handshake
    /// failed.
    pub async fn into_satellite_stream_with_handshake<F, Fut, Err, Primary>(
        self,
        cap: &Capability,
        handshake: F,
    ) -> Result<RlpxSatelliteStream<St, Primary>, Err>
    where
        F: FnOnce(ProtocolProxy) -> Fut,
        Fut: Future<Output = Result<Primary, Err>>,
        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
        P2PStreamError: Into<Err>,
    {
        self.into_satellite_stream_with_tuple_handshake(cap, async move |proxy| {
            let st = handshake(proxy).await?;
            Ok((st, ()))
        })
        .await
        .map(|(st, _)| st)
    }

    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with the given primary protocol.
    ///
    /// Returns an error if the primary protocol is not supported by the remote or the handshake
    /// failed.
    ///
    /// This accepts a closure that does a handshake with the remote peer and returns a tuple of the
    /// primary stream and extra data.
    ///
    /// See also [`UnauthedEthStream::handshake`](crate::UnauthedEthStream)
    pub async fn into_satellite_stream_with_tuple_handshake<F, Fut, Err, Primary, Extra>(
        mut self,
        cap: &Capability,
        handshake: F,
    ) -> Result<(RlpxSatelliteStream<St, Primary>, Extra), Err>
    where
        F: FnOnce(ProtocolProxy) -> Fut,
        Fut: Future<Output = Result<(Primary, Extra), Err>>,
        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
        P2PStreamError: Into<Err>,
    {
        let Ok(shared_cap) = self.shared_capabilities().ensure_matching_capability(cap).cloned()
        else {
            return Err(P2PStreamError::CapabilityNotShared.into())
        };

        let (to_primary, from_wire) =
            self.inner.inbound_channel(&shared_cap, ProtocolIngressLimits::default());
        let (to_wire, mut from_primary) = mpsc::unbounded_channel();
        let proxy = ProtocolProxy { shared_cap: shared_cap.clone(), from_wire, to_wire };

        let f = handshake(proxy);
        let mut f = pin!(f);
        let mut inbound_messages = 0;
        let mut handshake_satellite_bytes = 0usize;
        let mut handshake_satellite_messages = 0usize;

        // this polls the connection and the primary stream concurrently until the handshake is
        // complete
        loop {
            tokio::select! {
                biased;
                res = &mut f => {
                    let (st, extra) = res?;
                    return Ok((
                        RlpxSatelliteStream {
                            inner: self.inner,
                            primary: PrimaryProtocol {
                                to_primary,
                                from_primary: UnboundedReceiverStream::new(from_primary),
                                st,
                                shared_cap,
                            },
                            next_outbound: 0,
                            producer_polls_since_inbound: 0,
                        },
                        extra,
                    ))
                }
                Some(msg) = from_primary.recv() => {
                    self.inner.conn.send(msg).await.map_err(Into::into)?;
                }
                // Polling a satellite drives its inbound consumer even when it has no outbound
                // message. This must happen before reading another frame from the socket.
                msg = ProtocolsPoller::new(&mut self.inner.protocols) => {
                     self.inner.conn.send(msg.map_err(Into::into)?).await.map_err(Into::into)?;
                }
                incoming = self.inner.conn.next() => {
                    let msg = match incoming {
                        Some(Ok(msg)) => msg,
                        Some(Err(err)) => return Err(err.into()),
                        None => {
                            return Err(P2PStreamError::Io(io::ErrorKind::UnexpectedEof.into()).into())
                        }
                    };
                    // Ensure the message belongs to the primary protocol
                    let Some(offset) = msg.first().copied()
                    else {
                        return Err(P2PStreamError::EmptyProtocolMessage.into())
                    };
                    if let Some(cap) = self.shared_capabilities().find_by_relative_offset(offset).cloned() {
                            if cap == shared_cap {
                                // delegate to primary
                                to_primary.try_send(msg).map_err(Into::into)?;
                            } else {
                                // Satellite consumers may wait for authentication, so keep their
                                // aggregate pre-handshake retention much smaller than active queues.
                                handshake_satellite_bytes = handshake_satellite_bytes
                                    .saturating_add(msg.capacity().max(msg.len()));
                                handshake_satellite_messages += 1;
                                if handshake_satellite_bytes > MAX_HANDSHAKE_SATELLITE_BYTES ||
                                    handshake_satellite_messages > MAX_HANDSHAKE_SATELLITE_MESSAGES
                                {
                                    counter!("p2pstream.subprotocol_inbound_buffer_full").increment(1);
                                    return Err(P2PStreamError::SubprotocolInboundBufferFull {
                                        capability: cap.capability().into_owned(),
                                    }.into())
                                }
                                self.inner.delegate_message(&cap, msg).map_err(Into::into)?;
                            }
                        } else {
                           return Err(P2PStreamError::UnknownReservedMessageId(offset).into())
                        }

                    inbound_messages += 1;
                    if inbound_messages == MAX_INBOUND_MESSAGES_PER_POLL {
                        inbound_messages = 0;
                        yield_once().await;
                    }
                }
            }
        }
    }

    /// Converts this multiplexer into a [`RlpxSatelliteStream`] with eth protocol as the given
    /// primary protocol and the handshake implementation.
    pub async fn into_eth_satellite_stream<N: NetworkPrimitives>(
        self,
        status: UnifiedStatus,
        fork_filter: ForkFilter,
        handshake: Arc<dyn EthRlpxHandshake>,
        eth_max_message_size: usize,
    ) -> Result<(RlpxSatelliteStream<St, EthStream<ProtocolProxy, N>>, UnifiedStatus), EthStreamError>
    where
        St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
    {
        let eth_cap = self.inner.conn.shared_capabilities().eth_version()?;
        self.into_satellite_stream_with_tuple_handshake(
            &Capability::eth(eth_cap),
            async move |proxy| {
                let handshake = handshake.clone();
                let mut unauth = UnauthProxy { inner: proxy };
                let their_status = handshake
                    .handshake(&mut unauth, status, fork_filter, HANDSHAKE_TIMEOUT)
                    .await?;
                let eth_stream = EthStream::with_max_message_size(
                    eth_cap,
                    unauth.into_inner(),
                    eth_max_message_size,
                );
                Ok((eth_stream, their_status))
            },
        )
        .await
    }
}

#[derive(Debug)]
struct MultiplexInner<St> {
    /// The raw p2p stream
    conn: P2PStream<St>,
    /// All the subprotocols that are multiplexed on top of the raw p2p stream
    protocols: VecDeque<ProtocolStream>,
    /// Buffer for outgoing messages on the wire.
    out_buffer: OutBuffer,
    /// Byte budget shared by every inbound protocol queue on this connection.
    inbound_budget: Arc<InboundBudget>,
}

impl<St> MultiplexInner<St> {
    const fn shared_capabilities(&self) -> &SharedCapabilities {
        self.conn.shared_capabilities()
    }

    /// Delegates a message to the matching protocol.
    fn delegate_message(
        &self,
        cap: &SharedCapability,
        msg: BytesMut,
    ) -> Result<bool, P2PStreamError> {
        for proto in &self.protocols {
            if proto.shared_cap == *cap {
                proto.send_raw(msg)?;
                return Ok(true)
            }
        }
        Ok(false)
    }

    fn install_protocol<F, Proto>(
        &mut self,
        cap: &Capability,
        limits: ProtocolIngressLimits,
        f: F,
    ) -> Result<(), UnsupportedCapabilityError>
    where
        F: FnOnce(ProtocolConnection) -> Proto,
        Proto: Stream<Item = BytesMut> + Send + 'static,
    {
        let shared_cap =
            self.conn.shared_capabilities().ensure_matching_capability(cap).cloned()?;
        self.conn.set_protocol_ingress_limits(&shared_cap, limits);
        let (to_satellite, from_wire) = self.inbound_channel(&shared_cap, limits);
        let proto_conn = ProtocolConnection { from_wire };
        let st = f(proto_conn);
        let st = ProtocolStream { shared_cap, to_satellite, satellite_st: Box::pin(st) };
        self.protocols.push_back(st);
        Ok(())
    }

    fn inbound_channel(
        &self,
        capability: &SharedCapability,
        limits: ProtocolIngressLimits,
    ) -> (InboundSender, InboundReceiver) {
        inbound_channel(capability, limits, Arc::clone(&self.inbound_budget))
    }

    /// Hands buffered messages to the connection without flushing it.
    ///
    /// Returns `Ready` once the buffer is empty and the connection can accept another message.
    fn poll_send_buffered(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), P2PStreamError>>
    where
        St: Sink<Bytes, Error = io::Error> + Unpin,
    {
        loop {
            if let Err(err) = ready!(self.conn.poll_ready_unpin(cx)) {
                self.conn.start_disconnect(DisconnectReason::DisconnectRequested)?;
                return Poll::Ready(Err(err))
            }
            let Some(msg) = self.out_buffer.pop_front() else { return Poll::Ready(Ok(())) };
            self.conn.start_send_unpin(msg)?;
        }
    }
}

/// Represents a protocol in the multiplexer that is used as the primary protocol.
#[derive(Debug)]
struct PrimaryProtocol<Primary> {
    /// Channel to send messages to the primary protocol.
    to_primary: InboundSender,
    /// Receiver for messages from the primary protocol.
    from_primary: UnboundedReceiverStream<Bytes>,
    /// Shared capability of the primary protocol.
    shared_cap: SharedCapability,
    /// The primary stream.
    st: Primary,
}

/// A Stream and Sink type that acts as a wrapper around a primary `RLPx` subprotocol (e.g. "eth")
///
/// Only emits and sends _non-empty_ messages
#[derive(Debug)]
pub struct ProtocolProxy {
    shared_cap: SharedCapability,
    /// Receives _non-empty_ messages from the wire
    from_wire: InboundReceiver,
    /// Sends _non-empty_ messages from the wire
    to_wire: UnboundedSender<Bytes>,
}

impl ProtocolProxy {
    /// Sends a _non-empty_ message on the wire.
    fn try_send(&self, msg: Bytes) -> Result<(), io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }
        self.to_wire.send(self.mask_msg_id(msg)?).map_err(|_| io::ErrorKind::BrokenPipe.into())
    }

    /// Masks the message ID of a message to be sent on the wire.
    #[inline]
    fn mask_msg_id(&self, msg: Bytes) -> Result<Bytes, io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }

        let offset = self.shared_cap.relative_message_id_offset();
        if offset == 0 {
            return Ok(msg);
        }

        let mut masked: BytesMut = msg.into();
        masked[0] = masked[0].checked_add(offset).ok_or(io::ErrorKind::InvalidInput)?;
        Ok(masked.freeze())
    }

    /// Unmasks the message ID of a message received from the wire.
    #[inline]
    fn unmask_id(&self, mut msg: BytesMut) -> Result<BytesMut, io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }
        msg[0] = msg[0]
            .checked_sub(self.shared_cap.relative_message_id_offset())
            .ok_or(io::ErrorKind::InvalidInput)?;
        Ok(msg)
    }
}

impl Stream for ProtocolProxy {
    type Item = Result<BytesMut, io::Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let msg = ready!(self.from_wire.poll_next_unpin(cx));
        Poll::Ready(msg.map(|msg| self.get_mut().unmask_id(msg)))
    }
}

impl Sink<Bytes> for ProtocolProxy {
    type Error = io::Error;

    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
        self.get_mut().try_send(item)
    }

    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }
}

impl CanDisconnect<Bytes> for ProtocolProxy {
    fn disconnect(
        &mut self,
        _reason: DisconnectReason,
    ) -> Pin<Box<dyn Future<Output = Result<(), <Self as Sink<Bytes>>::Error>> + Send + '_>> {
        Box::pin(async move { Ok(()) })
    }
}

/// Adapter so the injected `EthRlpxHandshake` can run over a multiplexed `ProtocolProxy`
/// using the same error type expectations (`P2PStreamError`).
#[derive(Debug)]
struct UnauthProxy {
    inner: ProtocolProxy,
}

impl UnauthProxy {
    fn into_inner(self) -> ProtocolProxy {
        self.inner
    }
}

impl Stream for UnauthProxy {
    type Item = Result<BytesMut, P2PStreamError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.poll_next_unpin(cx).map(|opt| opt.map(|res| res.map_err(P2PStreamError::from)))
    }
}

impl Sink<Bytes> for UnauthProxy {
    type Error = P2PStreamError;

    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready_unpin(cx).map_err(P2PStreamError::from)
    }

    fn start_send(mut self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
        self.inner.start_send_unpin(item).map_err(P2PStreamError::from)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_flush_unpin(cx).map_err(P2PStreamError::from)
    }

    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_close_unpin(cx).map_err(P2PStreamError::from)
    }
}

impl CanDisconnect<Bytes> for UnauthProxy {
    fn disconnect(
        &mut self,
        reason: DisconnectReason,
    ) -> Pin<Box<dyn Future<Output = Result<(), <Self as Sink<Bytes>>::Error>> + Send + '_>> {
        let fut = self.inner.disconnect(reason);
        Box::pin(async move { fut.await.map_err(P2PStreamError::from) })
    }
}

#[derive(Debug)]
struct InboundSender {
    inner: mpsc::Sender<BudgetedInbound>,
    connection_budget: Arc<InboundBudget>,
    protocol_budget: Arc<InboundBudget>,
    capability: Capability,
}

impl InboundSender {
    fn try_send(&self, frame: BytesMut) -> Result<(), P2PStreamError> {
        let size = frame.capacity().max(frame.len());
        let Some(connection_guard) = InboundBudget::try_reserve(&self.connection_budget, size)
        else {
            return Err(self.buffer_full())
        };
        let Some(protocol_guard) = InboundBudget::try_reserve(&self.protocol_budget, size) else {
            return Err(self.buffer_full())
        };

        let frame = BudgetedInbound {
            frame,
            _connection_guard: connection_guard,
            _protocol_guard: protocol_guard,
        };
        self.inner.try_send(frame).map_err(|err| match err {
            mpsc::error::TrySendError::Full(_) => self.buffer_full(),
            mpsc::error::TrySendError::Closed(_) => {
                P2PStreamError::Io(io::ErrorKind::BrokenPipe.into())
            }
        })
    }

    fn buffer_full(&self) -> P2PStreamError {
        counter!("p2pstream.subprotocol_inbound_buffer_full").increment(1);
        P2PStreamError::SubprotocolInboundBufferFull { capability: self.capability.clone() }
    }
}

#[derive(Debug)]
struct InboundReceiver {
    inner: ReceiverStream<BudgetedInbound>,
}

impl Stream for InboundReceiver {
    type Item = BytesMut;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.poll_next_unpin(cx).map(|frame| frame.map(BudgetedInbound::into_frame))
    }
}

#[derive(Debug)]
struct BudgetedInbound {
    frame: BytesMut,
    _connection_guard: InboundBudgetGuard,
    _protocol_guard: InboundBudgetGuard,
}

impl BudgetedInbound {
    fn into_frame(self) -> BytesMut {
        let Self { frame, _connection_guard, _protocol_guard } = self;
        frame
    }
}

#[derive(Debug)]
struct InboundBudget {
    used: AtomicUsize,
    max: usize,
}

impl InboundBudget {
    const fn new(max: usize) -> Self {
        Self { used: AtomicUsize::new(0), max }
    }

    fn try_reserve(budget: &Arc<Self>, size: usize) -> Option<InboundBudgetGuard> {
        budget
            .used
            .try_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
                used.checked_add(size).filter(|next| *next <= budget.max)
            })
            .ok()
            .map(|_| InboundBudgetGuard { size, budget: Arc::clone(budget) })
    }
}

#[derive(Debug)]
struct InboundBudgetGuard {
    size: usize,
    budget: Arc<InboundBudget>,
}

impl Drop for InboundBudgetGuard {
    fn drop(&mut self) {
        self.budget.used.fetch_sub(self.size, Ordering::Relaxed);
    }
}

fn inbound_channel(
    capability: &SharedCapability,
    limits: ProtocolIngressLimits,
    connection_budget: Arc<InboundBudget>,
) -> (InboundSender, InboundReceiver) {
    let (tx, rx) = mpsc::channel(limits.max_buffered_messages());
    let sender = InboundSender {
        inner: tx,
        connection_budget,
        protocol_budget: Arc::new(InboundBudget::new(limits.max_buffered_bytes())),
        capability: capability.capability().into_owned(),
    };
    let receiver = InboundReceiver { inner: ReceiverStream::new(rx) };
    (sender, receiver)
}

/// A connection channel to receive _`non_empty`_ messages for the negotiated protocol.
///
/// This is a [Stream] that returns raw bytes of the received messages for this protocol.
#[derive(Debug)]
pub struct ProtocolConnection {
    from_wire: InboundReceiver,
}

impl Stream for ProtocolConnection {
    type Item = BytesMut;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.from_wire.poll_next_unpin(cx)
    }
}

/// A Stream and Sink type that acts as a wrapper around a primary `RLPx` subprotocol (e.g. "eth")
/// [`EthStream`] and can also handle additional subprotocols.
#[derive(Debug)]
pub struct RlpxSatelliteStream<St, Primary> {
    inner: MultiplexInner<St>,
    primary: PrimaryProtocol<Primary>,
    /// Round-robin cursor for the next outbound producer to poll.
    next_outbound: usize,
    /// Number of round-robin producers polled since the last inbound frame.
    producer_polls_since_inbound: usize,
}

impl<St, Primary> RlpxSatelliteStream<St, Primary> {
    /// Installs a new protocol on top of the raw p2p stream.
    ///
    /// This accepts a closure that receives a [`ProtocolConnection`] that will yield messages for
    /// the given capability.
    pub fn install_protocol<F, Proto>(
        &mut self,
        cap: &Capability,
        f: F,
    ) -> Result<(), UnsupportedCapabilityError>
    where
        F: FnOnce(ProtocolConnection) -> Proto,
        Proto: Stream<Item = BytesMut> + Send + 'static,
    {
        let result = self.inner.install_protocol(cap, ProtocolIngressLimits::default(), f);
        if result.is_ok() {
            self.producer_polls_since_inbound = 0;
        }
        result
    }

    /// Installs a new protocol with local limits for inbound messages.
    pub fn install_protocol_with_limits<F, Proto>(
        &mut self,
        cap: &Capability,
        limits: ProtocolIngressLimits,
        f: F,
    ) -> Result<(), UnsupportedCapabilityError>
    where
        F: FnOnce(ProtocolConnection) -> Proto,
        Proto: Stream<Item = BytesMut> + Send + 'static,
    {
        let result = self.inner.install_protocol(cap, limits, f);
        if result.is_ok() {
            self.producer_polls_since_inbound = 0;
        }
        result
    }

    /// Returns the primary protocol.
    #[inline]
    pub const fn primary(&self) -> &Primary {
        &self.primary.st
    }

    /// Returns mutable access to the primary protocol.
    #[inline]
    pub const fn primary_mut(&mut self) -> &mut Primary {
        &mut self.primary.st
    }

    /// Returns the underlying [`P2PStream`].
    #[inline]
    pub const fn inner(&self) -> &P2PStream<St> {
        &self.inner.conn
    }

    /// Returns mutable access to the underlying [`P2PStream`].
    #[inline]
    pub const fn inner_mut(&mut self) -> &mut P2PStream<St> {
        &mut self.inner.conn
    }

    /// Consumes this type and returns the wrapped [`P2PStream`].
    #[inline]
    pub fn into_inner(self) -> P2PStream<St> {
        self.inner.conn
    }

    /// Polls primary and satellite outbound producers round-robin until the `OutBuffer` is full or
    /// every producer is pending.
    ///
    /// The cursor advances after each producer poll, so a ready producer cannot drain repeatedly
    /// before later producers get a turn.
    fn poll_outbound_producers(&mut self, cx: &mut Context<'_>) -> Result<ProducerPoll, io::Error> {
        let producers = self.inner.protocols.len() + 1;
        let mut pending = 0;

        while pending < producers {
            if self.inner.out_buffer.is_full() {
                return Ok(ProducerPoll::Full)
            }

            if self.next_outbound >= producers {
                self.next_outbound = 0;
            }

            let producer = self.next_outbound;
            self.next_outbound = (self.next_outbound + 1) % producers;
            if self.producer_polls_since_inbound < producers {
                self.producer_polls_since_inbound += 1;
            }

            let msg = if producer == 0 {
                match self.primary.from_primary.poll_next_unpin(cx) {
                    Poll::Ready(Some(msg)) => msg,
                    Poll::Ready(None) => return Ok(ProducerPoll::Closed),
                    Poll::Pending => {
                        pending += 1;
                        continue
                    }
                }
            } else {
                let proto = self
                    .inner
                    .protocols
                    .get_mut(producer - 1)
                    .expect("outbound producer index checked against protocol count");
                match proto.poll_next_unpin(cx) {
                    Poll::Ready(Some(msg)) => msg?,
                    Poll::Ready(None) => return Ok(ProducerPoll::Closed),
                    Poll::Pending => {
                        pending += 1;
                        continue
                    }
                }
            };

            pending = 0;
            self.inner.out_buffer.push_back(msg);
        }

        Ok(ProducerPoll::Pending)
    }

    /// Moves messages the primary protocol already accepted into the connection.
    ///
    /// [`ProtocolProxy`] accepts outgoing messages into an unbounded channel, so the sink must not
    /// report readiness before that channel is drained. Otherwise a caller that sends while the
    /// sink is ready queues messages without limit when the remote reads slowly.
    fn poll_drain_primary(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), P2PStreamError>>
    where
        St: Sink<Bytes, Error = io::Error> + Unpin,
    {
        loop {
            ready!(self.inner.poll_send_buffered(cx))?;
            // Both errors, empty and disconnected, mean no accepted message is left.
            let Ok(msg) = self.primary.from_primary.as_mut().try_recv() else {
                return Poll::Ready(Ok(()))
            };
            self.inner.out_buffer.push_back(msg);
        }
    }
}

impl<St, Primary, PrimaryErr> Stream for RlpxSatelliteStream<St, Primary>
where
    St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
    Primary: TryStream<Error = PrimaryErr> + Unpin,
    P2PStreamError: Into<PrimaryErr>,
{
    type Item = Result<Primary::Ok, Primary::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        let mut inbound_messages = 0;

        loop {
            // first drain the primary stream
            if let Poll::Ready(Some(msg)) = this.primary.st.try_poll_next_unpin(cx) {
                return Poll::Ready(Some(msg))
            }

            let mut conn_ready = match this.inner.poll_send_buffered(cx) {
                Poll::Ready(Ok(())) => true,
                Poll::Ready(Err(err)) => return Poll::Ready(Some(Err(err.into()))),
                Poll::Pending => false,
            };
            // The connection only buffers frames on `start_send`; `poll_flush` performs the
            // actual writes and flushes the transport once for the batch handed to it above.
            // This also resumes a flush that returned pending on an earlier pass; a no-op if
            // nothing is buffered.
            match this.inner.conn.poll_flush_unpin(cx) {
                Poll::Ready(Ok(())) => {}
                Poll::Ready(Err(err)) => return Poll::Ready(Some(Err(err.into()))),
                Poll::Pending => {
                    conn_ready = false;
                }
            }

            match this.poll_outbound_producers(cx) {
                Ok(ProducerPoll::Full)
                    if this.producer_polls_since_inbound < this.inner.protocols.len() + 1 =>
                {
                    // A full buffer prevents protocol streams from being polled. Pause ingress
                    // until every consumer has had a chance to drain the previous frame.
                    if conn_ready {
                        cx.waker().wake_by_ref();
                    }
                    return Poll::Pending
                }
                Ok(ProducerPoll::Pending | ProducerPoll::Full) => {}
                Ok(ProducerPoll::Closed) => return Poll::Ready(None),
                Err(err) => return Poll::Ready(Some(Err(P2PStreamError::Io(err).into()))),
            }

            let mut delegated = false;
            // Pull one message before returning to the top of the loop, where every protocol gets
            // another chance to consume its inbound queue.
            match this.inner.conn.poll_next_unpin(cx) {
                Poll::Ready(Some(Ok(msg))) => {
                    delegated = true;
                    let Some(offset) = msg.first().copied() else {
                        return Poll::Ready(Some(Err(P2PStreamError::EmptyProtocolMessage.into())))
                    };
                    let Some(cap) = this
                        .inner
                        .conn
                        .shared_capabilities()
                        .find_by_relative_offset(offset)
                        .cloned()
                    else {
                        return Poll::Ready(Some(Err(P2PStreamError::UnknownReservedMessageId(
                            offset,
                        )
                        .into())))
                    };

                    let result = if cap == this.primary.shared_cap {
                        this.primary.to_primary.try_send(msg)
                    } else {
                        this.inner.delegate_message(&cap, msg).map(|_| ())
                    };
                    if let Err(err) = result {
                        return Poll::Ready(Some(Err(err.into())))
                    }
                    this.producer_polls_since_inbound = 0;

                    inbound_messages += 1;
                    if inbound_messages == MAX_INBOUND_MESSAGES_PER_POLL {
                        cx.waker().wake_by_ref();
                        return Poll::Pending
                    }
                }
                Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err.into()))),
                Poll::Ready(None) => return Poll::Ready(None),
                Poll::Pending => {}
            }

            if !conn_ready || (!delegated && this.inner.out_buffer.is_empty()) {
                return Poll::Pending
            }
        }
    }
}

impl<St, Primary, T> Sink<T> for RlpxSatelliteStream<St, Primary>
where
    St: Stream<Item = io::Result<BytesMut>> + Sink<Bytes, Error = io::Error> + Unpin,
    Primary: Sink<T> + Unpin,
    P2PStreamError: Into<<Primary as Sink<T>>::Error>,
{
    type Error = <Primary as Sink<T>>::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let this = self.get_mut();
        if let Err(err) = ready!(this.poll_drain_primary(cx)) {
            return Poll::Ready(Err(err.into()))
        }
        if let Err(err) = ready!(this.primary.st.poll_ready_unpin(cx)) {
            return Poll::Ready(Err(err))
        }
        Poll::Ready(Ok(()))
    }

    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
        self.get_mut().primary.st.start_send_unpin(item)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        let this = self.get_mut();
        if let Err(err) = ready!(this.poll_drain_primary(cx)) {
            return Poll::Ready(Err(err.into()))
        }
        this.inner.conn.poll_flush_unpin(cx).map_err(Into::into)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.get_mut().inner.conn.poll_close_unpin(cx).map_err(Into::into)
    }
}

/// Wraps a `RLPx` subprotocol and handles message ID multiplexing.
struct ProtocolStream {
    shared_cap: SharedCapability,
    /// the channel shared with the satellite stream
    to_satellite: InboundSender,
    satellite_st: Pin<Box<dyn Stream<Item = BytesMut> + Send>>,
}

impl ProtocolStream {
    /// Masks the message ID of a message to be sent on the wire.
    #[inline]
    fn mask_msg_id(&self, mut msg: BytesMut) -> Result<Bytes, io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }
        msg[0] = msg[0]
            .checked_add(self.shared_cap.relative_message_id_offset())
            .ok_or(io::ErrorKind::InvalidInput)?;
        Ok(msg.freeze())
    }

    /// Unmasks the message ID of a message received from the wire.
    #[inline]
    fn unmask_id(&self, mut msg: BytesMut) -> Result<BytesMut, io::Error> {
        if msg.is_empty() {
            // message must not be empty
            return Err(io::ErrorKind::InvalidInput.into())
        }
        msg[0] = msg[0]
            .checked_sub(self.shared_cap.relative_message_id_offset())
            .ok_or(io::ErrorKind::InvalidInput)?;
        Ok(msg)
    }

    /// Sends the message to the satellite stream.
    fn send_raw(&self, msg: BytesMut) -> Result<(), P2PStreamError> {
        let msg = self.unmask_id(msg)?;
        self.to_satellite.try_send(msg)
    }
}

impl Stream for ProtocolStream {
    type Item = Result<Bytes, io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        let msg = ready!(this.satellite_st.as_mut().poll_next(cx));
        Poll::Ready(msg.map(|msg| this.mask_msg_id(msg)))
    }
}

impl fmt::Debug for ProtocolStream {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ProtocolStream").field("cap", &self.shared_cap).finish_non_exhaustive()
    }
}

/// Helper to poll multiple protocol streams in a `tokio::select`! branch
struct ProtocolsPoller<'a> {
    protocols: &'a mut VecDeque<ProtocolStream>,
}

impl<'a> ProtocolsPoller<'a> {
    const fn new(protocols: &'a mut VecDeque<ProtocolStream>) -> Self {
        Self { protocols }
    }
}

impl<'a> Future for ProtocolsPoller<'a> {
    type Output = Result<Bytes, P2PStreamError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let protocols = self.protocols.len();
        for _ in 0..protocols {
            let mut proto = self.protocols.pop_front().expect("protocol count checked");
            match proto.poll_next_unpin(cx) {
                Poll::Ready(Some(Err(err))) => {
                    self.protocols.push_back(proto);
                    return Poll::Ready(Err(P2PStreamError::from(err)))
                }
                Poll::Ready(Some(Ok(msg))) => {
                    // Got a message, put protocol back and return the message
                    self.protocols.push_back(proto);
                    return Poll::Ready(Ok(msg));
                }
                _ => {
                    // push it back because we still want to complete the handshake first
                    self.protocols.push_back(proto);
                }
            }
        }

        // All protocols processed, nothing ready
        Poll::Pending
    }
}

fn yield_once() -> impl Future<Output = ()> {
    let mut yielded = false;
    std::future::poll_fn(move |cx| {
        if yielded {
            Poll::Ready(())
        } else {
            yielded = true;
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    })
}

/// Aggregate byte budget for inbound messages queued across all protocols on one connection.
const MAX_MUX_IN_BUFFER_BYTES: usize = 32 * 1024 * 1024;

/// Aggregate satellite traffic accepted while the primary handshake is pending. These limits
/// bound both payload retention and queue metadata before application consumers may start.
const MAX_HANDSHAKE_SATELLITE_BYTES: usize = 256 * 1024;
const MAX_HANDSHAKE_SATELLITE_MESSAGES: usize = 64;

/// Maximum network messages delegated before yielding the connection task.
const MAX_INBOUND_MESSAGES_PER_POLL: usize = 32;

/// Soft cap for per-connection outbound `RLPx` messages waiting in the multiplexer.
///
/// The cap is soft because the next message size is only known after polling a protocol stream.
/// The buffer may exceed this by at most one message before producer polling is paused.
///
/// The lower [`P2PStream`] sink admits two outbound messages and rejects uncompressed payloads
/// above 16 MiB, so 32 MiB mirrors the largest payload volume the lower p2p layer is already
/// prepared to buffer.
const MAX_MUX_OUT_BUFFER_BYTES: usize = 32 * 1024 * 1024;

#[derive(Debug)]
struct OutBuffer {
    messages: VecDeque<Bytes>,
    bytes: usize,
    max_bytes: usize,
}

impl Default for OutBuffer {
    fn default() -> Self {
        Self { messages: Default::default(), bytes: 0, max_bytes: MAX_MUX_OUT_BUFFER_BYTES }
    }
}

impl OutBuffer {
    fn push_back(&mut self, msg: Bytes) {
        self.bytes += msg.len();
        self.messages.push_back(msg);
    }

    fn pop_front(&mut self) -> Option<Bytes> {
        let msg = self.messages.pop_front()?;
        self.bytes -= msg.len();
        Some(msg)
    }

    fn is_empty(&self) -> bool {
        self.messages.is_empty()
    }

    const fn is_full(&self) -> bool {
        self.bytes >= self.max_bytes
    }
}

/// Result of polling outbound producers into the mux buffer.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ProducerPoll {
    /// All outbound producers are pending.
    Pending,
    /// The mux buffer reached its soft cap.
    Full,
    /// An outbound producer closed.
    Closed,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        handshake::EthHandshake,
        message::MAX_MESSAGE_SIZE,
        protocol::Protocol,
        test_utils::{
            connect_passthrough, eth_handshake, eth_hello,
            proto::{test_hello, TestProtoMessage},
        },
        UnauthedEthStream, UnauthedP2PStream,
    };
    use futures::{stream, task::noop_waker_ref};
    use reth_eth_wire_types::EthNetworkPrimitives;
    use std::{sync::atomic::AtomicBool, task::Poll};
    use tokio::{net::TcpListener, sync::oneshot};
    use tokio_util::codec::Decoder;

    #[derive(Debug)]
    struct PendingPrimary {
        _proxy: ProtocolProxy,
    }

    impl Stream for PendingPrimary {
        type Item = Result<(), P2PStreamError>;

        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Poll::Pending
        }
    }

    #[derive(Debug)]
    struct BurstingPrimary {
        proxy: ProtocolProxy,
        remaining: usize,
    }

    impl Stream for BurstingPrimary {
        type Item = Result<(), P2PStreamError>;

        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            if self.remaining > 0 {
                self.remaining -= 1;
                self.proxy.try_send(Bytes::from_static(&[0, 0])).unwrap();
            }
            Poll::Pending
        }
    }

    #[derive(Debug)]
    struct StalledTransport;

    impl Stream for StalledTransport {
        type Item = io::Result<BytesMut>;

        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Poll::Pending
        }
    }

    impl Sink<Bytes> for StalledTransport {
        type Error = io::Error;

        fn poll_ready(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Pending
        }

        fn start_send(self: Pin<&mut Self>, _item: Bytes) -> Result<(), Self::Error> {
            Ok(())
        }

        fn poll_flush(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Pending
        }

        fn poll_close(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            Poll::Pending
        }
    }

    #[derive(Debug)]
    struct InboundFramesTransport {
        frames: VecDeque<BytesMut>,
        writable: AtomicBool,
    }

    impl Stream for InboundFramesTransport {
        type Item = io::Result<BytesMut>;

        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Poll::Ready(self.frames.pop_front().map(Ok))
        }
    }

    impl Sink<Bytes> for InboundFramesTransport {
        type Error = io::Error;

        fn poll_ready(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), Self::Error>> {
            if self.writable.load(Ordering::Relaxed) {
                Poll::Ready(Ok(()))
            } else {
                Poll::Pending
            }
        }

        fn start_send(self: Pin<&mut Self>, _item: Bytes) -> Result<(), Self::Error> {
            Ok(())
        }

        fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.poll_ready(cx)
        }

        fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            self.poll_ready(cx)
        }
    }

    #[derive(Debug)]
    struct HoldingProtocol {
        _conn: ProtocolConnection,
    }

    impl Stream for HoldingProtocol {
        type Item = BytesMut;

        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Poll::Pending
        }
    }

    #[derive(Debug)]
    struct DrainingProtocol {
        conn: ProtocolConnection,
    }

    impl Stream for DrainingProtocol {
        type Item = BytesMut;

        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            while matches!(self.conn.poll_next_unpin(cx), Poll::Ready(Some(_))) {}
            Poll::Pending
        }
    }

    #[derive(Debug)]
    struct SignalingDrainingProtocol {
        conn: ProtocolConnection,
        drained: Option<oneshot::Sender<()>>,
    }

    impl Stream for SignalingDrainingProtocol {
        type Item = BytesMut;

        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            if matches!(self.conn.poll_next_unpin(cx), Poll::Ready(Some(_))) &&
                let Some(drained) = self.drained.take()
            {
                let _ = drained.send(());
            }
            Poll::Pending
        }
    }

    fn compressed_frame(message_id: u8, payload: &[u8]) -> BytesMut {
        let mut encoded = vec![0; snap::raw::max_compress_len(payload.len()) + 1];
        encoded[0] = message_id;
        let len = snap::raw::Encoder::new().compress(payload, &mut encoded[1..]).unwrap();
        encoded.truncate(len + 1);
        BytesMut::from(encoded.as_slice())
    }

    fn test_multiplexer(
        frame_count: usize,
    ) -> (RlpxProtocolMultiplexer<InboundFramesTransport>, SharedCapability, SharedCapability) {
        let (hello, _) = test_hello();
        let shared_capabilities =
            SharedCapabilities::try_new(hello.protocols.clone(), hello.message().capabilities)
                .unwrap();
        let eth = shared_capabilities.eth().unwrap().clone();
        let test = shared_capabilities.find(&TestProtoMessage::capability()).unwrap().clone();
        let frames =
            (0..frame_count).map(|_| compressed_frame(test.message_id_offset(), &[0])).collect();
        let conn = P2PStream::new(
            InboundFramesTransport { frames, writable: AtomicBool::new(true) },
            shared_capabilities,
        );
        (RlpxProtocolMultiplexer::new(conn), eth, test)
    }

    fn shared_test_capability(name: &'static str, offset: u8) -> SharedCapability {
        SharedCapability::UnknownCapability {
            cap: Capability::new_static(name, 1),
            offset,
            messages: 1,
        }
    }

    #[tokio::test]
    async fn inbound_channel_enforces_and_releases_protocol_limits() {
        let capability = shared_test_capability("test", 0x10);
        let limits =
            ProtocolIngressLimits::new(16).with_max_buffered_bytes(4).with_max_buffered_messages(1);
        let aggregate = Arc::new(InboundBudget::new(16));
        let (sender, mut receiver) = inbound_channel(&capability, limits, aggregate);

        sender.try_send(BytesMut::zeroed(4)).unwrap();
        assert!(matches!(
            sender.try_send(BytesMut::zeroed(1)),
            Err(P2PStreamError::SubprotocolInboundBufferFull { .. })
        ));

        assert_eq!(receiver.next().await.unwrap().len(), 4);
        sender.try_send(BytesMut::zeroed(1)).unwrap();
    }

    #[tokio::test]
    async fn inbound_channels_share_connection_byte_budget() {
        let aggregate = Arc::new(InboundBudget::new(6));
        let limits = ProtocolIngressLimits::new(16)
            .with_max_buffered_bytes(16)
            .with_max_buffered_messages(2);
        let (first, mut first_receiver) =
            inbound_channel(&shared_test_capability("aaa", 0x10), limits, Arc::clone(&aggregate));
        let (second, _second_receiver) =
            inbound_channel(&shared_test_capability("bbb", 0x11), limits, aggregate);

        first.try_send(BytesMut::zeroed(4)).unwrap();
        assert!(matches!(
            second.try_send(BytesMut::zeroed(4)),
            Err(P2PStreamError::SubprotocolInboundBufferFull { .. })
        ));

        let _ = first_receiver.next().await.unwrap();
        second.try_send(BytesMut::zeroed(4)).unwrap();
    }

    #[tokio::test]
    async fn stalled_satellite_fails_when_its_inbound_queue_is_full() {
        let (mut mux, eth, test) = test_multiplexer(2);
        let limits =
            ProtocolIngressLimits::new(2).with_max_buffered_bytes(2).with_max_buffered_messages(1);
        mux.install_protocol_with_limits(&TestProtoMessage::capability(), limits, |conn| {
            HoldingProtocol { _conn: conn }
        })
        .unwrap();
        let mut stream = mux
            .into_satellite_stream(eth.capability().as_ref(), |proxy| PendingPrimary {
                _proxy: proxy,
            })
            .unwrap();

        let mut cx = Context::from_waker(noop_waker_ref());
        assert!(matches!(
            Pin::new(&mut stream).poll_next(&mut cx),
            Poll::Ready(Some(Err(P2PStreamError::SubprotocolInboundBufferFull {
                capability,
            }))) if capability == test.capability().into_owned()
        ));
    }

    #[tokio::test]
    async fn satellite_is_polled_between_inbound_frames() {
        let (mut mux, eth, _) = test_multiplexer(2 * MAX_INBOUND_MESSAGES_PER_POLL);
        let limits =
            ProtocolIngressLimits::new(2).with_max_buffered_bytes(2).with_max_buffered_messages(1);
        mux.install_protocol_with_limits(&TestProtoMessage::capability(), limits, |conn| {
            DrainingProtocol { conn }
        })
        .unwrap();
        let mut stream = mux
            .into_satellite_stream(eth.capability().as_ref(), |proxy| PendingPrimary {
                _proxy: proxy,
            })
            .unwrap();

        let mut cx = Context::from_waker(noop_waker_ref());
        assert!(Pin::new(&mut stream).poll_next(&mut cx).is_pending());
        assert_eq!(stream.inner.conn.inner().frames.len(), MAX_INBOUND_MESSAGES_PER_POLL);
    }

    #[tokio::test]
    async fn outbound_backpressure_preserves_inbound_fairness() {
        let (hello, _) = test_hello();
        let shared_capabilities =
            SharedCapabilities::try_new(hello.protocols.clone(), hello.message().capabilities)
                .unwrap();
        let eth = shared_capabilities.eth().unwrap().clone();
        let test = shared_capabilities.find(&TestProtoMessage::capability()).unwrap().clone();
        let frames = (0..2)
            .map(|_| compressed_frame(test.message_id_offset(), &[0]))
            .collect::<VecDeque<_>>();
        let conn = P2PStream::new(
            InboundFramesTransport { frames, writable: AtomicBool::new(false) },
            shared_capabilities,
        );
        let mut mux = RlpxProtocolMultiplexer::new(conn);
        let limits =
            ProtocolIngressLimits::new(2).with_max_buffered_bytes(2).with_max_buffered_messages(1);
        mux.install_protocol_with_limits(&TestProtoMessage::capability(), limits, |conn| {
            DrainingProtocol { conn }
        })
        .unwrap();
        let mut stream = mux
            .into_satellite_stream(eth.capability().as_ref(), |proxy| BurstingPrimary {
                proxy,
                remaining: 3,
            })
            .unwrap();
        stream.inner.conn.start_send_unpin(Bytes::from_static(&[0, 0])).unwrap();
        stream.inner.conn.start_send_unpin(Bytes::from_static(&[0, 0])).unwrap();
        stream.inner.out_buffer.max_bytes = 1;
        stream.inner.out_buffer.push_back(Bytes::from_static(&[0, 0]));

        let mut cx = Context::from_waker(noop_waker_ref());
        assert!(Pin::new(&mut stream).poll_next(&mut cx).is_pending());
        assert_eq!(stream.inner.conn.inner().frames.len(), 2);

        assert!(Pin::new(&mut stream).poll_next(&mut cx).is_pending());
        assert_eq!(stream.inner.conn.inner().frames.len(), 2);

        stream.inner.conn.inner().writable.store(true, Ordering::Relaxed);
        assert!(Pin::new(&mut stream).poll_next(&mut cx).is_pending());
        assert_eq!(stream.inner.conn.inner().frames.len(), 2);

        assert!(matches!(Pin::new(&mut stream).poll_next(&mut cx), Poll::Ready(None)));
        assert_eq!(stream.inner.conn.inner().frames.len(), 0);
    }

    #[tokio::test]
    async fn satellite_is_polled_between_frames_during_primary_handshake() {
        let (mut mux, eth, _) = test_multiplexer(2);
        let limits =
            ProtocolIngressLimits::new(2).with_max_buffered_bytes(2).with_max_buffered_messages(1);
        let (drained_tx, drained_rx) = oneshot::channel();
        mux.install_protocol_with_limits(&TestProtoMessage::capability(), limits, |conn| {
            SignalingDrainingProtocol { conn, drained: Some(drained_tx) }
        })
        .unwrap();

        let stream = mux
            .into_satellite_stream_with_handshake(eth.capability().as_ref(), async move |proxy| {
                drained_rx.await.unwrap();
                Ok::<_, P2PStreamError>(PendingPrimary { _proxy: proxy })
            })
            .await;

        assert!(stream.is_ok());
    }

    #[tokio::test]
    async fn satellite_mux_stops_polling_protocols_when_out_buffer_is_full() {
        let (hello, _) = test_hello();
        let shared_capabilities =
            SharedCapabilities::try_new(hello.protocols.clone(), hello.message().capabilities)
                .unwrap();
        let conn = P2PStream::new(StalledTransport, shared_capabilities);
        let eth = conn.shared_capabilities().eth().unwrap().clone();

        let mut st = RlpxProtocolMultiplexer::new(conn)
            .into_satellite_stream(eth.capability().as_ref(), |proxy| PendingPrimary {
                _proxy: proxy,
            })
            .unwrap();
        const MESSAGE_COUNT: usize = 4096;
        const MESSAGE_BYTES: usize = 1024;
        st.inner.out_buffer.max_bytes = 4 * MESSAGE_BYTES + 1;
        st.install_protocol(&TestProtoMessage::capability(), |_conn| {
            stream::iter((0..MESSAGE_COUNT).map(|_| {
                let mut msg = BytesMut::zeroed(MESSAGE_BYTES);
                msg[0] = TestProtoMessage::ping().message_type as u8;
                msg
            }))
        })
        .unwrap();

        let mut cx = Context::from_waker(noop_waker_ref());
        assert!(Pin::new(&mut st).poll_next(&mut cx).is_pending());

        assert!(st.inner.out_buffer.bytes > st.inner.out_buffer.max_bytes);
        assert!(st.inner.out_buffer.bytes <= st.inner.out_buffer.max_bytes + MESSAGE_BYTES);
        assert!(st.inner.out_buffer.messages.len() < MESSAGE_COUNT);
    }

    #[tokio::test]
    async fn satellite_mux_round_robins_ready_protocols_when_out_buffer_fills() {
        let (mut hello, _) = eth_hello();
        let cap_a = Capability::new_static("aaa", 1);
        let cap_b = Capability::new_static("bbb", 1);
        hello.protocols.push(Protocol::new(cap_a.clone(), 1));
        hello.protocols.push(Protocol::new(cap_b.clone(), 1));

        let shared_capabilities =
            SharedCapabilities::try_new(hello.protocols.clone(), hello.message().capabilities)
                .unwrap();
        let conn = P2PStream::new(StalledTransport, shared_capabilities);
        let eth = conn.shared_capabilities().eth().unwrap().clone();
        let cap_a_offset =
            conn.shared_capabilities().find(&cap_a).unwrap().relative_message_id_offset();
        let cap_b_offset =
            conn.shared_capabilities().find(&cap_b).unwrap().relative_message_id_offset();

        let mut st = RlpxProtocolMultiplexer::new(conn)
            .into_satellite_stream(eth.capability().as_ref(), |proxy| PendingPrimary {
                _proxy: proxy,
            })
            .unwrap();
        st.inner.out_buffer.max_bytes = 5;
        st.install_protocol(&cap_a, |_conn| {
            stream::iter((0..16).map(|_| BytesMut::from(&[0, b'a'][..])))
        })
        .unwrap();
        st.install_protocol(&cap_b, |_conn| {
            stream::iter((0..16).map(|_| BytesMut::from(&[0, b'b'][..])))
        })
        .unwrap();

        let mut cx = Context::from_waker(noop_waker_ref());
        assert!(Pin::new(&mut st).poll_next(&mut cx).is_pending());

        let message_ids =
            st.inner.out_buffer.messages.iter().take(2).map(|msg| msg[0]).collect::<Vec<_>>();
        assert_eq!(message_ids.len(), 2);
        assert_ne!(message_ids[0], message_ids[1]);
        assert!(message_ids.contains(&cap_a_offset));
        assert!(message_ids.contains(&cap_b_offset));
    }

    #[tokio::test]
    async fn eth_satellite() {
        reth_tracing::init_test_tracing();
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let local_addr = listener.local_addr().unwrap();
        let (status, fork_filter) = eth_handshake();
        let other_status = status;
        let other_fork_filter = fork_filter.clone();
        let _handle = tokio::spawn(async move {
            let (incoming, _) = listener.accept().await.unwrap();
            let stream = crate::PassthroughCodec::default().framed(incoming);
            let (server_hello, _) = eth_hello();
            let (p2p_stream, _) =
                UnauthedP2PStream::new(stream).handshake(server_hello).await.unwrap();

            let (_eth_stream, _) = UnauthedEthStream::new(p2p_stream)
                .handshake::<EthNetworkPrimitives>(other_status, other_fork_filter)
                .await
                .unwrap();

            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        });

        let conn = connect_passthrough(local_addr, eth_hello().0).await;
        let eth = conn.shared_capabilities().eth().unwrap().clone();

        let multiplexer = RlpxProtocolMultiplexer::new(conn);
        let _satellite = multiplexer
            .into_satellite_stream_with_handshake(eth.capability().as_ref(), async move |proxy| {
                UnauthedEthStream::new(proxy)
                    .handshake::<EthNetworkPrimitives>(status, fork_filter)
                    .await
            })
            .await
            .unwrap();
    }

    /// A test that install a satellite stream eth+test protocol and sends messages between them.
    #[tokio::test(flavor = "multi_thread")]
    async fn eth_test_protocol_satellite() {
        reth_tracing::init_test_tracing();
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let local_addr = listener.local_addr().unwrap();
        let (status, fork_filter) = eth_handshake();
        let other_status = status;
        let other_fork_filter = fork_filter.clone();
        let _handle = tokio::spawn(async move {
            let (incoming, _) = listener.accept().await.unwrap();
            let stream = crate::PassthroughCodec::default().framed(incoming);
            let (server_hello, _) = test_hello();
            let (conn, _) = UnauthedP2PStream::new(stream).handshake(server_hello).await.unwrap();

            let (mut st, _their_status) = RlpxProtocolMultiplexer::new(conn)
                .into_eth_satellite_stream::<EthNetworkPrimitives>(
                    other_status,
                    other_fork_filter,
                    Arc::new(EthHandshake::default()),
                    MAX_MESSAGE_SIZE,
                )
                .await
                .unwrap();

            st.install_protocol(&TestProtoMessage::capability(), |mut conn| {
                async_stream::stream! {
                    yield TestProtoMessage::ping().encoded();
                    let msg = conn.next().await.unwrap();
                    let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                    assert_eq!(msg, TestProtoMessage::pong());

                    yield TestProtoMessage::message("hello").encoded();
                    let msg = conn.next().await.unwrap();
                    let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                    assert_eq!(msg, TestProtoMessage::message("good bye!"));

                    yield TestProtoMessage::message("good bye!").encoded();

                    futures::future::pending::<()>().await;
                    unreachable!()
                }
            })
            .unwrap();

            loop {
                let _ = st.next().await;
            }
        });

        let conn = connect_passthrough(local_addr, test_hello().0).await;
        let (mut st, _their_status) = RlpxProtocolMultiplexer::new(conn)
            .into_eth_satellite_stream::<EthNetworkPrimitives>(
                status,
                fork_filter,
                Arc::new(EthHandshake::default()),
                MAX_MESSAGE_SIZE,
            )
            .await
            .unwrap();

        let (tx, mut rx) = oneshot::channel();

        st.install_protocol(&TestProtoMessage::capability(), |mut conn| {
            async_stream::stream! {
                let msg = conn.next().await.unwrap();
                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                assert_eq!(msg, TestProtoMessage::ping());

                yield TestProtoMessage::pong().encoded();

                let msg = conn.next().await.unwrap();
                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                assert_eq!(msg, TestProtoMessage::message("hello"));

                yield TestProtoMessage::message("good bye!").encoded();

                let msg = conn.next().await.unwrap();
                let msg = TestProtoMessage::decode_message(&mut &msg[..]).unwrap();
                assert_eq!(msg, TestProtoMessage::message("good bye!"));

                tx.send(()).unwrap();

                futures::future::pending::<()>().await;
                unreachable!()
            }
        })
        .unwrap();

        loop {
            tokio::select! {
                _ = &mut rx => {
                    break
                }
               _ = st.next() => {
                }
            }
        }
    }

    #[tokio::test]
    async fn satellite_sink_drains_accepted_primary_messages() {
        let (mux, eth, _) = test_multiplexer(0);
        let mut stream = mux
            .into_satellite_stream(eth.capability().as_ref(), |proxy| UnauthProxy { inner: proxy })
            .unwrap();
        stream.inner.conn.inner().writable.store(false, Ordering::Relaxed);
        let mut cx = Context::from_waker(noop_waker_ref());

        // Mirrors the session, which hands over queued messages while the sink reports readiness.
        let mut sent = 0;
        while sent < 64 &&
            matches!(
                Sink::<Bytes>::poll_ready(Pin::new(&mut stream), &mut cx),
                Poll::Ready(Ok(()))
            )
        {
            Pin::new(&mut stream).start_send(Bytes::from_static(&[0, 0])).unwrap();
            sent += 1;
        }
        // A stalled transport only admits the two messages the lower p2p sink buffers.
        assert_eq!(sent, 2);
        assert!(stream.primary.from_primary.as_ref().is_empty());
        assert!(stream.inner.out_buffer.is_empty());

        stream.inner.conn.inner().writable.store(true, Ordering::Relaxed);
        assert!(matches!(
            Sink::<Bytes>::poll_ready(Pin::new(&mut stream), &mut cx),
            Poll::Ready(Ok(()))
        ));
        Pin::new(&mut stream).start_send(Bytes::from_static(&[0, 0])).unwrap();
        assert!(matches!(
            Sink::<Bytes>::poll_flush(Pin::new(&mut stream), &mut cx),
            Poll::Ready(Ok(()))
        ));
        assert!(stream.primary.from_primary.as_ref().is_empty());
    }
}