reth_ecies/
codec.rs

1//! This contains the main codec for `RLPx` ECIES messages
2
3use crate::{algorithm::ECIES, ECIESError, ECIESErrorImpl, EgressECIESValue, IngressECIESValue};
4use alloy_primitives::{bytes::BytesMut, B512 as PeerId};
5use secp256k1::SecretKey;
6use std::{fmt::Debug, io};
7use tokio_util::codec::{Decoder, Encoder};
8use tracing::{instrument, trace};
9
10/// The max size that the initial handshake packet can be. Currently 2KiB.
11const MAX_INITIAL_HANDSHAKE_SIZE: usize = 2048;
12
13/// The largest frame body the 24-bit size field of an `RLPx` frame header can describe.
14pub(crate) const MAX_FRAME_SIZE: usize = (1 << 24) - 1;
15
16/// Tokio codec for ECIES
17#[derive(Debug)]
18pub struct ECIESCodec {
19    ecies: ECIES,
20    state: ECIESState,
21}
22
23/// Current ECIES state of a connection
24#[derive(Clone, Copy, Debug, PartialEq, Eq)]
25pub enum ECIESState {
26    /// The first stage of the ECIES handshake, where each side of the connection sends an auth
27    /// message containing the ephemeral public key, signature of the public key, nonce, and other
28    /// metadata.
29    Auth,
30
31    /// The second stage of the ECIES handshake, where each side of the connection sends an ack
32    /// message containing the nonce and other metadata.
33    Ack,
34
35    /// This is the same as the [`ECIESState::Header`] stage, but occurs only after the first
36    /// [`ECIESState::Ack`] message. This is so that the initial handshake message can be properly
37    /// validated.
38    InitialHeader,
39
40    /// The third stage of the ECIES handshake, where header is parsed, message integrity checks
41    /// performed, and message is decrypted.
42    Header,
43
44    /// The final stage, where the ECIES message is actually read and returned by the ECIES codec.
45    Body,
46}
47
48impl ECIESCodec {
49    /// Create a new server codec using the given secret key
50    pub(crate) fn new_server(secret_key: SecretKey) -> Result<Self, ECIESError> {
51        Ok(Self { ecies: ECIES::new_server(secret_key)?, state: ECIESState::Auth })
52    }
53
54    /// Create a new client codec using the given secret key and the server's public id
55    pub(crate) fn new_client(secret_key: SecretKey, remote_id: PeerId) -> Result<Self, ECIESError> {
56        Ok(Self { ecies: ECIES::new_client(secret_key, remote_id)?, state: ECIESState::Auth })
57    }
58}
59
60impl Decoder for ECIESCodec {
61    type Item = IngressECIESValue;
62    type Error = ECIESError;
63
64    #[instrument(level = "trace", target = "net::ecies", skip_all, fields(peer=?self.ecies.remote_id, state=?self.state))]
65    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
66        loop {
67            match self.state {
68                ECIESState::Auth => {
69                    trace!("parsing auth");
70                    if buf.len() < 2 {
71                        return Ok(None)
72                    }
73
74                    let payload_size = u16::from_be_bytes([buf[0], buf[1]]) as usize;
75                    let total_size = payload_size + 2;
76
77                    if buf.len() < total_size {
78                        trace!("current len {}, need {}", buf.len(), total_size);
79                        return Ok(None)
80                    }
81
82                    self.ecies.read_auth(&mut buf.split_to(total_size))?;
83
84                    self.state = ECIESState::InitialHeader;
85                    return Ok(Some(IngressECIESValue::AuthReceive(self.ecies.remote_id())))
86                }
87                ECIESState::Ack => {
88                    trace!("parsing ack with len {}", buf.len());
89                    if buf.len() < 2 {
90                        return Ok(None)
91                    }
92
93                    let payload_size = u16::from_be_bytes([buf[0], buf[1]]) as usize;
94                    let total_size = payload_size + 2;
95
96                    if buf.len() < total_size {
97                        trace!("current len {}, need {}", buf.len(), total_size);
98                        return Ok(None)
99                    }
100
101                    self.ecies.read_ack(&mut buf.split_to(total_size))?;
102
103                    self.state = ECIESState::InitialHeader;
104                    return Ok(Some(IngressECIESValue::Ack))
105                }
106                ECIESState::InitialHeader => {
107                    if buf.len() < ECIES::header_len() {
108                        trace!("current len {}, need {}", buf.len(), ECIES::header_len());
109                        return Ok(None)
110                    }
111
112                    let body_size =
113                        self.ecies.read_header(&mut buf.split_to(ECIES::header_len()))?;
114
115                    if body_size > MAX_INITIAL_HANDSHAKE_SIZE {
116                        trace!(?body_size, max=?MAX_INITIAL_HANDSHAKE_SIZE, "Body exceeds max initial handshake size");
117                        return Err(ECIESErrorImpl::InitialHeaderBodyTooLarge {
118                            body_size,
119                            max_body_size: MAX_INITIAL_HANDSHAKE_SIZE,
120                        }
121                        .into())
122                    }
123
124                    self.state = ECIESState::Body;
125                }
126                ECIESState::Header => {
127                    if buf.len() < ECIES::header_len() {
128                        trace!("current len {}, need {}", buf.len(), ECIES::header_len());
129                        return Ok(None)
130                    }
131
132                    self.ecies.read_header(&mut buf.split_to(ECIES::header_len()))?;
133
134                    self.state = ECIESState::Body;
135                }
136                ECIESState::Body => {
137                    if buf.len() < self.ecies.body_len() {
138                        return Ok(None)
139                    }
140
141                    let mut data = buf.split_to(self.ecies.body_len());
142                    let body_len = {
143                        let body = self.ecies.read_body(&mut data)?;
144                        body.len()
145                    };
146                    data.truncate(body_len);
147
148                    self.state = ECIESState::Header;
149                    return Ok(Some(IngressECIESValue::Message(data)))
150                }
151            }
152        }
153    }
154}
155
156impl Encoder<EgressECIESValue> for ECIESCodec {
157    type Error = io::Error;
158
159    #[instrument(level = "trace", target = "net::ecies", skip(self, buf), fields(peer=?self.ecies.remote_id, state=?self.state))]
160    fn encode(&mut self, item: EgressECIESValue, buf: &mut BytesMut) -> Result<(), Self::Error> {
161        match item {
162            EgressECIESValue::Auth => {
163                self.state = ECIESState::Ack;
164                self.ecies.write_auth(buf);
165            }
166            EgressECIESValue::Ack => {
167                self.state = ECIESState::InitialHeader;
168                self.ecies.write_ack(buf);
169            }
170            EgressECIESValue::Message(data) => {
171                if data.len() > MAX_FRAME_SIZE {
172                    return Err(io::Error::new(
173                        io::ErrorKind::InvalidInput,
174                        ECIESError::from(ECIESErrorImpl::FrameTooLarge {
175                            size: data.len(),
176                            max_size: MAX_FRAME_SIZE,
177                        }),
178                    ))
179                }
180                self.ecies.write_header(buf, data.len());
181                self.ecies.write_body(buf, &data);
182            }
183        }
184        Ok(())
185    }
186}
reth_ecies/codec.rs

reth_ecies/
codec.rs
