use std::{
collections::VecDeque,
fmt,
future::Future,
io,
pin::{pin, Pin},
task::{ready, Context, Poll},
};
use crate::{
capability::{SharedCapabilities, SharedCapability, UnsupportedCapabilityError},
errors::{EthStreamError, P2PStreamError},
p2pstream::DisconnectP2P,
CanDisconnect, Capability, DisconnectReason, EthStream, P2PStream, Status, UnauthedEthStream,
};
use bytes::{Bytes, BytesMut};
use futures::{Sink, SinkExt, Stream, StreamExt, TryStream, TryStreamExt};
use reth_eth_wire_types::NetworkPrimitives;
use reth_ethereum_forks::ForkFilter;
use tokio::sync::{mpsc, mpsc::UnboundedSender};
use tokio_stream::wrappers::UnboundedReceiverStream;
#[derive(Debug)]
pub struct RlpxProtocolMultiplexer<St> {
inner: MultiplexInner<St>,
}
impl<St> RlpxProtocolMultiplexer<St> {
pub fn new(conn: P2PStream<St>) -> Self {
Self {
inner: MultiplexInner {
conn,
protocols: Default::default(),
out_buffer: Default::default(),
},
}
}
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, f)
}
pub const fn shared_capabilities(&self) -> &SharedCapabilities {
self.inner.shared_capabilities()
}
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) = mpsc::unbounded_channel();
let (to_wire, from_primary) = mpsc::unbounded_channel();
let proxy = ProtocolProxy {
shared_cap: shared_cap.clone(),
from_wire: UnboundedReceiverStream::new(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,
},
})
}
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, move |proxy| async move {
let st = handshake(proxy).await?;
Ok((st, ()))
})
.await
.map(|(st, _)| st)
}
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) = mpsc::unbounded_channel();
let (to_wire, mut from_primary) = mpsc::unbounded_channel();
let proxy = ProtocolProxy {
shared_cap: shared_cap.clone(),
from_wire: UnboundedReceiverStream::new(from_wire),
to_wire,
};
let f = handshake(proxy);
let mut f = pin!(f);
loop {
tokio::select! {
Some(Ok(msg)) = self.inner.conn.next() => {
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 {
let _ = to_primary.send(msg);
} else {
self.inner.delegate_message(&cap, msg);
}
} else {
return Err(P2PStreamError::UnknownReservedMessageId(offset).into())
}
}
Some(msg) = from_primary.recv() => {
self.inner.conn.send(msg).await.map_err(Into::into)?;
}
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,
}
}, extra))
}
}
}
}
pub async fn into_eth_satellite_stream<N: NetworkPrimitives>(
self,
status: Status,
fork_filter: ForkFilter,
) -> Result<(RlpxSatelliteStream<St, EthStream<ProtocolProxy, N>>, Status), 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),
move |proxy| async move {
UnauthedEthStream::new(proxy).handshake(status, fork_filter).await
},
)
.await
}
}
#[derive(Debug)]
struct MultiplexInner<St> {
conn: P2PStream<St>,
protocols: Vec<ProtocolStream>,
out_buffer: VecDeque<Bytes>,
}
impl<St> MultiplexInner<St> {
const fn shared_capabilities(&self) -> &SharedCapabilities {
self.conn.shared_capabilities()
}
fn delegate_message(&self, cap: &SharedCapability, msg: BytesMut) -> bool {
for proto in &self.protocols {
if proto.shared_cap == *cap {
proto.send_raw(msg);
return true
}
}
false
}
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 shared_cap =
self.conn.shared_capabilities().ensure_matching_capability(cap).cloned()?;
let (to_satellite, rx) = mpsc::unbounded_channel();
let proto_conn = ProtocolConnection { from_wire: UnboundedReceiverStream::new(rx) };
let st = f(proto_conn);
let st = ProtocolStream { shared_cap, to_satellite, satellite_st: Box::pin(st) };
self.protocols.push(st);
Ok(())
}
}
#[derive(Debug)]
struct PrimaryProtocol<Primary> {
to_primary: UnboundedSender<BytesMut>,
from_primary: UnboundedReceiverStream<Bytes>,
shared_cap: SharedCapability,
st: Primary,
}
#[derive(Debug)]
pub struct ProtocolProxy {
shared_cap: SharedCapability,
from_wire: UnboundedReceiverStream<BytesMut>,
to_wire: UnboundedSender<Bytes>,
}
impl ProtocolProxy {
fn try_send(&self, msg: Bytes) -> Result<(), io::Error> {
if msg.is_empty() {
return Err(io::ErrorKind::InvalidInput.into())
}
self.to_wire.send(self.mask_msg_id(msg)?).map_err(|_| io::ErrorKind::BrokenPipe.into())
}
#[inline]
fn mask_msg_id(&self, msg: Bytes) -> Result<Bytes, io::Error> {
if msg.is_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 = Vec::from(msg);
masked[0] = masked[0].checked_add(offset).ok_or(io::ErrorKind::InvalidInput)?;
Ok(masked.into())
}
#[inline]
fn unmask_id(&self, mut msg: BytesMut) -> Result<BytesMut, io::Error> {
if msg.is_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 {
async fn disconnect(
&mut self,
_reason: DisconnectReason,
) -> Result<(), <Self as Sink<Bytes>>::Error> {
Ok(())
}
}
#[derive(Debug)]
pub struct ProtocolConnection {
from_wire: UnboundedReceiverStream<BytesMut>,
}
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)
}
}
#[derive(Debug)]
pub struct RlpxSatelliteStream<St, Primary> {
inner: MultiplexInner<St>,
primary: PrimaryProtocol<Primary>,
}
impl<St, Primary> RlpxSatelliteStream<St, Primary> {
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, f)
}
#[inline]
pub const fn primary(&self) -> &Primary {
&self.primary.st
}
#[inline]
pub fn primary_mut(&mut self) -> &mut Primary {
&mut self.primary.st
}
#[inline]
pub const fn inner(&self) -> &P2PStream<St> {
&self.inner.conn
}
#[inline]
pub fn inner_mut(&mut self) -> &mut P2PStream<St> {
&mut self.inner.conn
}
#[inline]
pub fn into_inner(self) -> P2PStream<St> {
self.inner.conn
}
}
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();
loop {
if let Poll::Ready(Some(msg)) = this.primary.st.try_poll_next_unpin(cx) {
return Poll::Ready(Some(msg))
}
let mut conn_ready = true;
loop {
match this.inner.conn.poll_ready_unpin(cx) {
Poll::Ready(Ok(())) => {
if let Some(msg) = this.inner.out_buffer.pop_front() {
if let Err(err) = this.inner.conn.start_send_unpin(msg) {
return Poll::Ready(Some(Err(err.into())))
}
} else {
break
}
}
Poll::Ready(Err(err)) => {
if let Err(disconnect_err) =
this.inner.conn.start_disconnect(DisconnectReason::DisconnectRequested)
{
return Poll::Ready(Some(Err(disconnect_err.into())))
}
return Poll::Ready(Some(Err(err.into())))
}
Poll::Pending => {
conn_ready = false;
break
}
}
}
loop {
match this.primary.from_primary.poll_next_unpin(cx) {
Poll::Ready(Some(msg)) => {
this.inner.out_buffer.push_back(msg);
}
Poll::Ready(None) => {
return Poll::Ready(None)
}
Poll::Pending => break,
}
}
for idx in (0..this.inner.protocols.len()).rev() {
let mut proto = this.inner.protocols.swap_remove(idx);
loop {
match proto.poll_next_unpin(cx) {
Poll::Ready(Some(Err(err))) => {
return Poll::Ready(Some(Err(P2PStreamError::Io(err).into())))
}
Poll::Ready(Some(Ok(msg))) => {
this.inner.out_buffer.push_back(msg);
}
Poll::Ready(None) => return Poll::Ready(None),
Poll::Pending => {
this.inner.protocols.push(proto);
break
}
}
}
}
let mut delegated = false;
loop {
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()
)))
};
if let Some(cap) =
this.inner.conn.shared_capabilities().find_by_relative_offset(offset)
{
if cap == &this.primary.shared_cap {
let _ = this.primary.to_primary.send(msg);
} else {
for proto in &this.inner.protocols {
if proto.shared_cap == *cap {
proto.send_raw(msg);
break
}
}
}
} else {
return Poll::Ready(Some(Err(P2PStreamError::UnknownReservedMessageId(
offset,
)
.into())))
}
}
Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err.into()))),
Poll::Ready(None) => {
return Poll::Ready(None)
}
Poll::Pending => break,
}
}
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.inner.conn.poll_ready_unpin(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>> {
self.get_mut().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)
}
}
struct ProtocolStream {
shared_cap: SharedCapability,
to_satellite: UnboundedSender<BytesMut>,
satellite_st: Pin<Box<dyn Stream<Item = BytesMut> + Send>>,
}
impl ProtocolStream {
#[inline]
fn mask_msg_id(&self, mut msg: BytesMut) -> Result<Bytes, io::Error> {
if msg.is_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())
}
#[inline]
fn unmask_id(&self, mut msg: BytesMut) -> Result<BytesMut, io::Error> {
if msg.is_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)
}
fn send_raw(&self, msg: BytesMut) {
let _ = self.unmask_id(msg).map(|msg| self.to_satellite.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()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
test_utils::{
connect_passthrough, eth_handshake, eth_hello,
proto::{test_hello, TestProtoMessage},
},
UnauthedP2PStream,
};
use reth_eth_wire_types::EthNetworkPrimitives;
use tokio::{net::TcpListener, sync::oneshot};
use tokio_util::codec::Decoder;
#[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(),
move |proxy| async move {
UnauthedEthStream::new(proxy)
.handshake::<EthNetworkPrimitives>(status, fork_filter)
.await
},
)
.await
.unwrap();
}
#[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)
.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)
.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() => {
}
}
}
}
}