1use crate::{Capability, EthVersion, ProtocolVersion};
2use alloy_rlp::{RlpDecodable, RlpEncodable};
3use reth_codecs::add_arbitrary_tests;
4use reth_network_peers::PeerId;
5use reth_primitives_traits::constants::RETH_CLIENT_VERSION;
6
7pub(crate) const DEFAULT_TCP_PORT: u16 = 30303;
11
12use crate::protocol::Protocol;
13#[cfg(feature = "serde")]
14use serde::{Deserialize, Serialize};
15
16#[derive(Debug, Clone, PartialEq, Eq)]
26#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
27pub struct HelloMessageWithProtocols {
28 pub protocol_version: ProtocolVersion,
30 pub client_version: String,
33 pub protocols: Vec<Protocol>,
35 pub port: u16,
39 pub id: PeerId,
41}
42
43impl HelloMessageWithProtocols {
44 pub const fn builder(id: PeerId) -> HelloMessageBuilder {
55 HelloMessageBuilder::new(id)
56 }
57
58 #[inline]
60 pub fn message(&self) -> HelloMessage {
61 HelloMessage {
62 protocol_version: self.protocol_version,
63 client_version: self.client_version.clone(),
64 capabilities: self.protocols.iter().map(|p| p.cap.clone()).collect(),
65 port: self.port,
66 id: self.id,
67 }
68 }
69
70 pub fn into_message(self) -> HelloMessage {
72 HelloMessage {
73 protocol_version: self.protocol_version,
74 client_version: self.client_version,
75 capabilities: self.protocols.into_iter().map(|p| p.cap).collect(),
76 port: self.port,
77 id: self.id,
78 }
79 }
80
81 #[inline]
83 pub fn contains_protocol(&self, protocol: &Protocol) -> bool {
84 self.protocols.iter().any(|p| p.cap == protocol.cap)
85 }
86
87 #[inline]
91 pub fn try_add_protocol(&mut self, protocol: Protocol) -> Result<(), Protocol> {
92 if self.contains_protocol(&protocol) {
93 Err(protocol)
94 } else {
95 self.protocols.push(protocol);
96 Ok(())
97 }
98 }
99}
100
101#[derive(Clone, Debug, PartialEq, Eq, RlpEncodable, RlpDecodable)]
107#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
108#[cfg_attr(any(test, feature = "arbitrary"), derive(arbitrary::Arbitrary))]
109#[add_arbitrary_tests(rlp)]
110pub struct HelloMessage {
111 pub protocol_version: ProtocolVersion,
113 pub client_version: String,
116 pub capabilities: Vec<Capability>,
118 pub port: u16,
120 pub id: PeerId,
122}
123
124impl HelloMessage {
127 pub const fn builder(id: PeerId) -> HelloMessageBuilder {
138 HelloMessageBuilder::new(id)
139 }
140}
141
142#[derive(Debug)]
144pub struct HelloMessageBuilder {
145 pub protocol_version: Option<ProtocolVersion>,
147 pub client_version: Option<String>,
150 pub protocols: Option<Vec<Protocol>>,
152 pub port: Option<u16>,
154 pub id: PeerId,
156}
157
158impl HelloMessageBuilder {
161 pub const fn new(id: PeerId) -> Self {
163 Self { protocol_version: None, client_version: None, protocols: None, port: None, id }
164 }
165
166 pub const fn port(mut self, port: u16) -> Self {
168 self.port = Some(port);
169 self
170 }
171
172 pub fn protocol(mut self, protocols: impl Into<Protocol>) -> Self {
174 self.protocols.get_or_insert_with(Vec::new).push(protocols.into());
175 self
176 }
177
178 pub fn protocols(mut self, protocols: impl IntoIterator<Item = Protocol>) -> Self {
180 self.protocols.get_or_insert_with(Vec::new).extend(protocols);
181 self
182 }
183
184 pub fn client_version(mut self, client_version: impl Into<String>) -> Self {
186 self.client_version = Some(client_version.into());
187 self
188 }
189
190 pub const fn protocol_version(mut self, protocol_version: ProtocolVersion) -> Self {
192 self.protocol_version = Some(protocol_version);
193 self
194 }
195
196 pub fn build(self) -> HelloMessageWithProtocols {
203 let Self { protocol_version, client_version, protocols, port, id } = self;
204 HelloMessageWithProtocols {
205 protocol_version: protocol_version.unwrap_or_default(),
206 client_version: client_version.unwrap_or_else(|| RETH_CLIENT_VERSION.to_string()),
207 protocols: protocols.unwrap_or_else(|| {
208 vec![EthVersion::Eth68.into(), EthVersion::Eth67.into(), EthVersion::Eth66.into()]
209 }),
211 port: port.unwrap_or(DEFAULT_TCP_PORT),
212 id,
213 }
214 }
215}
216
217#[cfg(test)]
218mod tests {
219 use crate::{p2pstream::P2PMessage, Capability, EthVersion, HelloMessage, ProtocolVersion};
220 use alloy_rlp::{Decodable, Encodable, EMPTY_STRING_CODE};
221 use reth_network_peers::pk2id;
222 use secp256k1::{SecretKey, SECP256K1};
223
224 #[test]
225 fn test_hello_encoding_round_trip() {
226 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
227 let id = pk2id(&secret_key.public_key(SECP256K1));
228 let hello = P2PMessage::Hello(HelloMessage {
229 protocol_version: ProtocolVersion::V5,
230 client_version: "reth/0.1.0".to_string(),
231 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
232 port: 30303,
233 id,
234 });
235
236 let mut hello_encoded = Vec::new();
237 hello.encode(&mut hello_encoded);
238
239 let hello_decoded = P2PMessage::decode(&mut &hello_encoded[..]).unwrap();
240
241 assert_eq!(hello, hello_decoded);
242 }
243
244 #[test]
245 fn hello_encoding_length() {
246 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
247 let id = pk2id(&secret_key.public_key(SECP256K1));
248 let hello = P2PMessage::Hello(HelloMessage {
249 protocol_version: ProtocolVersion::V5,
250 client_version: "reth/0.1.0".to_string(),
251 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
252 port: 30303,
253 id,
254 });
255
256 let mut hello_encoded = Vec::new();
257 hello.encode(&mut hello_encoded);
258
259 assert_eq!(hello_encoded.len(), hello.length());
260 }
261
262 #[test]
263 fn hello_message_id_prefix() {
264 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
266 let id = pk2id(&secret_key.public_key(SECP256K1));
267 let hello = P2PMessage::Hello(HelloMessage {
268 protocol_version: ProtocolVersion::V5,
269 client_version: "reth/0.1.0".to_string(),
270 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
271 port: 30303,
272 id,
273 });
274
275 let mut hello_encoded = Vec::new();
276 hello.encode(&mut hello_encoded);
277
278 assert_eq!(hello_encoded[0], EMPTY_STRING_CODE);
280 }
281}