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 pub fn with_snap(mut self, enable: bool) -> Self {
105 if enable {
106 let _ = self.try_add_protocol(Protocol::snap_2());
107 } else {
108 self.protocols.retain(|p| p.cap.name != "snap");
109 }
110 self
111 }
112}
113
114#[derive(Clone, Debug, PartialEq, Eq, RlpEncodable, RlpDecodable)]
120#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
121#[cfg_attr(any(test, feature = "arbitrary"), derive(arbitrary::Arbitrary))]
122#[add_arbitrary_tests(rlp)]
123pub struct HelloMessage {
124 pub protocol_version: ProtocolVersion,
126 pub client_version: String,
129 pub capabilities: Vec<Capability>,
131 pub port: u16,
133 pub id: PeerId,
135}
136
137impl HelloMessage {
140 pub const fn builder(id: PeerId) -> HelloMessageBuilder {
151 HelloMessageBuilder::new(id)
152 }
153}
154
155#[derive(Debug)]
157pub struct HelloMessageBuilder {
158 pub protocol_version: Option<ProtocolVersion>,
160 pub client_version: Option<String>,
163 pub protocols: Option<Vec<Protocol>>,
165 pub port: Option<u16>,
167 pub id: PeerId,
169}
170
171impl HelloMessageBuilder {
174 pub const fn new(id: PeerId) -> Self {
176 Self { protocol_version: None, client_version: None, protocols: None, port: None, id }
177 }
178
179 pub const fn port(mut self, port: u16) -> Self {
181 self.port = Some(port);
182 self
183 }
184
185 pub fn protocol(mut self, protocols: impl Into<Protocol>) -> Self {
187 self.protocols.get_or_insert_with(Vec::new).push(protocols.into());
188 self
189 }
190
191 pub fn protocols(mut self, protocols: impl IntoIterator<Item = Protocol>) -> Self {
193 self.protocols.get_or_insert_with(Vec::new).extend(protocols);
194 self
195 }
196
197 pub fn client_version(mut self, client_version: impl Into<String>) -> Self {
199 self.client_version = Some(client_version.into());
200 self
201 }
202
203 pub const fn protocol_version(mut self, protocol_version: ProtocolVersion) -> Self {
205 self.protocol_version = Some(protocol_version);
206 self
207 }
208
209 pub fn build(self) -> HelloMessageWithProtocols {
216 let Self { protocol_version, client_version, protocols, port, id } = self;
217 HelloMessageWithProtocols {
218 protocol_version: protocol_version.unwrap_or_default(),
219 client_version: client_version.unwrap_or_else(|| RETH_CLIENT_VERSION.to_string()),
220 protocols: protocols.unwrap_or_else(|| {
221 EthVersion::ALL_VERSIONS.iter().copied().map(Into::into).collect()
222 }),
223 port: port.unwrap_or(DEFAULT_TCP_PORT),
224 id,
225 }
226 }
227}
228
229#[cfg(test)]
230mod tests {
231 use crate::{
232 p2pstream::P2PMessage, Capability, EthVersion, HelloMessage, HelloMessageWithProtocols,
233 ProtocolVersion,
234 };
235 use alloy_rlp::{Decodable, Encodable, EMPTY_STRING_CODE};
236 use reth_network_peers::pk2id;
237 use secp256k1::{SecretKey, SECP256K1};
238
239 #[test]
240 fn test_hello_encoding_round_trip() {
241 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
242 let id = pk2id(&secret_key.public_key(SECP256K1));
243 let hello = P2PMessage::Hello(HelloMessage {
244 protocol_version: ProtocolVersion::V5,
245 client_version: "reth/0.1.0".to_string(),
246 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
247 port: 30303,
248 id,
249 });
250
251 let mut hello_encoded = Vec::new();
252 hello.encode(&mut hello_encoded);
253
254 let hello_decoded = P2PMessage::decode(&mut &hello_encoded[..]).unwrap();
255
256 assert_eq!(hello, hello_decoded);
257 }
258
259 #[test]
260 fn hello_encoding_length() {
261 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
262 let id = pk2id(&secret_key.public_key(SECP256K1));
263 let hello = P2PMessage::Hello(HelloMessage {
264 protocol_version: ProtocolVersion::V5,
265 client_version: "reth/0.1.0".to_string(),
266 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
267 port: 30303,
268 id,
269 });
270
271 let mut hello_encoded = Vec::new();
272 hello.encode(&mut hello_encoded);
273
274 assert_eq!(hello_encoded.len(), hello.length());
275 }
276 #[test]
279 fn test_default_protocols_still_include_eth69() {
280 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
282 let id = pk2id(&secret_key.public_key(SECP256K1));
283 let hello = HelloMessageWithProtocols::builder(id).build();
284
285 let has_eth69 = hello
286 .protocols
287 .iter()
288 .any(|p| p.cap.name == "eth" && p.cap.version == EthVersion::Eth69 as usize);
289 assert!(has_eth69, "Default protocols should include Eth69");
290 }
291
292 #[test]
293 fn hello_message_id_prefix() {
294 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
296 let id = pk2id(&secret_key.public_key(SECP256K1));
297 let hello = P2PMessage::Hello(HelloMessage {
298 protocol_version: ProtocolVersion::V5,
299 client_version: "reth/0.1.0".to_string(),
300 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
301 port: 30303,
302 id,
303 });
304
305 let mut hello_encoded = Vec::new();
306 hello.encode(&mut hello_encoded);
307
308 assert_eq!(hello_encoded[0], EMPTY_STRING_CODE);
310 }
311}