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        let started_len = b.len();
213        if started_len < payload_length {
214            return Err(alloy_rlp::Error::InputTooShort);
215        }
216
217        // Geth allows the ipaddr to be possibly empty:
218        // <https://github.com/ethereum/go-ethereum/blob/380688c636a654becc8f114438c2a5d93d2db032/p2p/discover/v4_udp.go#L206-L209>
219        // <https://github.com/ethereum/go-ethereum/blob/380688c636a654becc8f114438c2a5d93d2db032/p2p/enode/node.go#L189-L189>
220        //
221        // Therefore, if we see an empty list instead of a properly formed `IpAddr`, we will
222        // instead use `IpV4Addr::UNSPECIFIED`
223        let address =
224            if *b.first().ok_or(alloy_rlp::Error::InputTooShort)? == alloy_rlp::EMPTY_STRING_CODE {
225                let addr = IpAddr::V4(Ipv4Addr::UNSPECIFIED);
226                b.advance(1);
227                addr
228            } else {
229                alloy_rlp::Decodable::decode(b)?
230            };
231
232        let this = NodeEndpoint {
233            address,
234            udp_port: alloy_rlp::Decodable::decode(b)?,
235            tcp_port: alloy_rlp::Decodable::decode(b)?,
236        };
237        let consumed = started_len - b.len();
238        if consumed != payload_length {
239            return Err(alloy_rlp::Error::ListLengthMismatch {
240                expected: payload_length,
241                got: consumed,
242            });
243        }
244        Ok(Self(this))
245    }
246}
247
248/// Represents the `from`, `to` fields in the packets
249#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, RlpEncodable, RlpDecodable)]
250pub struct NodeEndpoint {
251    /// The IP address of the network endpoint. It can be either IPv4 or IPv6.
252    pub address: IpAddr,
253    /// The UDP port used for communication in the discovery protocol.
254    pub udp_port: u16,
255    /// The TCP port used for communication in the `RLPx` protocol.
256    pub tcp_port: u16,
257}
258
259impl From<NodeRecord> for NodeEndpoint {
260    fn from(NodeRecord { address, tcp_port, udp_port, .. }: NodeRecord) -> Self {
261        Self { address, tcp_port, udp_port }
262    }
263}
264
265impl NodeEndpoint {
266    /// Creates a new [`NodeEndpoint`] from a given UDP address and TCP port.
267    pub const fn from_udp_address(udp_address: &std::net::SocketAddr, tcp_port: u16) -> Self {
268        Self { address: udp_address.ip(), udp_port: udp_address.port(), tcp_port }
269    }
270}
271
272/// A [FindNode packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#findnode-packet-0x03).
273#[derive(Clone, Copy, Debug, Eq, PartialEq, RlpEncodable)]
274pub struct FindNode {
275    /// The target node's ID, a 64-byte secp256k1 public key.
276    pub id: PeerId,
277    /// The expiration timestamp of the packet, an absolute UNIX time stamp.
278    pub expire: u64,
279}
280
281impl Decodable for FindNode {
282    // NOTE(onbjerg): Manual implementation to satisfy EIP-8.
283    //
284    // See https://eips.ethereum.org/EIPS/eip-8
285    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
286        let b = &mut &**buf;
287        let rlp_head = Header::decode(b)?;
288        if !rlp_head.list {
289            return Err(RlpError::UnexpectedString)
290        }
291        let started_len = b.len();
292
293        let this = Self { id: Decodable::decode(b)?, expire: Decodable::decode(b)? };
294
295        // NOTE(onbjerg): Because of EIP-8, we only check that we did not consume *more* than the
296        // payload length, i.e. it is ok if payload length is greater than what we consumed, as we
297        // just discard the remaining list items
298        let consumed = started_len - b.len();
299        if consumed > rlp_head.payload_length {
300            return Err(RlpError::ListLengthMismatch {
301                expected: rlp_head.payload_length,
302                got: consumed,
303            })
304        }
305
306        let rem = rlp_head.payload_length - consumed;
307        b.advance(rem);
308        *buf = *b;
309
310        Ok(this)
311    }
312}
313
314/// A [Neighbours packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#neighbors-packet-0x04).
315#[derive(Clone, Debug, Eq, PartialEq, RlpEncodable)]
316pub struct Neighbours {
317    /// The list of nodes containing IP, UDP port, TCP port, and node ID.
318    pub nodes: Vec<NodeRecord>,
319    /// The expiration timestamp of the packet, an absolute UNIX time stamp.
320    pub expire: u64,
321}
322
323impl Decodable for Neighbours {
324    // NOTE(onbjerg): Manual implementation to satisfy EIP-8.
325    //
326    // See https://eips.ethereum.org/EIPS/eip-8
327    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
328        let b = &mut &**buf;
329        let rlp_head = Header::decode(b)?;
330        if !rlp_head.list {
331            return Err(RlpError::UnexpectedString)
332        }
333        let started_len = b.len();
334
335        let this = Self { nodes: Decodable::decode(b)?, expire: Decodable::decode(b)? };
336
337        // NOTE(onbjerg): Because of EIP-8, we only check that we did not consume *more* than the
338        // payload length, i.e. it is ok if payload length is greater than what we consumed, as we
339        // just discard the remaining list items
340        let consumed = started_len - b.len();
341        if consumed > rlp_head.payload_length {
342            return Err(RlpError::ListLengthMismatch {
343                expected: rlp_head.payload_length,
344                got: consumed,
345            })
346        }
347
348        let rem = rlp_head.payload_length - consumed;
349        b.advance(rem);
350        *buf = *b;
351
352        Ok(this)
353    }
354}
355
356/// A [ENRRequest packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#enrrequest-packet-0x05).
357///
358/// This packet is used to request the current version of a node's Ethereum Node Record (ENR).
359#[derive(Clone, Copy, Debug, Eq, PartialEq, RlpEncodable)]
360pub struct EnrRequest {
361    /// The expiration timestamp for the request. No reply should be sent if it refers to a time in
362    /// the past.
363    pub expire: u64,
364}
365
366impl Decodable for EnrRequest {
367    // NOTE(onbjerg): Manual implementation to satisfy EIP-8.
368    //
369    // See https://eips.ethereum.org/EIPS/eip-8
370    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
371        let b = &mut &**buf;
372        let rlp_head = Header::decode(b)?;
373        if !rlp_head.list {
374            return Err(RlpError::UnexpectedString)
375        }
376        let started_len = b.len();
377
378        let this = Self { expire: Decodable::decode(b)? };
379
380        // NOTE(onbjerg): Because of EIP-8, we only check that we did not consume *more* than the
381        // payload length, i.e. it is ok if payload length is greater than what we consumed, as we
382        // just discard the remaining list items
383        let consumed = started_len - b.len();
384        if consumed > rlp_head.payload_length {
385            return Err(RlpError::ListLengthMismatch {
386                expected: rlp_head.payload_length,
387                got: consumed,
388            })
389        }
390
391        let rem = rlp_head.payload_length - consumed;
392        b.advance(rem);
393        *buf = *b;
394
395        Ok(this)
396    }
397}
398
399/// A [ENRResponse packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#enrresponse-packet-0x06).
400///
401/// This packet is used to respond to an `ENRRequest` packet and includes the requested ENR along
402/// with the hash of the original request.
403#[derive(Clone, Debug, Eq, PartialEq, RlpEncodable, RlpDecodable)]
404pub struct EnrResponse {
405    /// The hash of the `ENRRequest` packet being replied to.
406    pub request_hash: B256,
407    /// The ENR (Ethereum Node Record) for the responding node.
408    pub enr: Enr<SecretKey>,
409}
410
411// === impl EnrResponse ===
412
413impl EnrResponse {
414    /// Returns the [`ForkId`] if set
415    ///
416    /// See also <https://github.com/ethereum/go-ethereum/blob/9244d5cd61f3ea5a7645fdf2a1a96d53421e412f/eth/protocols/eth/discovery.go#L36>
417    pub fn eth_fork_id(&self) -> Option<ForkId> {
418        let mut maybe_fork_id = self.enr.get_raw_rlp(b"eth")?;
419        EnrForkIdEntry::decode(&mut maybe_fork_id).ok().map(Into::into)
420    }
421}
422
423/// Represents a Ping packet.
424///
425/// A [Ping packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#ping-packet-0x01).
426#[derive(Debug, Clone, Eq, PartialEq)]
427pub struct Ping {
428    /// The sender's endpoint.
429    pub from: NodeEndpoint,
430    /// The recipient's endpoint.
431    pub to: NodeEndpoint,
432    /// The expiration timestamp.
433    pub expire: u64,
434    /// Optional `enr_seq` for <https://eips.ethereum.org/EIPS/eip-868>
435    pub enr_sq: Option<u64>,
436}
437
438impl Encodable for Ping {
439    fn encode(&self, out: &mut dyn BufMut) {
440        #[derive(RlpEncodable)]
441        struct V4PingMessage<'a> {
442            version: u32,
443            from: &'a NodeEndpoint,
444            to: &'a NodeEndpoint,
445            expire: u64,
446        }
447
448        #[derive(RlpEncodable)]
449        struct V4PingMessageEIP868<'a> {
450            version: u32,
451            from: &'a NodeEndpoint,
452            to: &'a NodeEndpoint,
453            expire: u64,
454            enr_seq: u64,
455        }
456        if let Some(enr_seq) = self.enr_sq {
457            V4PingMessageEIP868 {
458                version: 4, // version 4
459                from: &self.from,
460                to: &self.to,
461                expire: self.expire,
462                enr_seq,
463            }
464            .encode(out);
465        } else {
466            V4PingMessage {
467                version: 4, // version 4
468                from: &self.from,
469                to: &self.to,
470                expire: self.expire,
471            }
472            .encode(out);
473        }
474    }
475}
476
477impl Decodable for Ping {
478    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
479        let b = &mut &**buf;
480        let rlp_head = Header::decode(b)?;
481        if !rlp_head.list {
482            return Err(RlpError::UnexpectedString)
483        }
484        let started_len = b.len();
485
486        // > Implementations should ignore any mismatches in version:
487        // <https://github.com/ethereum/devp2p/blob/master/discv4.md#ping-packet-0x01>
488        let _version = u32::decode(b)?;
489
490        // see `Decodable` implementation in `PingNodeEndpoint` for why this is needed
491        let from = PingNodeEndpoint::decode(b)?.0;
492
493        let mut this =
494            Self { from, to: Decodable::decode(b)?, expire: Decodable::decode(b)?, enr_sq: None };
495
496        // only decode the ENR sequence if there's more data in the datagram to decode else skip
497        if b.has_remaining() {
498            this.enr_sq = Some(Decodable::decode(b)?);
499        }
500
501        let consumed = started_len - b.len();
502        if consumed > rlp_head.payload_length {
503            return Err(RlpError::ListLengthMismatch {
504                expected: rlp_head.payload_length,
505                got: consumed,
506            })
507        }
508        let rem = rlp_head.payload_length - consumed;
509        b.advance(rem);
510        *buf = *b;
511        Ok(this)
512    }
513}
514
515/// Represents a Pong packet.
516///
517/// A [Pong packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#pong-packet-0x02).
518#[derive(Clone, Debug, Eq, PartialEq)]
519pub struct Pong {
520    /// The recipient's endpoint.
521    pub to: NodeEndpoint,
522    /// The hash of the corresponding ping packet.
523    pub echo: B256,
524    /// The expiration timestamp.
525    pub expire: u64,
526    /// Optional `enr_seq` for <https://eips.ethereum.org/EIPS/eip-868>
527    pub enr_sq: Option<u64>,
528}
529
530impl Encodable for Pong {
531    fn encode(&self, out: &mut dyn BufMut) {
532        #[derive(RlpEncodable)]
533        struct PongMessageEIP868<'a> {
534            to: &'a NodeEndpoint,
535            echo: &'a B256,
536            expire: u64,
537            enr_seq: u64,
538        }
539
540        #[derive(RlpEncodable)]
541        struct PongMessage<'a> {
542            to: &'a NodeEndpoint,
543            echo: &'a B256,
544            expire: u64,
545        }
546
547        if let Some(enr_seq) = self.enr_sq {
548            PongMessageEIP868 { to: &self.to, echo: &self.echo, expire: self.expire, enr_seq }
549                .encode(out);
550        } else {
551            PongMessage { to: &self.to, echo: &self.echo, expire: self.expire }.encode(out);
552        }
553    }
554}
555
556impl Decodable for Pong {
557    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
558        let b = &mut &**buf;
559        let rlp_head = Header::decode(b)?;
560        if !rlp_head.list {
561            return Err(RlpError::UnexpectedString)
562        }
563        let started_len = b.len();
564        let mut this = Self {
565            to: Decodable::decode(b)?,
566            echo: Decodable::decode(b)?,
567            expire: Decodable::decode(b)?,
568            enr_sq: None,
569        };
570
571        // only decode the ENR sequence if there's more data in the datagram to decode else skip
572        if b.has_remaining() {
573            this.enr_sq = Some(Decodable::decode(b)?);
574        }
575
576        let consumed = started_len - b.len();
577        if consumed > rlp_head.payload_length {
578            return Err(RlpError::ListLengthMismatch {
579                expected: rlp_head.payload_length,
580                got: consumed,
581            })
582        }
583        let rem = rlp_head.payload_length - consumed;
584        b.advance(rem);
585        *buf = *b;
586
587        Ok(this)
588    }
589}
590
591#[cfg(test)]
592mod tests {
593    use super::*;
594    use crate::{
595        test_utils::{rng_endpoint, rng_ipv4_record, rng_ipv6_record, rng_message},
596        DEFAULT_DISCOVERY_PORT, SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS,
597    };
598    use alloy_primitives::hex;
599    use assert_matches::assert_matches;
600    use enr::EnrPublicKey;
601    use rand_08::{thread_rng as rng, Rng, RngCore};
602    use reth_ethereum_forks::ForkHash;
603
604    #[test]
605    fn test_endpoint_ipv_v4() {
606        let mut rng = rng();
607        for _ in 0..100 {
608            let mut ip = [0u8; 4];
609            rng.fill_bytes(&mut ip);
610            let msg = NodeEndpoint {
611                address: IpAddr::V4(ip.into()),
612                tcp_port: rng.r#gen(),
613                udp_port: rng.r#gen(),
614            };
615
616            let decoded = NodeEndpoint::decode(&mut alloy_rlp::encode(msg).as_slice()).unwrap();
617            assert_eq!(msg, decoded);
618        }
619    }
620
621    #[test]
622    fn test_endpoint_ipv_64() {
623        let mut rng = rng();
624        for _ in 0..100 {
625            let mut ip = [0u8; 16];
626            rng.fill_bytes(&mut ip);
627            let msg = NodeEndpoint {
628                address: IpAddr::V6(ip.into()),
629                tcp_port: rng.r#gen(),
630                udp_port: rng.r#gen(),
631            };
632
633            let decoded = NodeEndpoint::decode(&mut alloy_rlp::encode(msg).as_slice()).unwrap();
634            assert_eq!(msg, decoded);
635        }
636    }
637
638    #[test]
639    fn test_ping_message() {
640        let mut rng = rng();
641        for _ in 0..100 {
642            let mut ip = [0u8; 16];
643            rng.fill_bytes(&mut ip);
644            let msg = Ping {
645                from: rng_endpoint(&mut rng),
646                to: rng_endpoint(&mut rng),
647                expire: 0,
648                enr_sq: None,
649            };
650
651            let decoded = Ping::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
652            assert_eq!(msg, decoded);
653        }
654    }
655
656    #[test]
657    fn test_ping_message_with_enr() {
658        let mut rng = rng();
659        for _ in 0..100 {
660            let mut ip = [0u8; 16];
661            rng.fill_bytes(&mut ip);
662            let msg = Ping {
663                from: rng_endpoint(&mut rng),
664                to: rng_endpoint(&mut rng),
665                expire: 0,
666                enr_sq: Some(rng.r#gen()),
667            };
668
669            let decoded = Ping::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
670            assert_eq!(msg, decoded);
671        }
672    }
673
674    #[test]
675    fn test_pong_message() {
676        let mut rng = rng();
677        for _ in 0..100 {
678            let mut ip = [0u8; 16];
679            rng.fill_bytes(&mut ip);
680            let msg = Pong {
681                to: rng_endpoint(&mut rng),
682                echo: B256::random(),
683                expire: rng.r#gen(),
684                enr_sq: None,
685            };
686
687            let decoded = Pong::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
688            assert_eq!(msg, decoded);
689        }
690    }
691
692    #[test]
693    fn test_pong_message_with_enr() {
694        let mut rng = rng();
695        for _ in 0..100 {
696            let mut ip = [0u8; 16];
697            rng.fill_bytes(&mut ip);
698            let msg = Pong {
699                to: rng_endpoint(&mut rng),
700                echo: B256::random(),
701                expire: rng.r#gen(),
702                enr_sq: Some(rng.r#gen()),
703            };
704
705            let decoded = Pong::decode(&mut alloy_rlp::encode(&msg).as_slice()).unwrap();
706            assert_eq!(msg, decoded);
707        }
708    }
709
710    #[test]
711    fn test_hash_mismatch() {
712        let mut rng = rng();
713        let msg = rng_message(&mut rng);
714        let (secret_key, _) = SECP256K1.generate_keypair(&mut rng);
715        let (buf, _) = msg.encode(&secret_key);
716
717        let mut buf_vec = buf.to_vec();
718        buf_vec.push(0);
719        match Message::decode(buf_vec.as_slice()).unwrap_err() {
720            DecodePacketError::HashMismatch => {}
721            err => {
722                unreachable!("unexpected err {}", err)
723            }
724        }
725    }
726
727    #[test]
728    fn neighbours_max_ipv4() {
729        let mut rng = rng();
730        let msg = Message::Neighbours(Neighbours {
731            nodes: std::iter::repeat_with(|| rng_ipv4_record(&mut rng)).take(16).collect(),
732            expire: rng.r#gen(),
733        });
734        let (secret_key, _) = SECP256K1.generate_keypair(&mut rng);
735
736        let (encoded, _) = msg.encode(&secret_key);
737        // Assert that 16 nodes never fit into one packet
738        assert!(encoded.len() > MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
739    }
740
741    #[test]
742    fn neighbours_max_nodes() {
743        let mut rng = rng();
744        for _ in 0..1000 {
745            let msg = Message::Neighbours(Neighbours {
746                nodes: std::iter::repeat_with(|| rng_ipv6_record(&mut rng))
747                    .take(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS)
748                    .collect(),
749                expire: rng.r#gen(),
750            });
751            let (secret_key, _) = SECP256K1.generate_keypair(&mut rng);
752
753            let (encoded, _) = msg.encode(&secret_key);
754            assert!(encoded.len() <= MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
755
756            let mut neighbours = Neighbours {
757                nodes: std::iter::repeat_with(|| rng_ipv6_record(&mut rng))
758                    .take(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS - 1)
759                    .collect(),
760                expire: rng.r#gen(),
761            };
762            neighbours.nodes.push(rng_ipv4_record(&mut rng));
763            let msg = Message::Neighbours(neighbours);
764            let (encoded, _) = msg.encode(&secret_key);
765            assert!(encoded.len() <= MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
766        }
767    }
768
769    #[test]
770    fn test_encode_decode_message() {
771        let mut rng = rng();
772        for _ in 0..100 {
773            let msg = rng_message(&mut rng);
774            let (secret_key, pk) = SECP256K1.generate_keypair(&mut rng);
775            let sender_id = pk2id(&pk);
776
777            let (buf, _) = msg.encode(&secret_key);
778
779            let packet = Message::decode(buf.as_ref()).unwrap();
780
781            assert_eq!(msg, packet.msg);
782            assert_eq!(sender_id, packet.node_id);
783        }
784    }
785
786    #[test]
787    fn decode_pong_packet() {
788        let packet = "2ad84c37327a06c2522cf7bc039621da89f68907441b755935bb308dc4cd17d6fe550e90329ad6a516ca7db18e08900067928a0dfa3b5c75d55a42c984497373698d98616662c048983ea85895ea2da765eabeb15525478384e106337bfd8ed50002f3c9843ed8cae682fd1c80a008ad4dead0922211df47593e7d837b2b23d13954285871ca23250ea594993ded84635690e5829670";
789        let data = hex::decode(packet).unwrap();
790        Message::decode(&data).unwrap();
791    }
792    #[test]
793    fn decode_ping_packet() {
794        let packet = "05ae5bf922cf2a93f97632a4ab0943dc252a0dab0c42d86dd62e5d91e1a0966e9b628fbf4763fdfbb928540460b797e6be2e7058a82f6083f6d2e7391bb021741459976d4152aa16bbee0c3609dcfac6668db1ef78b7ee9f8b4ced10dd5ae2900101df04cb8403d12d4f82765f82765fc9843ed8cae6828aa6808463569916829670";
795        let data = hex::decode(packet).unwrap();
796        Message::decode(&data).unwrap();
797    }
798
799    #[test]
800    fn encode_decode_enr_msg() {
801        use alloy_rlp::Decodable;
802        use enr::secp256k1::SecretKey;
803        use std::net::Ipv4Addr;
804
805        let mut rng = rand_08::rngs::OsRng;
806        let key = SecretKey::new(&mut rng);
807        let ip = Ipv4Addr::new(127, 0, 0, 1);
808        let tcp = 3000;
809
810        let fork_id: ForkId = ForkId { hash: ForkHash([220, 233, 108, 45]), next: 0u64 };
811
812        let enr = {
813            let mut builder = Enr::builder();
814            builder.ip(ip.into());
815            builder.tcp4(tcp);
816            let mut buf = Vec::new();
817            let forkentry = EnrForkIdEntry { fork_id };
818            forkentry.encode(&mut buf);
819            builder.add_value_rlp("eth", buf.into());
820            builder.build(&key).unwrap()
821        };
822
823        let enr_response = EnrResponse { request_hash: B256::random(), enr };
824
825        let mut buf = Vec::new();
826        enr_response.encode(&mut buf);
827
828        let decoded = EnrResponse::decode(&mut &buf[..]).unwrap();
829
830        let fork_id_decoded = decoded.eth_fork_id().unwrap();
831        assert_eq!(fork_id, fork_id_decoded);
832    }
833
834    // test vector from the enr library rlp encoding tests
835    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#L1019>
836
837    #[test]
838    fn encode_known_rlp_enr() {
839        use alloy_rlp::Decodable;
840        use enr::{secp256k1::SecretKey, EnrPublicKey};
841        use std::net::Ipv4Addr;
842
843        let valid_record = hex!(
844            "f884b8407098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c01826964827634826970847f00000189736563703235366b31a103ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31388375647082765f"
845        );
846        let signature = hex!(
847            "7098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c"
848        );
849        let expected_pubkey =
850            hex!("03ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd3138");
851
852        let enr = Enr::<SecretKey>::decode(&mut &valid_record[..]).unwrap();
853        let pubkey = enr.public_key().encode();
854
855        assert_eq!(enr.ip4(), Some(Ipv4Addr::new(127, 0, 0, 1)));
856        assert_eq!(enr.id(), Some(String::from("v4")));
857        assert_eq!(enr.udp4(), Some(DEFAULT_DISCOVERY_PORT));
858        assert_eq!(enr.tcp4(), None);
859        assert_eq!(enr.signature(), &signature[..]);
860        assert_eq!(pubkey.to_vec(), expected_pubkey);
861        assert!(enr.verify());
862
863        assert_eq!(&alloy_rlp::encode(&enr)[..], &valid_record[..]);
864
865        // ensure the length is equal
866        assert_eq!(enr.length(), valid_record.len());
867    }
868
869    // test vector from the enr library rlp encoding tests
870    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#L1019>
871    #[test]
872    fn decode_enr_rlp() {
873        use enr::secp256k1::SecretKey;
874        use std::net::Ipv4Addr;
875
876        let valid_record = hex!(
877            "f884b8407098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c01826964827634826970847f00000189736563703235366b31a103ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31388375647082765f"
878        );
879        let signature = hex!(
880            "7098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c"
881        );
882        let expected_pubkey =
883            hex!("03ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd3138");
884
885        let mut valid_record_buf = valid_record.as_slice();
886        let enr = Enr::<SecretKey>::decode(&mut valid_record_buf).unwrap();
887        let pubkey = enr.public_key().encode();
888
889        // Byte array must be consumed after enr has finished decoding
890        assert!(valid_record_buf.is_empty());
891
892        assert_eq!(enr.ip4(), Some(Ipv4Addr::new(127, 0, 0, 1)));
893        assert_eq!(enr.id(), Some(String::from("v4")));
894        assert_eq!(enr.udp4(), Some(DEFAULT_DISCOVERY_PORT));
895        assert_eq!(enr.tcp4(), None);
896        assert_eq!(enr.signature(), &signature[..]);
897        assert_eq!(pubkey.to_vec(), expected_pubkey);
898        assert!(enr.verify());
899    }
900
901    // test for failing message decode
902    #[test]
903    fn decode_failing_packet() {
904        let packet = hex!(
905            "2467ab56952aedf4cfb8bb7830ddc8922d0f992185229919dad9de3841fe95d9b3a7b52459398235f6d3805644666d908b45edb3670414ed97f357afba51f71f7d35c1f45878ba732c3868b04ca42ff0ed347c99efcf3a5768afed68eb21ef960001db04c3808080c9840a480e8f82765f808466a9a06386019106833efe"
906        );
907
908        let _message = Message::decode(&packet[..]).unwrap();
909    }
910
911    // test for failing message decode
912    #[test]
913    fn decode_node() {
914        let packet = hex!("cb840000000082115c82115d");
915        let _message = NodeEndpoint::decode(&mut &packet[..]).unwrap();
916    }
917
918    // test vector from the enr library rlp encoding tests
919    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#LL1206C35-L1206C35>
920    #[test]
921    fn encode_decode_enr_rlp() {
922        use enr::{secp256k1::SecretKey, EnrPublicKey};
923        use std::net::Ipv4Addr;
924
925        let key = SecretKey::new(&mut rand_08::rngs::OsRng);
926        let ip = Ipv4Addr::new(127, 0, 0, 1);
927        let tcp = 3000;
928
929        let enr = {
930            let mut builder = Enr::builder();
931            builder.ip(ip.into());
932            builder.tcp4(tcp);
933            builder.build(&key).unwrap()
934        };
935
936        let mut encoded_bytes = &alloy_rlp::encode(&enr)[..];
937        let decoded_enr = Enr::<SecretKey>::decode(&mut encoded_bytes).unwrap();
938
939        // Byte array must be consumed after enr has finished decoding
940        assert!(encoded_bytes.is_empty());
941
942        assert_eq!(decoded_enr, enr);
943        assert_eq!(decoded_enr.id(), Some("v4".into()));
944        assert_eq!(decoded_enr.ip4(), Some(ip));
945        assert_eq!(decoded_enr.tcp4(), Some(tcp));
946        assert_eq!(
947            decoded_enr.public_key().encode(),
948            key.public_key(secp256k1::SECP256K1).encode()
949        );
950        assert!(decoded_enr.verify());
951    }
952
953    mod eip8 {
954        use super::*;
955
956        fn junk_enr_request() -> Vec<u8> {
957            let mut buf = Vec::new();
958            // enr request is just an expiration
959            let expire: u64 = 123456;
960
961            // add some junk
962            let junk: u64 = 112233;
963
964            // rlp header encoding
965            let payload_length = expire.length() + junk.length();
966            alloy_rlp::Header { list: true, payload_length }.encode(&mut buf);
967
968            // fields
969            expire.encode(&mut buf);
970            junk.encode(&mut buf);
971
972            buf
973        }
974
975        // checks that junk data at the end of the packet is discarded according to eip-8
976        #[test]
977        fn eip8_decode_enr_request() {
978            let enr_request_with_junk = junk_enr_request();
979
980            let mut buf = enr_request_with_junk.as_slice();
981            let decoded = EnrRequest::decode(&mut buf).unwrap();
982            assert_eq!(decoded.expire, 123456);
983        }
984
985        // checks that junk data at the end of the packet is discarded according to eip-8
986        //
987        // test vector from eip-8: https://eips.ethereum.org/EIPS/eip-8
988        #[test]
989        fn eip8_decode_findnode() {
990            let findnode_with_junk = hex!(
991                "c7c44041b9f7c7e41934417ebac9a8e1a4c6298f74553f2fcfdcae6ed6fe53163eb3d2b52e39fe91831b8a927bf4fc222c3902202027e5e9eb812195f95d20061ef5cd31d502e47ecb61183f74a504fe04c51e73df81f25c4d506b26db4517490103f84eb840ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31387574077f301b421bc84df7266c44e9e6d569fc56be00812904767bf5ccd1fc7f8443b9a35582999983999999280dc62cc8255c73471e0a61da0c89acdc0e035e260add7fc0c04ad9ebf3919644c91cb247affc82b69bd2ca235c71eab8e49737c937a2c396"
992            );
993
994            let buf = findnode_with_junk.as_slice();
995            let decoded = Message::decode(buf).unwrap();
996
997            let expected_id = hex!(
998                "ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31387574077f301b421bc84df7266c44e9e6d569fc56be00812904767bf5ccd1fc7f"
999            );
1000            assert_matches!(decoded.msg, Message::FindNode(FindNode { id, expire: 1136239445 }) if id == expected_id);
1001        }
1002
1003        // checks that junk data at the end of the packet is discarded according to eip-8
1004        //
1005        // test vector from eip-8: https://eips.ethereum.org/EIPS/eip-8
1006        #[test]
1007        fn eip8_decode_neighbours() {
1008            let neighbours_with_junk = hex!(
1009                "c679fc8fe0b8b12f06577f2e802d34f6fa257e6137a995f6f4cbfc9ee50ed3710faf6e66f932c4c8d81d64343f429651328758b47d3dbc02c4042f0fff6946a50f4a49037a72bb550f3a7872363a83e1b9ee6469856c24eb4ef80b7535bcf99c0004f9015bf90150f84d846321163782115c82115db8403155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32f84984010203040101b840312c55512422cf9b8a4097e9a6ad79402e87a15ae909a4bfefa22398f03d20951933beea1e4dfa6f968212385e829f04c2d314fc2d4e255e0d3bc08792b069dbf8599020010db83c4d001500000000abcdef12820d05820d05b84038643200b172dcfef857492156971f0e6aa2c538d8b74010f8e140811d53b98c765dd2d96126051913f44582e8c199ad7c6d6819e9a56483f637feaac9448aacf8599020010db885a308d313198a2e037073488203e78203e8b8408dcab8618c3253b558d459da53bd8fa68935a719aff8b811197101a4b2b47dd2d47295286fc00cc081bb542d760717d1bdd6bec2c37cd72eca367d6dd3b9df738443b9a355010203b525a138aa34383fec3d2719a0"
1010            );
1011
1012            let buf = neighbours_with_junk.as_slice();
1013            let decoded = Message::decode(buf).unwrap();
1014
1015            let _ = NodeRecord {
1016                address: "99.33.22.55".parse().unwrap(),
1017                tcp_port: 4444,
1018                udp_port: 4445,
1019                id: hex!("3155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32").into(),
1020            }.length();
1021
1022            let expected_nodes: Vec<NodeRecord> = vec![
1023                NodeRecord {
1024                    address: "99.33.22.55".parse().unwrap(),
1025                    udp_port: 4444,
1026                    tcp_port: 4445,
1027                    id: hex!("3155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32").into(),
1028                },
1029                NodeRecord {
1030                    address: "1.2.3.4".parse().unwrap(),
1031                    udp_port: 1,
1032                    tcp_port: 1,
1033                    id: hex!("312c55512422cf9b8a4097e9a6ad79402e87a15ae909a4bfefa22398f03d20951933beea1e4dfa6f968212385e829f04c2d314fc2d4e255e0d3bc08792b069db").into(),
1034                },
1035                NodeRecord {
1036                    address: "2001:db8:3c4d:15::abcd:ef12".parse().unwrap(),
1037                    udp_port: 3333,
1038                    tcp_port: 3333,
1039                    id: hex!("38643200b172dcfef857492156971f0e6aa2c538d8b74010f8e140811d53b98c765dd2d96126051913f44582e8c199ad7c6d6819e9a56483f637feaac9448aac").into(),
1040                },
1041                NodeRecord {
1042                    address: "2001:db8:85a3:8d3:1319:8a2e:370:7348".parse().unwrap(),
1043                    udp_port: 999,
1044                    tcp_port: 1000,
1045                    id: hex!("8dcab8618c3253b558d459da53bd8fa68935a719aff8b811197101a4b2b47dd2d47295286fc00cc081bb542d760717d1bdd6bec2c37cd72eca367d6dd3b9df73").into(),
1046                },
1047            ];
1048            assert_matches!(decoded.msg, Message::Neighbours(Neighbours { nodes, expire: 1136239445 }) if nodes == expected_nodes);
1049        }
1050    }
1051}