reth_ecies/
stream.rs

1//! The ECIES Stream implementation which wraps over [`AsyncRead`] and [`AsyncWrite`].
2
3use crate::{
4    codec::ECIESCodec, error::ECIESErrorImpl, ECIESError, EgressECIESValue, IngressECIESValue,
5};
6use alloy_primitives::{
7    bytes::{Bytes, BytesMut},
8    B512 as PeerId,
9};
10use futures::{ready, Sink, SinkExt};
11use secp256k1::SecretKey;
12use std::{
13    fmt::Debug,
14    io,
15    pin::Pin,
16    task::{Context, Poll},
17    time::Duration,
18};
19use tokio::{
20    io::{AsyncRead, AsyncWrite},
21    time::timeout,
22};
23use tokio_stream::{Stream, StreamExt};
24use tokio_util::codec::{Decoder, Framed};
25use tracing::{instrument, trace};
26
27const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
28
29/// Write buffer size at which the underlying `Framed` transport starts flushing frames to the
30/// socket from `poll_ready`, replacing the tokio-util default of 8KiB.
31///
32/// `RLPx` callers queue batches of messages (e.g. transaction broadcasts) and drive one flush per
33/// batch; the 8KiB default splits such a batch into one write syscall per few messages. A larger
34/// boundary batches them into fewer, larger writes. Once grown, the write buffer's capacity is
35/// retained for the connection's lifetime, so this also bounds the resident write buffer per
36/// connection.
37pub const DEFAULT_BACKPRESSURE_BOUNDARY: usize = 64 * 1024;
38
39/// `ECIES` stream over TCP exchanging raw bytes
40#[derive(Debug)]
41#[pin_project::pin_project]
42pub struct ECIESStream<Io> {
43    #[pin]
44    stream: Framed<Io, ECIESCodec>,
45    remote_id: PeerId,
46}
47
48impl<Io> ECIESStream<Io>
49where
50    Io: AsyncRead + AsyncWrite + Unpin,
51{
52    /// Connect to an `ECIES` server
53    #[instrument(level = "trace", target = "net::ecies", skip(transport, secret_key))]
54    pub async fn connect(
55        transport: Io,
56        secret_key: SecretKey,
57        remote_id: PeerId,
58    ) -> Result<Self, ECIESError> {
59        Self::connect_with_timeout(transport, secret_key, remote_id, HANDSHAKE_TIMEOUT).await
60    }
61
62    /// Wrapper around `connect_no_timeout` which enforces a timeout.
63    pub async fn connect_with_timeout(
64        transport: Io,
65        secret_key: SecretKey,
66        remote_id: PeerId,
67        timeout_limit: Duration,
68    ) -> Result<Self, ECIESError> {
69        timeout(timeout_limit, Self::connect_without_timeout(transport, secret_key, remote_id))
70            .await
71            .map_err(|_| ECIESError::from(ECIESErrorImpl::StreamTimeout))?
72    }
73
74    /// Connect to an `ECIES` server with no timeout.
75    pub async fn connect_without_timeout(
76        transport: Io,
77        secret_key: SecretKey,
78        remote_id: PeerId,
79    ) -> Result<Self, ECIESError> {
80        let ecies = ECIESCodec::new_client(secret_key, remote_id)?;
81
82        let mut transport = ecies.framed(transport);
83        transport.set_backpressure_boundary(DEFAULT_BACKPRESSURE_BOUNDARY);
84
85        trace!("sending ecies auth ...");
86        transport.send(EgressECIESValue::Auth).await?;
87
88        trace!("waiting for ecies ack ...");
89
90        let msg = transport.try_next().await?;
91
92        // `Framed` returns `None` if the underlying stream is no longer readable, and the codec is
93        // unable to decode another message from the (partially filled) buffer. This usually happens
94        // if the remote drops the TcpStream.
95        let msg = msg.ok_or(ECIESErrorImpl::UnreadableStream)?;
96
97        trace!("parsing ecies ack ...");
98        if matches!(msg, IngressECIESValue::Ack) {
99            Ok(Self { stream: transport, remote_id })
100        } else {
101            Err(ECIESErrorImpl::InvalidHandshake {
102                expected: IngressECIESValue::Ack,
103                msg: Some(msg),
104            }
105            .into())
106        }
107    }
108
109    /// Listen on a just connected ECIES client
110    pub async fn incoming(transport: Io, secret_key: SecretKey) -> Result<Self, ECIESError> {
111        let ecies = ECIESCodec::new_server(secret_key)?;
112
113        trace!("incoming ecies stream");
114        let mut transport = ecies.framed(transport);
115        transport.set_backpressure_boundary(DEFAULT_BACKPRESSURE_BOUNDARY);
116        let msg = transport.try_next().await?;
117
118        trace!("receiving ecies auth");
119        let remote_id = match &msg {
120            Some(IngressECIESValue::AuthReceive(remote_id)) => *remote_id,
121            _ => {
122                return Err(ECIESErrorImpl::InvalidHandshake {
123                    expected: IngressECIESValue::AuthReceive(Default::default()),
124                    msg,
125                }
126                .into())
127            }
128        };
129
130        trace!("sending ecies ack");
131        transport.send(EgressECIESValue::Ack).await?;
132
133        Ok(Self { stream: transport, remote_id })
134    }
135
136    /// Get the remote id
137    pub const fn remote_id(&self) -> PeerId {
138        self.remote_id
139    }
140}
141
142impl<Io> ECIESStream<Io> {
143    /// Sets the write buffer size at which the underlying transport starts flushing to the socket
144    /// from `poll_ready`, overriding [`DEFAULT_BACKPRESSURE_BOUNDARY`].
145    ///
146    /// A larger boundary batches more frames into a single write syscall when many messages are
147    /// sent back to back, at the cost of buffering more encrypted data in memory per connection.
148    pub fn set_backpressure_boundary(&mut self, boundary: usize) {
149        self.stream.set_backpressure_boundary(boundary);
150    }
151}
152
153impl<Io> Stream for ECIESStream<Io>
154where
155    Io: AsyncRead + Unpin,
156{
157    type Item = Result<BytesMut, io::Error>;
158
159    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
160        match ready!(self.project().stream.poll_next(cx)) {
161            Some(Ok(IngressECIESValue::Message(body))) => Poll::Ready(Some(Ok(body))),
162            Some(other) => Poll::Ready(Some(Err(io::Error::other(format!(
163                "ECIES stream protocol error: expected message, received {other:?}"
164            ))))),
165            None => Poll::Ready(None),
166        }
167    }
168}
169
170impl<Io> Sink<Bytes> for ECIESStream<Io>
171where
172    Io: AsyncWrite + Unpin,
173{
174    type Error = io::Error;
175
176    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
177        self.project().stream.poll_ready(cx)
178    }
179
180    fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
181        self.project().stream.start_send(EgressECIESValue::Message(item))?;
182        Ok(())
183    }
184
185    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
186        self.project().stream.poll_flush(cx)
187    }
188
189    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
190        self.project().stream.poll_close(cx)
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use reth_network_peers::pk2id;
198    use secp256k1::SECP256K1;
199    use tokio::net::{TcpListener, TcpStream};
200
201    #[tokio::test]
202    async fn can_write_and_read() {
203        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
204        let addr = listener.local_addr().unwrap();
205        let server_key = SecretKey::new(&mut rand_08::thread_rng());
206
207        let handle = tokio::spawn(async move {
208            // roughly based off of the design of tokio::net::TcpListener
209            let (incoming, _) = listener.accept().await.unwrap();
210            let mut stream = ECIESStream::incoming(incoming, server_key).await.unwrap();
211
212            // use the stream to get the next message
213            let message = stream.next().await.unwrap().unwrap();
214            assert_eq!(message, Bytes::from("hello"));
215        });
216
217        // create the server pubkey
218        let server_id = pk2id(&server_key.public_key(SECP256K1));
219
220        let client_key = SecretKey::new(&mut rand_08::thread_rng());
221        let outgoing = TcpStream::connect(addr).await.unwrap();
222        let mut client_stream =
223            ECIESStream::connect(outgoing, client_key, server_id).await.unwrap();
224        client_stream.send(Bytes::from("hello")).await.unwrap();
225
226        // make sure the server receives the message and asserts before ending the test
227        handle.await.unwrap();
228    }
229
230    #[tokio::test]
231    async fn connection_should_timeout() {
232        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
233        let addr = listener.local_addr().unwrap();
234        let server_key = SecretKey::new(&mut rand_08::thread_rng());
235
236        let _handle = tokio::spawn(async move {
237            // Delay accepting the connection for longer than the client's timeout period
238            tokio::time::sleep(Duration::from_secs(11)).await;
239            let (incoming, _) = listener.accept().await.unwrap();
240            let mut stream = ECIESStream::incoming(incoming, server_key).await.unwrap();
241
242            // use the stream to get the next message
243            let message = stream.next().await.unwrap().unwrap();
244            assert_eq!(message, Bytes::from("hello"));
245        });
246
247        // create the server pubkey
248        let server_id = pk2id(&server_key.public_key(SECP256K1));
249
250        let client_key = SecretKey::new(&mut rand_08::thread_rng());
251        let outgoing = TcpStream::connect(addr).await.unwrap();
252
253        // Attempt to connect, expecting a timeout due to the server's delayed response
254        let connect_result = ECIESStream::connect_with_timeout(
255            outgoing,
256            client_key,
257            server_id,
258            Duration::from_secs(1),
259        )
260        .await;
261
262        // Assert that a timeout error occurred
263        assert!(
264            matches!(connect_result, Err(e) if e.to_string() == ECIESErrorImpl::StreamTimeout.to_string())
265        );
266    }
267}
reth_ecies/stream.rs

reth_ecies/
stream.rs
