Skip to main content

reth_discv4/
proto.rs

1//! Discovery v4 protocol implementation.
2
3use crate::{error::DecodePacketError, MAX_PACKET_SIZE, MIN_PACKET_SIZE};
4use alloy_primitives::{
5    bytes::{Buf, BufMut, Bytes, BytesMut},
6    keccak256, B256,
7};
8use alloy_rlp::{
9    Decodable, Encodable, Error as RlpError, Header, RlpDecodable, RlpEncodable,
10    RlpEncodableWrapper,
11};
12use enr::Enr;
13use reth_ethereum_forks::{EnrForkIdEntry, ForkId};
14use reth_network_peers::{pk2id, NodeRecord, PeerId};
15use secp256k1::{
16    ecdsa::{RecoverableSignature, RecoveryId},
17    SecretKey, SECP256K1,
18};
19use std::net::{IpAddr, Ipv4Addr};
20
21// Note: this is adapted from https://github.com/vorot93/discv4
22
23/// Represents the identifier for message variants.
24///
25/// This enumeration assigns unique identifiers (u8 values) to different message types.
26#[derive(Debug)]
27#[repr(u8)]
28pub enum MessageId {
29    /// Ping message identifier.
30    Ping = 1,
31    /// Pong message identifier.
32    Pong = 2,
33    /// Find node message identifier.
34    FindNode = 3,
35    /// Neighbours message identifier.
36    Neighbours = 4,
37    /// ENR request message identifier.
38    EnrRequest = 5,
39    /// ENR response message identifier.
40    EnrResponse = 6,
41}
42
43impl MessageId {
44    /// Converts the byte that represents the message id to the enum.
45    const fn from_u8(msg: u8) -> Result<Self, u8> {
46        Ok(match msg {
47            1 => Self::Ping,
48            2 => Self::Pong,
49            3 => Self::FindNode,
50            4 => Self::Neighbours,
51            5 => Self::EnrRequest,
52            6 => Self::EnrResponse,
53            _ => return Err(msg),
54        })
55    }
56}
57
58/// Enum representing various message types exchanged in the Discovery v4 protocol.
59#[derive(Debug, Eq, PartialEq)]
60pub enum Message {
61    /// Represents a ping message sent during liveness checks.
62    Ping(Ping),
63    /// Represents a pong message, which is a reply to a PING message.
64    Pong(Pong),
65    /// Represents a query for nodes in the given bucket.
66    FindNode(FindNode),
67    /// Represents a neighbour message, providing information about nearby nodes.
68    Neighbours(Neighbours),
69    /// Represents an ENR request message, a request for Ethereum Node Records (ENR) as per [EIP-778](https://eips.ethereum.org/EIPS/eip-778).
70    EnrRequest(EnrRequest),
71    /// Represents an ENR response message, a response to an ENR request with Ethereum Node Records (ENR) as per [EIP-778](https://eips.ethereum.org/EIPS/eip-778).
72    EnrResponse(EnrResponse),
73}
74
75// === impl Message ===
76
77impl Message {
78    /// Returns the id for this type
79    pub const fn msg_type(&self) -> MessageId {
80        match self {
81            Self::Ping(_) => MessageId::Ping,
82            Self::Pong(_) => MessageId::Pong,
83            Self::FindNode(_) => MessageId::FindNode,
84            Self::Neighbours(_) => MessageId::Neighbours,
85            Self::EnrRequest(_) => MessageId::EnrRequest,
86            Self::EnrResponse(_) => MessageId::EnrResponse,
87        }
88    }
89
90    /// Encodes the UDP datagram, See <https://github.com/ethereum/devp2p/blob/master/discv4.md#wire-protocol>
91    ///
92    /// The datagram is `header || payload`
93    /// where header is `hash || signature || packet-type`
94    pub fn encode(&self, secret_key: &SecretKey) -> (Bytes, B256) {
95        // allocate max packet size
96        let mut datagram = BytesMut::with_capacity(MAX_PACKET_SIZE);
97
98        // since signature has fixed len, we can split and fill the datagram buffer at fixed
99        // positions, this way we can encode the message directly in the datagram buffer
100        let mut sig_bytes = datagram.split_off(B256::len_bytes());
101        let mut payload = sig_bytes.split_off(secp256k1::constants::COMPACT_SIGNATURE_SIZE + 1);
102
103        // Put the message type at the beginning of the payload
104        payload.put_u8(self.msg_type() as u8);
105
106        // Match the message type and encode the corresponding message into the payload
107        match self {
108            Self::Ping(message) => message.encode(&mut payload),
109            Self::Pong(message) => message.encode(&mut payload),
110            Self::FindNode(message) => message.encode(&mut payload),
111            Self::Neighbours(message) => message.encode(&mut payload),
112            Self::EnrRequest(message) => message.encode(&mut payload),
113            Self::EnrResponse(message) => message.encode(&mut payload),
114        }
115
116        // Sign the payload with the secret key using recoverable ECDSA
117        let signature: RecoverableSignature = SECP256K1.sign_ecdsa_recoverable(
118            &secp256k1::Message::from_digest(keccak256(&payload).0),
119            secret_key,
120        );
121
122        // Serialize the signature and append it to the signature bytes
123        let (rec, sig) = signature.serialize_compact();
124        sig_bytes.extend_from_slice(&sig);
125        sig_bytes.put_u8(i32::from(rec) as u8);
126        sig_bytes.unsplit(payload);
127
128        // Calculate the hash of the signature bytes and append it to the datagram
129        let hash = keccak256(&sig_bytes);
130        datagram.extend_from_slice(hash.as_slice());
131
132        // Append the signature bytes to the datagram
133        datagram.unsplit(sig_bytes);
134
135        // Return the frozen datagram and the hash
136        (datagram.freeze(), hash)
137    }
138
139    /// Decodes the [`Message`] from the given buffer.
140    ///
141    /// Returns the decoded message and the public key of the sender.
142    pub fn decode(packet: &[u8]) -> Result<Packet, DecodePacketError> {
143        if packet.len() < MIN_PACKET_SIZE {
144            return Err(DecodePacketError::PacketTooShort)
145        }
146
147        // parses the wire-protocol, every packet starts with a header:
148        // packet-header = hash || signature || packet-type
149        // hash = keccak256(signature || packet-type || packet-data)
150        // signature = sign(packet-type || packet-data)
151
152        let header_hash = keccak256(&packet[32..]);
153        let data_hash = B256::from_slice(&packet[..32]);
154        if data_hash != header_hash {
155            return Err(DecodePacketError::HashMismatch)
156        }
157
158        // Resolve the message type before recovering the public key: recovery is by far the most
159        // expensive step of decoding, and a packet we have no handler for is rejected either way.
160        let message_id =
161            MessageId::from_u8(packet[97]).map_err(DecodePacketError::UnknownMessage)?;
162
163        let signature = &packet[32..96];
164        let recovery_id = RecoveryId::try_from(packet[96] as i32)?;
165        let recoverable_sig = RecoverableSignature::from_compact(signature, recovery_id)?;
166
167        // recover the public key
168        let msg = secp256k1::Message::from_digest(keccak256(&packet[97..]).0);
169
170        let pk = SECP256K1.recover_ecdsa(&msg, &recoverable_sig)?;
171        let node_id = pk2id(&pk);
172
173        let payload = &mut &packet[98..];
174
175        let msg = match message_id {
176            MessageId::Ping => Self::Ping(Ping::decode(payload)?),
177            MessageId::Pong => Self::Pong(Pong::decode(payload)?),
178            MessageId::FindNode => Self::FindNode(FindNode::decode(payload)?),
179            MessageId::Neighbours => Self::Neighbours(Neighbours::decode(payload)?),
180            MessageId::EnrRequest => Self::EnrRequest(EnrRequest::decode(payload)?),
181            MessageId::EnrResponse => Self::EnrResponse(EnrResponse::decode(payload)?),
182        };
183
184        Ok(Packet { msg, node_id, hash: header_hash })
185    }
186}
187
188/// Represents a decoded packet.
189///
190/// This struct holds information about a decoded packet, including the message, node ID, and hash.
191#[derive(Debug)]
192pub struct Packet {
193    /// The decoded message from the packet.
194    pub msg: Message,
195    /// The ID of the peer that sent the packet.
196    pub node_id: PeerId,
197    /// The hash of the packet.
198    pub hash: B256,
199}
200
201/// Represents the `from` field in the `Ping` packet
202#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, RlpEncodableWrapper)]
203struct PingNodeEndpoint(NodeEndpoint);
204
205impl alloy_rlp::Decodable for PingNodeEndpoint {
206    #[inline]
207    fn decode(b: &mut &[u8]) -> alloy_rlp::Result<Self> {
208        let alloy_rlp::Header { list, payload_length } = alloy_rlp::Header::decode(b)?;
209        if !list {
210            return Err(alloy_rlp::Error::UnexpectedString);
211        }
212        // Payload length checked by Header::decode.
213        let started_len = b.len();
214
215        // Geth allows the ipaddr to be possibly empty:
216        // <https://github.com/ethereum/go-ethereum/blob/380688c636a654becc8f114438c2a5d93d2db032/p2p/discover/v4_udp.go#L206-L209>
217        // <https://github.com/ethereum/go-ethereum/blob/380688c636a654becc8f114438c2a5d93d2db032/p2p/enode/node.go#L189-L189>
218        //
219        // Therefore, if we see an empty list instead of a properly formed `IpAddr`, we will
220        // instead use `IpV4Addr::UNSPECIFIED`
221        let address =
222            if *b.first().ok_or(alloy_rlp::Error::InputTooShort)? == alloy_rlp::EMPTY_STRING_CODE {
223                let addr = IpAddr::V4(Ipv4Addr::UNSPECIFIED);
224                b.advance(1);
225                addr
226            } else {
227                alloy_rlp::Decodable::decode(b)?
228            };
229
230        let this = NodeEndpoint {
231            address,
232            udp_port: alloy_rlp::Decodable::decode(b)?,
233            tcp_port: alloy_rlp::Decodable::decode(b)?,
234        };
235        let consumed = started_len - b.len();
236        if consumed != payload_length {
237            return Err(alloy_rlp::Error::ListLengthMismatch {
238                expected: payload_length,
239                got: consumed,
240            });
241        }
242        Ok(Self(this))
243    }
244}
245
246/// Represents the `from`, `to` fields in the packets
247#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, RlpEncodable, RlpDecodable)]
248pub struct NodeEndpoint {
249    /// The IP address of the network endpoint. It can be either IPv4 or IPv6.
250    pub address: IpAddr,
251    /// The UDP port used for communication in the discovery protocol.
252    pub udp_port: u16,
253    /// The TCP port used for communication in the `RLPx` protocol.
254    pub tcp_port: u16,
255}
256
257impl From<NodeRecord> for NodeEndpoint {
258    fn from(NodeRecord { address, tcp_port, udp_port, .. }: NodeRecord) -> Self {
259        Self { address, tcp_port, udp_port }
260    }
261}
262
263impl NodeEndpoint {
264    /// Creates a new [`NodeEndpoint`] from a given UDP address and TCP port.
265    pub const fn from_udp_address(udp_address: &std::net::SocketAddr, tcp_port: u16) -> Self {
266        Self { address: udp_address.ip(), udp_port: udp_address.port(), tcp_port }
267    }
268}
269
270/// A [FindNode packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#findnode-packet-0x03).
271#[derive(Clone, Copy, Debug, Eq, PartialEq, RlpEncodable)]
272pub struct FindNode {
273    /// The target node's ID, a 64-byte secp256k1 public key.
274    pub id: PeerId,
275    /// The expiration timestamp of the packet, an absolute UNIX time stamp.
276    pub expire: u64,
277}
278
279impl Decodable for FindNode {
280    // NOTE(onbjerg): Manual implementation to satisfy EIP-8.
281    //
282    // See https://eips.ethereum.org/EIPS/eip-8
283    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
284        let b = &mut &**buf;
285        let rlp_head = Header::decode(b)?;
286        if !rlp_head.list {
287            return Err(RlpError::UnexpectedString)
288        }
289        let started_len = b.len();
290
291        let this = Self { id: Decodable::decode(b)?, expire: Decodable::decode(b)? };
292
293        // NOTE(onbjerg): Because of EIP-8, we only check that we did not consume *more* than the
294        // payload length, i.e. it is ok if payload length is greater than what we consumed, as we
295        // just discard the remaining list items
296        let consumed = started_len - b.len();
297        if consumed > rlp_head.payload_length {
298            return Err(RlpError::ListLengthMismatch {
299                expected: rlp_head.payload_length,
300                got: consumed,
301            })
302        }
303
304        let rem = rlp_head.payload_length - consumed;
305        b.advance(rem);
306        *buf = *b;
307
308        Ok(this)
309    }
310}
311
312/// A [Neighbours packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#neighbors-packet-0x04).
313#[derive(Clone, Debug, Eq, PartialEq, RlpEncodable)]
314pub struct Neighbours {
315    /// The list of nodes containing IP, UDP port, TCP port, and node ID.
316    pub nodes: Vec<NodeRecord>,
317    /// The expiration timestamp of the packet, an absolute UNIX time stamp.
318    pub expire: u64,
319}
320
321impl Decodable for Neighbours {
322    // NOTE(onbjerg): Manual implementation to satisfy EIP-8.
323    //
324    // See https://eips.ethereum.org/EIPS/eip-8
325    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
326        let b = &mut &**buf;
327        let rlp_head = Header::decode(b)?;
328        if !rlp_head.list {
329            return Err(RlpError::UnexpectedString)
330        }
331        let started_len = b.len();
332
333        let this = Self { nodes: Decodable::decode(b)?, expire: Decodable::decode(b)? };
334
335        // NOTE(onbjerg): Because of EIP-8, we only check that we did not consume *more* than the
336        // payload length, i.e. it is ok if payload length is greater than what we consumed, as we
337        // just discard the remaining list items
338        let consumed = started_len - b.len();
339        if consumed > rlp_head.payload_length {
340            return Err(RlpError::ListLengthMismatch {
341                expected: rlp_head.payload_length,
342                got: consumed,
343            })
344        }
345
346        let rem = rlp_head.payload_length - consumed;
347        b.advance(rem);
348        *buf = *b;
349
350        Ok(this)
351    }
352}
353
354/// A [ENRRequest packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#enrrequest-packet-0x05).
355///
356/// This packet is used to request the current version of a node's Ethereum Node Record (ENR).
357#[derive(Clone, Copy, Debug, Eq, PartialEq, RlpEncodable)]
358pub struct EnrRequest {
359    /// The expiration timestamp for the request. No reply should be sent if it refers to a time in
360    /// the past.
361    pub expire: u64,
362}
363
364impl Decodable for EnrRequest {
365    // NOTE(onbjerg): Manual implementation to satisfy EIP-8.
366    //
367    // See https://eips.ethereum.org/EIPS/eip-8
368    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
369        let b = &mut &**buf;
370        let rlp_head = Header::decode(b)?;
371        if !rlp_head.list {
372            return Err(RlpError::UnexpectedString)
373        }
374        let started_len = b.len();
375
376        let this = Self { expire: Decodable::decode(b)? };
377
378        // NOTE(onbjerg): Because of EIP-8, we only check that we did not consume *more* than the
379        // payload length, i.e. it is ok if payload length is greater than what we consumed, as we
380        // just discard the remaining list items
381        let consumed = started_len - b.len();
382        if consumed > rlp_head.payload_length {
383            return Err(RlpError::ListLengthMismatch {
384                expected: rlp_head.payload_length,
385                got: consumed,
386            })
387        }
388
389        let rem = rlp_head.payload_length - consumed;
390        b.advance(rem);
391        *buf = *b;
392
393        Ok(this)
394    }
395}
396
397/// A [ENRResponse packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#enrresponse-packet-0x06).
398///
399/// This packet is used to respond to an `ENRRequest` packet and includes the requested ENR along
400/// with the hash of the original request.
401#[derive(Clone, Debug, Eq, PartialEq, RlpEncodable, RlpDecodable)]
402pub struct EnrResponse {
403    /// The hash of the `ENRRequest` packet being replied to.
404    pub request_hash: B256,
405    /// The ENR (Ethereum Node Record) for the responding node.
406    pub enr: Enr<SecretKey>,
407}
408
409// === impl EnrResponse ===
410
411impl EnrResponse {
412    /// Returns the [`ForkId`] if set
413    ///
414    /// See also <https://github.com/ethereum/go-ethereum/blob/9244d5cd61f3ea5a7645fdf2a1a96d53421e412f/eth/protocols/eth/discovery.go#L36>
415    pub fn eth_fork_id(&self) -> Option<ForkId> {
416        let mut maybe_fork_id = self.enr.get_raw_rlp(b"eth")?;
417        EnrForkIdEntry::decode(&mut maybe_fork_id).ok().map(Into::into)
418    }
419}
420
421/// Represents a Ping packet.
422///
423/// A [Ping packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#ping-packet-0x01).
424#[derive(Debug, Clone, Eq, PartialEq)]
425pub struct Ping {
426    /// The sender's endpoint.
427    pub from: NodeEndpoint,
428    /// The recipient's endpoint.
429    pub to: NodeEndpoint,
430    /// The expiration timestamp.
431    pub expire: u64,
432    /// Optional `enr_seq` for <https://eips.ethereum.org/EIPS/eip-868>
433    pub enr_sq: Option<u64>,
434}
435
436impl Encodable for Ping {
437    fn length(&self) -> usize {
438        let mut payload_length =
439            4u32.length() + self.from.length() + self.to.length() + self.expire.length();
440        if let Some(enr_seq) = self.enr_sq {
441            payload_length += enr_seq.length();
442        }
443        payload_length + alloy_rlp::length_of_length(payload_length)
444    }
445
446    fn encode(&self, out: &mut dyn BufMut) {
447        #[derive(RlpEncodable)]
448        struct V4PingMessage<'a> {
449            version: u32,
450            from: &'a NodeEndpoint,
451            to: &'a NodeEndpoint,
452            expire: u64,
453        }
454
455        #[derive(RlpEncodable)]
456        struct V4PingMessageEIP868<'a> {
457            version: u32,
458            from: &'a NodeEndpoint,
459            to: &'a NodeEndpoint,
460            expire: u64,
461            enr_seq: u64,
462        }
463        if let Some(enr_seq) = self.enr_sq {
464            V4PingMessageEIP868 {
465                version: 4, // version 4
466                from: &self.from,
467                to: &self.to,
468                expire: self.expire,
469                enr_seq,
470            }
471            .encode(out);
472        } else {
473            V4PingMessage {
474                version: 4, // version 4
475                from: &self.from,
476                to: &self.to,
477                expire: self.expire,
478            }
479            .encode(out);
480        }
481    }
482}
483
484impl Decodable for Ping {
485    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
486        let b = &mut &**buf;
487        let rlp_head = Header::decode(b)?;
488        if !rlp_head.list {
489            return Err(RlpError::UnexpectedString)
490        }
491        let started_len = b.len();
492
493        // > Implementations should ignore any mismatches in version:
494        // <https://github.com/ethereum/devp2p/blob/master/discv4.md#ping-packet-0x01>
495        let _version = u32::decode(b)?;
496
497        // see `Decodable` implementation in `PingNodeEndpoint` for why this is needed
498        let from = PingNodeEndpoint::decode(b)?.0;
499
500        let mut this =
501            Self { from, to: Decodable::decode(b)?, expire: Decodable::decode(b)?, enr_sq: None };
502
503        // only decode the ENR sequence if there's more data in the datagram to decode else skip
504        if b.has_remaining() {
505            this.enr_sq = Some(Decodable::decode(b)?);
506        }
507
508        let consumed = started_len - b.len();
509        if consumed > rlp_head.payload_length {
510            return Err(RlpError::ListLengthMismatch {
511                expected: rlp_head.payload_length,
512                got: consumed,
513            })
514        }
515        let rem = rlp_head.payload_length - consumed;
516        b.advance(rem);
517        *buf = *b;
518        Ok(this)
519    }
520}
521
522/// Represents a Pong packet.
523///
524/// A [Pong packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#pong-packet-0x02).
525#[derive(Clone, Debug, Eq, PartialEq)]
526pub struct Pong {
527    /// The recipient's endpoint.
528    pub to: NodeEndpoint,
529    /// The hash of the corresponding ping packet.
530    pub echo: B256,
531    /// The expiration timestamp.
532    pub expire: u64,
533    /// Optional `enr_seq` for <https://eips.ethereum.org/EIPS/eip-868>
534    pub enr_sq: Option<u64>,
535}
536
537impl Encodable for Pong {
538    fn length(&self) -> usize {
539        let mut payload_length = self.to.length() + self.echo.length() + self.expire.length();
540        if let Some(enr_seq) = self.enr_sq {
541            payload_length += enr_seq.length();
542        }
543        payload_length + alloy_rlp::length_of_length(payload_length)
544    }
545
546    fn encode(&self, out: &mut dyn BufMut) {
547        #[derive(RlpEncodable)]
548        struct PongMessageEIP868<'a> {
549            to: &'a NodeEndpoint,
550            echo: &'a B256,
551            expire: u64,
552            enr_seq: u64,
553        }
554
555        #[derive(RlpEncodable)]
556        struct PongMessage<'a> {
557            to: &'a NodeEndpoint,
558            echo: &'a B256,
559            expire: u64,
560        }
561
562        if let Some(enr_seq) = self.enr_sq {
563            PongMessageEIP868 { to: &self.to, echo: &self.echo, expire: self.expire, enr_seq }
564                .encode(out);
565        } else {
566            PongMessage { to: &self.to, echo: &self.echo, expire: self.expire }.encode(out);
567        }
568    }
569}
570
571impl Decodable for Pong {
572    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
573        let b = &mut &**buf;
574        let rlp_head = Header::decode(b)?;
575        if !rlp_head.list {
576            return Err(RlpError::UnexpectedString)
577        }
578        let started_len = b.len();
579        let mut this = Self {
580            to: Decodable::decode(b)?,
581            echo: Decodable::decode(b)?,
582            expire: Decodable::decode(b)?,
583            enr_sq: None,
584        };
585
586        // only decode the ENR sequence if there's more data in the datagram to decode else skip
587        if b.has_remaining() {
588            this.enr_sq = Some(Decodable::decode(b)?);
589        }
590
591        let consumed = started_len - b.len();
592        if consumed > rlp_head.payload_length {
593            return Err(RlpError::ListLengthMismatch {
594                expected: rlp_head.payload_length,
595                got: consumed,
596            })
597        }
598        let rem = rlp_head.payload_length - consumed;
599        b.advance(rem);
600        *buf = *b;
601
602        Ok(this)
603    }
604}
605
606#[cfg(test)]
607mod tests {
608    use super::*;
609    use crate::{
610        test_utils::{rng_endpoint, rng_ipv4_record, rng_ipv6_record, rng_message},
611        DEFAULT_DISCOVERY_PORT, SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS,
612    };
613    use alloy_primitives::hex;
614    use assert_matches::assert_matches;
615    use enr::EnrPublicKey;
616    use rand_08::{thread_rng as rng, Rng, RngCore};
617    use reth_ethereum_forks::ForkHash;
618
619    #[test]
620    fn test_endpoint_ipv_v4() {
621        let mut rng = rng();
622        for _ in 0..100 {
623            let mut ip = [0u8; 4];
624            rng.fill_bytes(&mut ip);
625            let msg = NodeEndpoint {
626                address: IpAddr::V4(ip.into()),
627                tcp_port: rng.r#gen(),
628                udp_port: rng.r#gen(),
629            };
630
631            let decoded = NodeEndpoint::decode(&mut alloy_rlp::encode(msg).as_slice()).unwrap();
632            assert_eq!(msg, decoded);
633        }
634    }
635
636    #[test]
637    fn test_endpoint_ipv_64() {
638        let mut rng = rng();
639        for _ in 0..100 {
640            let mut ip = [0u8; 16];
641            rng.fill_bytes(&mut ip);
642            let msg = NodeEndpoint {
643                address: IpAddr::V6(ip.into()),
644                tcp_port: rng.r#gen(),
645                udp_port: rng.r#gen(),
646            };
647
648            let decoded = NodeEndpoint::decode(&mut alloy_rlp::encode(msg).as_slice()).unwrap();
649            assert_eq!(msg, decoded);
650        }
651    }
652
653    #[test]
654    fn test_ping_message() {
655        let mut rng = rng();
656        for _ in 0..100 {
657            let mut ip = [0u8; 16];
658            rng.fill_bytes(&mut ip);
659            let msg = Ping {
660                from: rng_endpoint(&mut rng),
661                to: rng_endpoint(&mut rng),
662                expire: 0,
663                enr_sq: None,
664            };
665
666            let encoded = alloy_rlp::encode(&msg);
667            assert_eq!(msg.length(), encoded.len());
668            let mut buf = encoded.as_slice();
669            let decoded = Ping::decode(&mut buf).unwrap();
670            assert_eq!(msg, decoded);
671            assert!(buf.is_empty());
672        }
673    }
674
675    #[test]
676    fn test_ping_message_with_enr() {
677        let mut rng = rng();
678        for _ in 0..100 {
679            let mut ip = [0u8; 16];
680            rng.fill_bytes(&mut ip);
681            let msg = Ping {
682                from: rng_endpoint(&mut rng),
683                to: rng_endpoint(&mut rng),
684                expire: 0,
685                enr_sq: Some(rng.r#gen()),
686            };
687
688            let encoded = alloy_rlp::encode(&msg);
689            assert_eq!(msg.length(), encoded.len());
690            let mut buf = encoded.as_slice();
691            let decoded = Ping::decode(&mut buf).unwrap();
692            assert_eq!(msg, decoded);
693            assert!(buf.is_empty());
694        }
695    }
696
697    #[test]
698    fn test_pong_message() {
699        let mut rng = rng();
700        for _ in 0..100 {
701            let mut ip = [0u8; 16];
702            rng.fill_bytes(&mut ip);
703            let msg = Pong {
704                to: rng_endpoint(&mut rng),
705                echo: B256::random(),
706                expire: rng.r#gen(),
707                enr_sq: None,
708            };
709
710            let encoded = alloy_rlp::encode(&msg);
711            assert_eq!(msg.length(), encoded.len());
712            let mut buf = encoded.as_slice();
713            let decoded = Pong::decode(&mut buf).unwrap();
714            assert_eq!(msg, decoded);
715            assert!(buf.is_empty());
716        }
717    }
718
719    #[test]
720    fn test_pong_message_with_enr() {
721        let mut rng = rng();
722        for _ in 0..100 {
723            let mut ip = [0u8; 16];
724            rng.fill_bytes(&mut ip);
725            let msg = Pong {
726                to: rng_endpoint(&mut rng),
727                echo: B256::random(),
728                expire: rng.r#gen(),
729                enr_sq: Some(rng.r#gen()),
730            };
731
732            let encoded = alloy_rlp::encode(&msg);
733            assert_eq!(msg.length(), encoded.len());
734            let mut buf = encoded.as_slice();
735            let decoded = Pong::decode(&mut buf).unwrap();
736            assert_eq!(msg, decoded);
737            assert!(buf.is_empty());
738        }
739    }
740
741    #[test]
742    fn test_hash_mismatch() {
743        let mut rng = rng();
744        let msg = rng_message(&mut rng);
745        let (secret_key, _) = SECP256K1.generate_keypair(&mut rng);
746        let (buf, _) = msg.encode(&secret_key);
747
748        let mut buf_vec = buf.to_vec();
749        buf_vec.push(0);
750        match Message::decode(buf_vec.as_slice()).unwrap_err() {
751            DecodePacketError::HashMismatch => {}
752            err => {
753                unreachable!("unexpected err {}", err)
754            }
755        }
756    }
757
758    #[test]
759    fn neighbours_max_ipv4() {
760        let mut rng = rng();
761        let msg = Message::Neighbours(Neighbours {
762            nodes: std::iter::repeat_with(|| rng_ipv4_record(&mut rng)).take(16).collect(),
763            expire: rng.r#gen(),
764        });
765        let (secret_key, _) = SECP256K1.generate_keypair(&mut rng);
766
767        let (encoded, _) = msg.encode(&secret_key);
768        // Assert that 16 nodes never fit into one packet
769        assert!(encoded.len() > MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
770    }
771
772    #[test]
773    fn neighbours_max_nodes() {
774        let mut rng = rng();
775        for _ in 0..1000 {
776            let msg = Message::Neighbours(Neighbours {
777                nodes: std::iter::repeat_with(|| rng_ipv6_record(&mut rng))
778                    .take(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS)
779                    .collect(),
780                expire: rng.r#gen(),
781            });
782            let (secret_key, _) = SECP256K1.generate_keypair(&mut rng);
783
784            let (encoded, _) = msg.encode(&secret_key);
785            assert!(encoded.len() <= MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
786
787            let mut neighbours = Neighbours {
788                nodes: std::iter::repeat_with(|| rng_ipv6_record(&mut rng))
789                    .take(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS - 1)
790                    .collect(),
791                expire: rng.r#gen(),
792            };
793            neighbours.nodes.push(rng_ipv4_record(&mut rng));
794            let msg = Message::Neighbours(neighbours);
795            let (encoded, _) = msg.encode(&secret_key);
796            assert!(encoded.len() <= MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
797        }
798    }
799
800    #[test]
801    fn test_encode_decode_message() {
802        let mut rng = rng();
803        for _ in 0..100 {
804            let msg = rng_message(&mut rng);
805            let (secret_key, pk) = SECP256K1.generate_keypair(&mut rng);
806            let sender_id = pk2id(&pk);
807
808            let (buf, _) = msg.encode(&secret_key);
809
810            let packet = Message::decode(buf.as_ref()).unwrap();
811
812            assert_eq!(msg, packet.msg);
813            assert_eq!(sender_id, packet.node_id);
814        }
815    }
816
817    #[test]
818    fn decode_pong_packet() {
819        let packet = "2ad84c37327a06c2522cf7bc039621da89f68907441b755935bb308dc4cd17d6fe550e90329ad6a516ca7db18e08900067928a0dfa3b5c75d55a42c984497373698d98616662c048983ea85895ea2da765eabeb15525478384e106337bfd8ed50002f3c9843ed8cae682fd1c80a008ad4dead0922211df47593e7d837b2b23d13954285871ca23250ea594993ded84635690e5829670";
820        let data = hex::decode(packet).unwrap();
821        Message::decode(&data).unwrap();
822    }
823    #[test]
824    fn decode_ping_packet() {
825        let packet = "05ae5bf922cf2a93f97632a4ab0943dc252a0dab0c42d86dd62e5d91e1a0966e9b628fbf4763fdfbb928540460b797e6be2e7058a82f6083f6d2e7391bb021741459976d4152aa16bbee0c3609dcfac6668db1ef78b7ee9f8b4ced10dd5ae2900101df04cb8403d12d4f82765f82765fc9843ed8cae6828aa6808463569916829670";
826        let data = hex::decode(packet).unwrap();
827        Message::decode(&data).unwrap();
828    }
829
830    #[test]
831    fn encode_decode_enr_msg() {
832        use alloy_rlp::Decodable;
833        use enr::secp256k1::SecretKey;
834        use std::net::Ipv4Addr;
835
836        let mut rng = rand_08::rngs::OsRng;
837        let key = SecretKey::new(&mut rng);
838        let ip = Ipv4Addr::new(127, 0, 0, 1);
839        let tcp = 3000;
840
841        let fork_id: ForkId = ForkId { hash: ForkHash([220, 233, 108, 45]), next: 0u64 };
842
843        let enr = {
844            let mut builder = Enr::builder();
845            builder.ip(ip.into());
846            builder.tcp4(tcp);
847            let mut buf = Vec::new();
848            let forkentry = EnrForkIdEntry { fork_id };
849            forkentry.encode(&mut buf);
850            builder.add_value_rlp("eth", buf.into());
851            builder.build(&key).unwrap()
852        };
853
854        let enr_response = EnrResponse { request_hash: B256::random(), enr };
855
856        let mut buf = Vec::new();
857        enr_response.encode(&mut buf);
858
859        let decoded = EnrResponse::decode(&mut &buf[..]).unwrap();
860
861        let fork_id_decoded = decoded.eth_fork_id().unwrap();
862        assert_eq!(fork_id, fork_id_decoded);
863    }
864
865    // test vector from the enr library rlp encoding tests
866    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#L1019>
867
868    #[test]
869    fn encode_known_rlp_enr() {
870        use alloy_rlp::Decodable;
871        use enr::{secp256k1::SecretKey, EnrPublicKey};
872        use std::net::Ipv4Addr;
873
874        let valid_record = hex!(
875            "f884b8407098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c01826964827634826970847f00000189736563703235366b31a103ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31388375647082765f"
876        );
877        let signature = hex!(
878            "7098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c"
879        );
880        let expected_pubkey =
881            hex!("03ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd3138");
882
883        let enr = Enr::<SecretKey>::decode(&mut &valid_record[..]).unwrap();
884        let pubkey = enr.public_key().encode();
885
886        assert_eq!(enr.ip4(), Some(Ipv4Addr::new(127, 0, 0, 1)));
887        assert_eq!(enr.id(), Some(String::from("v4")));
888        assert_eq!(enr.udp4(), Some(DEFAULT_DISCOVERY_PORT));
889        assert_eq!(enr.tcp4(), None);
890        assert_eq!(enr.signature(), &signature[..]);
891        assert_eq!(pubkey.to_vec(), expected_pubkey);
892        assert!(enr.verify());
893
894        assert_eq!(&alloy_rlp::encode(&enr)[..], &valid_record[..]);
895
896        // ensure the length is equal
897        assert_eq!(enr.length(), valid_record.len());
898    }
899
900    // test vector from the enr library rlp encoding tests
901    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#L1019>
902    #[test]
903    fn decode_enr_rlp() {
904        use enr::secp256k1::SecretKey;
905        use std::net::Ipv4Addr;
906
907        let valid_record = hex!(
908            "f884b8407098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c01826964827634826970847f00000189736563703235366b31a103ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31388375647082765f"
909        );
910        let signature = hex!(
911            "7098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c"
912        );
913        let expected_pubkey =
914            hex!("03ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd3138");
915
916        let mut valid_record_buf = valid_record.as_slice();
917        let enr = Enr::<SecretKey>::decode(&mut valid_record_buf).unwrap();
918        let pubkey = enr.public_key().encode();
919
920        // Byte array must be consumed after enr has finished decoding
921        assert!(valid_record_buf.is_empty());
922
923        assert_eq!(enr.ip4(), Some(Ipv4Addr::new(127, 0, 0, 1)));
924        assert_eq!(enr.id(), Some(String::from("v4")));
925        assert_eq!(enr.udp4(), Some(DEFAULT_DISCOVERY_PORT));
926        assert_eq!(enr.tcp4(), None);
927        assert_eq!(enr.signature(), &signature[..]);
928        assert_eq!(pubkey.to_vec(), expected_pubkey);
929        assert!(enr.verify());
930    }
931
932    // test for failing message decode
933    #[test]
934    fn decode_failing_packet() {
935        let packet = hex!(
936            "2467ab56952aedf4cfb8bb7830ddc8922d0f992185229919dad9de3841fe95d9b3a7b52459398235f6d3805644666d908b45edb3670414ed97f357afba51f71f7d35c1f45878ba732c3868b04ca42ff0ed347c99efcf3a5768afed68eb21ef960001db04c3808080c9840a480e8f82765f808466a9a06386019106833efe"
937        );
938
939        let _message = Message::decode(&packet[..]).unwrap();
940    }
941
942    // test for failing message decode
943    #[test]
944    fn decode_node() {
945        let packet = hex!("cb840000000082115c82115d");
946        let _message = NodeEndpoint::decode(&mut &packet[..]).unwrap();
947    }
948
949    // test vector from the enr library rlp encoding tests
950    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#LL1206C35-L1206C35>
951    #[test]
952    fn encode_decode_enr_rlp() {
953        use enr::{secp256k1::SecretKey, EnrPublicKey};
954        use std::net::Ipv4Addr;
955
956        let key = SecretKey::new(&mut rand_08::rngs::OsRng);
957        let ip = Ipv4Addr::new(127, 0, 0, 1);
958        let tcp = 3000;
959
960        let enr = {
961            let mut builder = Enr::builder();
962            builder.ip(ip.into());
963            builder.tcp4(tcp);
964            builder.build(&key).unwrap()
965        };
966
967        let mut encoded_bytes = &alloy_rlp::encode(&enr)[..];
968        let decoded_enr = Enr::<SecretKey>::decode(&mut encoded_bytes).unwrap();
969
970        // Byte array must be consumed after enr has finished decoding
971        assert!(encoded_bytes.is_empty());
972
973        assert_eq!(decoded_enr, enr);
974        assert_eq!(decoded_enr.id(), Some("v4".into()));
975        assert_eq!(decoded_enr.ip4(), Some(ip));
976        assert_eq!(decoded_enr.tcp4(), Some(tcp));
977        assert_eq!(
978            decoded_enr.public_key().encode(),
979            key.public_key(secp256k1::SECP256K1).encode()
980        );
981        assert!(decoded_enr.verify());
982    }
983
984    mod eip8 {
985        use super::*;
986
987        fn junk_enr_request() -> Vec<u8> {
988            let mut buf = Vec::new();
989            // enr request is just an expiration
990            let expire: u64 = 123456;
991
992            // add some junk
993            let junk: u64 = 112233;
994
995            // rlp header encoding
996            let payload_length = expire.length() + junk.length();
997            alloy_rlp::Header { list: true, payload_length }.encode(&mut buf);
998
999            // fields
1000            expire.encode(&mut buf);
1001            junk.encode(&mut buf);
1002
1003            buf
1004        }
1005
1006        // checks that junk data at the end of the packet is discarded according to eip-8
1007        #[test]
1008        fn eip8_decode_enr_request() {
1009            let enr_request_with_junk = junk_enr_request();
1010
1011            let mut buf = enr_request_with_junk.as_slice();
1012            let decoded = EnrRequest::decode(&mut buf).unwrap();
1013            assert_eq!(decoded.expire, 123456);
1014        }
1015
1016        // checks that junk data at the end of the packet is discarded according to eip-8
1017        //
1018        // test vector from eip-8: https://eips.ethereum.org/EIPS/eip-8
1019        #[test]
1020        fn eip8_decode_findnode() {
1021            let findnode_with_junk = hex!(
1022                "c7c44041b9f7c7e41934417ebac9a8e1a4c6298f74553f2fcfdcae6ed6fe53163eb3d2b52e39fe91831b8a927bf4fc222c3902202027e5e9eb812195f95d20061ef5cd31d502e47ecb61183f74a504fe04c51e73df81f25c4d506b26db4517490103f84eb840ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31387574077f301b421bc84df7266c44e9e6d569fc56be00812904767bf5ccd1fc7f8443b9a35582999983999999280dc62cc8255c73471e0a61da0c89acdc0e035e260add7fc0c04ad9ebf3919644c91cb247affc82b69bd2ca235c71eab8e49737c937a2c396"
1023            );
1024
1025            let buf = findnode_with_junk.as_slice();
1026            let decoded = Message::decode(buf).unwrap();
1027
1028            let expected_id = hex!(
1029                "ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31387574077f301b421bc84df7266c44e9e6d569fc56be00812904767bf5ccd1fc7f"
1030            );
1031            assert_matches!(decoded.msg, Message::FindNode(FindNode { id, expire: 1136239445 }) if id == expected_id);
1032        }
1033
1034        // checks that junk data at the end of the packet is discarded according to eip-8
1035        //
1036        // test vector from eip-8: https://eips.ethereum.org/EIPS/eip-8
1037        #[test]
1038        fn eip8_decode_neighbours() {
1039            let neighbours_with_junk = hex!(
1040                "c679fc8fe0b8b12f06577f2e802d34f6fa257e6137a995f6f4cbfc9ee50ed3710faf6e66f932c4c8d81d64343f429651328758b47d3dbc02c4042f0fff6946a50f4a49037a72bb550f3a7872363a83e1b9ee6469856c24eb4ef80b7535bcf99c0004f9015bf90150f84d846321163782115c82115db8403155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32f84984010203040101b840312c55512422cf9b8a4097e9a6ad79402e87a15ae909a4bfefa22398f03d20951933beea1e4dfa6f968212385e829f04c2d314fc2d4e255e0d3bc08792b069dbf8599020010db83c4d001500000000abcdef12820d05820d05b84038643200b172dcfef857492156971f0e6aa2c538d8b74010f8e140811d53b98c765dd2d96126051913f44582e8c199ad7c6d6819e9a56483f637feaac9448aacf8599020010db885a308d313198a2e037073488203e78203e8b8408dcab8618c3253b558d459da53bd8fa68935a719aff8b811197101a4b2b47dd2d47295286fc00cc081bb542d760717d1bdd6bec2c37cd72eca367d6dd3b9df738443b9a355010203b525a138aa34383fec3d2719a0"
1041            );
1042
1043            let buf = neighbours_with_junk.as_slice();
1044            let decoded = Message::decode(buf).unwrap();
1045
1046            let _ = NodeRecord {
1047                address: "99.33.22.55".parse().unwrap(),
1048                tcp_port: 4444,
1049                udp_port: 4445,
1050                id: hex!("3155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32").into(),
1051            }.length();
1052
1053            let expected_nodes: Vec<NodeRecord> = vec![
1054                NodeRecord {
1055                    address: "99.33.22.55".parse().unwrap(),
1056                    udp_port: 4444,
1057                    tcp_port: 4445,
1058                    id: hex!("3155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32").into(),
1059                },
1060                NodeRecord {
1061                    address: "1.2.3.4".parse().unwrap(),
1062                    udp_port: 1,
1063                    tcp_port: 1,
1064                    id: hex!("312c55512422cf9b8a4097e9a6ad79402e87a15ae909a4bfefa22398f03d20951933beea1e4dfa6f968212385e829f04c2d314fc2d4e255e0d3bc08792b069db").into(),
1065                },
1066                NodeRecord {
1067                    address: "2001:db8:3c4d:15::abcd:ef12".parse().unwrap(),
1068                    udp_port: 3333,
1069                    tcp_port: 3333,
1070                    id: hex!("38643200b172dcfef857492156971f0e6aa2c538d8b74010f8e140811d53b98c765dd2d96126051913f44582e8c199ad7c6d6819e9a56483f637feaac9448aac").into(),
1071                },
1072                NodeRecord {
1073                    address: "2001:db8:85a3:8d3:1319:8a2e:370:7348".parse().unwrap(),
1074                    udp_port: 999,
1075                    tcp_port: 1000,
1076                    id: hex!("8dcab8618c3253b558d459da53bd8fa68935a719aff8b811197101a4b2b47dd2d47295286fc00cc081bb542d760717d1bdd6bec2c37cd72eca367d6dd3b9df73").into(),
1077                },
1078            ];
1079            assert_matches!(decoded.msg, Message::Neighbours(Neighbours { nodes, expire: 1136239445 }) if nodes == expected_nodes);
1080        }
1081    }
1082}