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 EthVersion::ALL_VERSIONS.iter().copied().map(Into::into).collect()
209 }),
210 port: port.unwrap_or(DEFAULT_TCP_PORT),
211 id,
212 }
213 }
214}
215
216#[cfg(test)]
217mod tests {
218 use crate::{
219 p2pstream::P2PMessage, Capability, EthVersion, HelloMessage, HelloMessageWithProtocols,
220 ProtocolVersion,
221 };
222 use alloy_rlp::{Decodable, Encodable, EMPTY_STRING_CODE};
223 use reth_network_peers::pk2id;
224 use secp256k1::{SecretKey, SECP256K1};
225
226 #[test]
227 fn test_hello_encoding_round_trip() {
228 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
229 let id = pk2id(&secret_key.public_key(SECP256K1));
230 let hello = P2PMessage::Hello(HelloMessage {
231 protocol_version: ProtocolVersion::V5,
232 client_version: "reth/0.1.0".to_string(),
233 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
234 port: 30303,
235 id,
236 });
237
238 let mut hello_encoded = Vec::new();
239 hello.encode(&mut hello_encoded);
240
241 let hello_decoded = P2PMessage::decode(&mut &hello_encoded[..]).unwrap();
242
243 assert_eq!(hello, hello_decoded);
244 }
245
246 #[test]
247 fn hello_encoding_length() {
248 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
249 let id = pk2id(&secret_key.public_key(SECP256K1));
250 let hello = P2PMessage::Hello(HelloMessage {
251 protocol_version: ProtocolVersion::V5,
252 client_version: "reth/0.1.0".to_string(),
253 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
254 port: 30303,
255 id,
256 });
257
258 let mut hello_encoded = Vec::new();
259 hello.encode(&mut hello_encoded);
260
261 assert_eq!(hello_encoded.len(), hello.length());
262 }
263
264 #[test]
265 fn test_default_protocols_include_eth69() {
266 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
268 let id = pk2id(&secret_key.public_key(SECP256K1));
269 let hello = HelloMessageWithProtocols::builder(id).build();
270
271 let has_eth69 = hello
272 .protocols
273 .iter()
274 .any(|p| p.cap.name == "eth" && p.cap.version == EthVersion::Eth69 as usize);
275 assert!(has_eth69, "Default protocols should include Eth69");
276 }
277
278 #[test]
279 fn hello_message_id_prefix() {
280 let secret_key = SecretKey::new(&mut rand_08::thread_rng());
282 let id = pk2id(&secret_key.public_key(SECP256K1));
283 let hello = P2PMessage::Hello(HelloMessage {
284 protocol_version: ProtocolVersion::V5,
285 client_version: "reth/0.1.0".to_string(),
286 capabilities: vec![Capability::new_static("eth", EthVersion::Eth67 as usize)],
287 port: 30303,
288 id,
289 });
290
291 let mut hello_encoded = Vec::new();
292 hello.encode(&mut hello_encoded);
293
294 assert_eq!(hello_encoded[0], EMPTY_STRING_CODE);
296 }
297}