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,
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::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(rec.to_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 = recoverable_sig.recover_ecdsa(msg)?;
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)]
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
421impl Decodable for EnrResponse {
422    // Manual implementation to satisfy EIP-8: additional list elements are ignored.
423    //
424    // See https://eips.ethereum.org/EIPS/eip-8
425    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
426        let mut payload = Header::decode_bytes(buf, true)?;
427        let request_hash = Decodable::decode(&mut payload)?;
428
429        // `Enr::decode` applies the ENR size limit to its whole input, so it only gets the
430        // record's own bytes and not the additional list elements after it.
431        let mut rest = payload;
432        Header::decode_bytes(&mut rest, true)?;
433        let mut enr = &payload[..payload.len() - rest.len()];
434
435        Ok(Self { request_hash, enr: Decodable::decode(&mut enr)? })
436    }
437}
438
439/// Represents a Ping packet.
440///
441/// A [Ping packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#ping-packet-0x01).
442#[derive(Debug, Clone, Eq, PartialEq)]
443pub struct Ping {
444    /// The sender's endpoint.
445    pub from: NodeEndpoint,
446    /// The recipient's endpoint.
447    pub to: NodeEndpoint,
448    /// The expiration timestamp.
449    pub expire: u64,
450    /// Optional `enr_seq` for <https://eips.ethereum.org/EIPS/eip-868>
451    pub enr_sq: Option<u64>,
452}
453
454impl Encodable for Ping {
455    fn length(&self) -> usize {
456        let mut payload_length =
457            4u32.length() + self.from.length() + self.to.length() + self.expire.length();
458        if let Some(enr_seq) = self.enr_sq {
459            payload_length += enr_seq.length();
460        }
461        payload_length + alloy_rlp::length_of_length(payload_length)
462    }
463
464    fn encode(&self, out: &mut dyn BufMut) {
465        #[derive(RlpEncodable)]
466        struct V4PingMessage<'a> {
467            version: u32,
468            from: &'a NodeEndpoint,
469            to: &'a NodeEndpoint,
470            expire: u64,
471        }
472
473        #[derive(RlpEncodable)]
474        struct V4PingMessageEIP868<'a> {
475            version: u32,
476            from: &'a NodeEndpoint,
477            to: &'a NodeEndpoint,
478            expire: u64,
479            enr_seq: u64,
480        }
481        if let Some(enr_seq) = self.enr_sq {
482            V4PingMessageEIP868 {
483                version: 4, // version 4
484                from: &self.from,
485                to: &self.to,
486                expire: self.expire,
487                enr_seq,
488            }
489            .encode(out);
490        } else {
491            V4PingMessage {
492                version: 4, // version 4
493                from: &self.from,
494                to: &self.to,
495                expire: self.expire,
496            }
497            .encode(out);
498        }
499    }
500}
501
502impl Decodable for Ping {
503    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
504        let b = &mut &**buf;
505        let rlp_head = Header::decode(b)?;
506        if !rlp_head.list {
507            return Err(RlpError::UnexpectedString)
508        }
509        let started_len = b.len();
510
511        // > Implementations should ignore any mismatches in version:
512        // <https://github.com/ethereum/devp2p/blob/master/discv4.md#ping-packet-0x01>
513        let _version = u32::decode(b)?;
514
515        // see `Decodable` implementation in `PingNodeEndpoint` for why this is needed
516        let from = PingNodeEndpoint::decode(b)?.0;
517
518        let mut this =
519            Self { from, to: Decodable::decode(b)?, expire: Decodable::decode(b)?, enr_sq: None };
520
521        // only decode the ENR sequence if there's more data in the datagram to decode else skip
522        if b.has_remaining() {
523            this.enr_sq = Some(Decodable::decode(b)?);
524        }
525
526        let consumed = started_len - b.len();
527        if consumed > rlp_head.payload_length {
528            return Err(RlpError::ListLengthMismatch {
529                expected: rlp_head.payload_length,
530                got: consumed,
531            })
532        }
533        let rem = rlp_head.payload_length - consumed;
534        b.advance(rem);
535        *buf = *b;
536        Ok(this)
537    }
538}
539
540/// Represents a Pong packet.
541///
542/// A [Pong packet](https://github.com/ethereum/devp2p/blob/master/discv4.md#pong-packet-0x02).
543#[derive(Clone, Debug, Eq, PartialEq)]
544pub struct Pong {
545    /// The recipient's endpoint.
546    pub to: NodeEndpoint,
547    /// The hash of the corresponding ping packet.
548    pub echo: B256,
549    /// The expiration timestamp.
550    pub expire: u64,
551    /// Optional `enr_seq` for <https://eips.ethereum.org/EIPS/eip-868>
552    pub enr_sq: Option<u64>,
553}
554
555impl Encodable for Pong {
556    fn length(&self) -> usize {
557        let mut payload_length = self.to.length() + self.echo.length() + self.expire.length();
558        if let Some(enr_seq) = self.enr_sq {
559            payload_length += enr_seq.length();
560        }
561        payload_length + alloy_rlp::length_of_length(payload_length)
562    }
563
564    fn encode(&self, out: &mut dyn BufMut) {
565        #[derive(RlpEncodable)]
566        struct PongMessageEIP868<'a> {
567            to: &'a NodeEndpoint,
568            echo: &'a B256,
569            expire: u64,
570            enr_seq: u64,
571        }
572
573        #[derive(RlpEncodable)]
574        struct PongMessage<'a> {
575            to: &'a NodeEndpoint,
576            echo: &'a B256,
577            expire: u64,
578        }
579
580        if let Some(enr_seq) = self.enr_sq {
581            PongMessageEIP868 { to: &self.to, echo: &self.echo, expire: self.expire, enr_seq }
582                .encode(out);
583        } else {
584            PongMessage { to: &self.to, echo: &self.echo, expire: self.expire }.encode(out);
585        }
586    }
587}
588
589impl Decodable for Pong {
590    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
591        let b = &mut &**buf;
592        let rlp_head = Header::decode(b)?;
593        if !rlp_head.list {
594            return Err(RlpError::UnexpectedString)
595        }
596        let started_len = b.len();
597        let mut this = Self {
598            to: Decodable::decode(b)?,
599            echo: Decodable::decode(b)?,
600            expire: Decodable::decode(b)?,
601            enr_sq: None,
602        };
603
604        // only decode the ENR sequence if there's more data in the datagram to decode else skip
605        if b.has_remaining() {
606            this.enr_sq = Some(Decodable::decode(b)?);
607        }
608
609        let consumed = started_len - b.len();
610        if consumed > rlp_head.payload_length {
611            return Err(RlpError::ListLengthMismatch {
612                expected: rlp_head.payload_length,
613                got: consumed,
614            })
615        }
616        let rem = rlp_head.payload_length - consumed;
617        b.advance(rem);
618        *buf = *b;
619
620        Ok(this)
621    }
622}
623
624#[cfg(test)]
625mod tests {
626    use super::*;
627    use crate::{
628        test_utils::{rng_endpoint, rng_ipv4_record, rng_ipv6_record, rng_message},
629        DEFAULT_DISCOVERY_PORT, SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS,
630    };
631    use alloy_primitives::hex;
632    use alloy_rlp::Decodable;
633    use assert_matches::assert_matches;
634    use enr::{secp256k1::SecretKey, EnrPublicKey};
635    use rand::{rng, Rng, RngCore};
636    use reth_ethereum_forks::ForkHash;
637    use std::net::Ipv4Addr;
638
639    #[test]
640    fn test_endpoint_ipv_v4() {
641        let mut rng = rng();
642        for _ in 0..100 {
643            let mut ip = [0u8; 4];
644            rng.fill_bytes(&mut ip);
645            let msg = NodeEndpoint {
646                address: IpAddr::V4(ip.into()),
647                tcp_port: rng.random(),
648                udp_port: rng.random(),
649            };
650
651            let decoded = NodeEndpoint::decode(&mut alloy_rlp::encode(msg).as_slice()).unwrap();
652            assert_eq!(msg, decoded);
653        }
654    }
655
656    #[test]
657    fn test_endpoint_ipv_64() {
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 = NodeEndpoint {
663                address: IpAddr::V6(ip.into()),
664                tcp_port: rng.random(),
665                udp_port: rng.random(),
666            };
667
668            let decoded = NodeEndpoint::decode(&mut alloy_rlp::encode(msg).as_slice()).unwrap();
669            assert_eq!(msg, decoded);
670        }
671    }
672
673    #[test]
674    fn test_ping_message() {
675        let mut rng = rng();
676        for _ in 0..100 {
677            let mut ip = [0u8; 16];
678            rng.fill_bytes(&mut ip);
679            let msg = Ping {
680                from: rng_endpoint(&mut rng),
681                to: rng_endpoint(&mut rng),
682                expire: 0,
683                enr_sq: None,
684            };
685
686            let encoded = alloy_rlp::encode(&msg);
687            assert_eq!(msg.length(), encoded.len());
688            let mut buf = encoded.as_slice();
689            let decoded = Ping::decode(&mut buf).unwrap();
690            assert_eq!(msg, decoded);
691            assert!(buf.is_empty());
692        }
693    }
694
695    #[test]
696    fn test_ping_message_with_enr() {
697        let mut rng = rng();
698        for _ in 0..100 {
699            let mut ip = [0u8; 16];
700            rng.fill_bytes(&mut ip);
701            let msg = Ping {
702                from: rng_endpoint(&mut rng),
703                to: rng_endpoint(&mut rng),
704                expire: 0,
705                enr_sq: Some(rng.random()),
706            };
707
708            let encoded = alloy_rlp::encode(&msg);
709            assert_eq!(msg.length(), encoded.len());
710            let mut buf = encoded.as_slice();
711            let decoded = Ping::decode(&mut buf).unwrap();
712            assert_eq!(msg, decoded);
713            assert!(buf.is_empty());
714        }
715    }
716
717    #[test]
718    fn test_pong_message() {
719        let mut rng = rng();
720        for _ in 0..100 {
721            let mut ip = [0u8; 16];
722            rng.fill_bytes(&mut ip);
723            let msg = Pong {
724                to: rng_endpoint(&mut rng),
725                echo: B256::random(),
726                expire: rng.random(),
727                enr_sq: None,
728            };
729
730            let encoded = alloy_rlp::encode(&msg);
731            assert_eq!(msg.length(), encoded.len());
732            let mut buf = encoded.as_slice();
733            let decoded = Pong::decode(&mut buf).unwrap();
734            assert_eq!(msg, decoded);
735            assert!(buf.is_empty());
736        }
737    }
738
739    #[test]
740    fn test_pong_message_with_enr() {
741        let mut rng = rng();
742        for _ in 0..100 {
743            let mut ip = [0u8; 16];
744            rng.fill_bytes(&mut ip);
745            let msg = Pong {
746                to: rng_endpoint(&mut rng),
747                echo: B256::random(),
748                expire: rng.random(),
749                enr_sq: Some(rng.random()),
750            };
751
752            let encoded = alloy_rlp::encode(&msg);
753            assert_eq!(msg.length(), encoded.len());
754            let mut buf = encoded.as_slice();
755            let decoded = Pong::decode(&mut buf).unwrap();
756            assert_eq!(msg, decoded);
757            assert!(buf.is_empty());
758        }
759    }
760
761    #[test]
762    fn test_hash_mismatch() {
763        let mut rng = rng();
764        let msg = rng_message(&mut rng);
765        let (secret_key, _) = secp256k1::generate_keypair(&mut rng);
766        let (buf, _) = msg.encode(&secret_key);
767
768        let mut buf_vec = buf.to_vec();
769        buf_vec.push(0);
770        match Message::decode(buf_vec.as_slice()).unwrap_err() {
771            DecodePacketError::HashMismatch => {}
772            err => {
773                unreachable!("unexpected err {}", err)
774            }
775        }
776    }
777
778    #[test]
779    fn neighbours_max_ipv4() {
780        let mut rng = rng();
781        let msg = Message::Neighbours(Neighbours {
782            nodes: std::iter::repeat_with(|| rng_ipv4_record(&mut rng)).take(16).collect(),
783            expire: rng.random(),
784        });
785        let (secret_key, _) = secp256k1::generate_keypair(&mut rng);
786
787        let (encoded, _) = msg.encode(&secret_key);
788        // Assert that 16 nodes never fit into one packet
789        assert!(encoded.len() > MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
790    }
791
792    #[test]
793    fn neighbours_max_nodes() {
794        let mut rng = rng();
795        for _ in 0..1000 {
796            let msg = Message::Neighbours(Neighbours {
797                nodes: std::iter::repeat_with(|| rng_ipv6_record(&mut rng))
798                    .take(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS)
799                    .collect(),
800                expire: rng.random(),
801            });
802            let (secret_key, _) = secp256k1::generate_keypair(&mut rng);
803
804            let (encoded, _) = msg.encode(&secret_key);
805            assert!(encoded.len() <= MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
806
807            let mut neighbours = Neighbours {
808                nodes: std::iter::repeat_with(|| rng_ipv6_record(&mut rng))
809                    .take(SAFE_MAX_DATAGRAM_NEIGHBOUR_RECORDS - 1)
810                    .collect(),
811                expire: rng.random(),
812            };
813            neighbours.nodes.push(rng_ipv4_record(&mut rng));
814            let msg = Message::Neighbours(neighbours);
815            let (encoded, _) = msg.encode(&secret_key);
816            assert!(encoded.len() <= MAX_PACKET_SIZE, "{} {msg:?}", encoded.len());
817        }
818    }
819
820    #[test]
821    fn test_encode_decode_message() {
822        let mut rng = rng();
823        for _ in 0..100 {
824            let msg = rng_message(&mut rng);
825            let (secret_key, pk) = secp256k1::generate_keypair(&mut rng);
826            let sender_id = pk2id(&pk);
827
828            let (buf, _) = msg.encode(&secret_key);
829
830            let packet = Message::decode(buf.as_ref()).unwrap();
831
832            assert_eq!(msg, packet.msg);
833            assert_eq!(sender_id, packet.node_id);
834        }
835    }
836
837    #[test]
838    fn decode_pong_packet() {
839        let packet = "2ad84c37327a06c2522cf7bc039621da89f68907441b755935bb308dc4cd17d6fe550e90329ad6a516ca7db18e08900067928a0dfa3b5c75d55a42c984497373698d98616662c048983ea85895ea2da765eabeb15525478384e106337bfd8ed50002f3c9843ed8cae682fd1c80a008ad4dead0922211df47593e7d837b2b23d13954285871ca23250ea594993ded84635690e5829670";
840        let data = hex::decode(packet).unwrap();
841        Message::decode(&data).unwrap();
842    }
843    #[test]
844    fn decode_ping_packet() {
845        let packet = "05ae5bf922cf2a93f97632a4ab0943dc252a0dab0c42d86dd62e5d91e1a0966e9b628fbf4763fdfbb928540460b797e6be2e7058a82f6083f6d2e7391bb021741459976d4152aa16bbee0c3609dcfac6668db1ef78b7ee9f8b4ced10dd5ae2900101df04cb8403d12d4f82765f82765fc9843ed8cae6828aa6808463569916829670";
846        let data = hex::decode(packet).unwrap();
847        Message::decode(&data).unwrap();
848    }
849
850    #[test]
851    fn encode_decode_enr_msg() {
852        let mut rng = rand::rng();
853        let key = SecretKey::new(&mut rng);
854        let ip = Ipv4Addr::new(127, 0, 0, 1);
855        let tcp = 3000;
856
857        let fork_id: ForkId = ForkId { hash: ForkHash([220, 233, 108, 45]), next: 0u64 };
858
859        let enr = {
860            let mut builder = Enr::builder();
861            builder.ip(ip.into());
862            builder.tcp4(tcp);
863            let mut buf = Vec::new();
864            let forkentry = EnrForkIdEntry { fork_id };
865            forkentry.encode(&mut buf);
866            builder.add_value_rlp("eth", buf.into());
867            builder.build(&key).unwrap()
868        };
869
870        let enr_response = EnrResponse { request_hash: B256::random(), enr };
871
872        let mut buf = Vec::new();
873        enr_response.encode(&mut buf);
874
875        let decoded = EnrResponse::decode(&mut &buf[..]).unwrap();
876
877        let fork_id_decoded = decoded.eth_fork_id().unwrap();
878        assert_eq!(fork_id, fork_id_decoded);
879    }
880
881    // test vector from the enr library rlp encoding tests
882    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#L1019>
883
884    #[test]
885    fn encode_known_rlp_enr() {
886        let valid_record = hex!(
887            "f884b8407098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c01826964827634826970847f00000189736563703235366b31a103ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31388375647082765f"
888        );
889        let signature = hex!(
890            "7098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c"
891        );
892        let expected_pubkey =
893            hex!("03ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd3138");
894
895        let enr = Enr::<SecretKey>::decode(&mut &valid_record[..]).unwrap();
896        let pubkey = enr.public_key().encode();
897
898        assert_eq!(enr.ip4(), Some(Ipv4Addr::new(127, 0, 0, 1)));
899        assert_eq!(enr.id(), Some(String::from("v4")));
900        assert_eq!(enr.udp4(), Some(DEFAULT_DISCOVERY_PORT));
901        assert_eq!(enr.tcp4(), None);
902        assert_eq!(enr.signature(), &signature[..]);
903        assert_eq!(pubkey.to_vec(), expected_pubkey);
904        assert!(enr.verify());
905
906        assert_eq!(&alloy_rlp::encode(&enr)[..], &valid_record[..]);
907
908        // ensure the length is equal
909        assert_eq!(enr.length(), valid_record.len());
910    }
911
912    // test vector from the enr library rlp encoding tests
913    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#L1019>
914    #[test]
915    fn decode_enr_rlp() {
916        let valid_record = hex!(
917            "f884b8407098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c01826964827634826970847f00000189736563703235366b31a103ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31388375647082765f"
918        );
919        let signature = hex!(
920            "7098ad865b00a582051940cb9cf36836572411a47278783077011599ed5cd16b76f2635f4e234738f30813a89eb9137e3e3df5266e3a1f11df72ecf1145ccb9c"
921        );
922        let expected_pubkey =
923            hex!("03ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd3138");
924
925        let mut valid_record_buf = valid_record.as_slice();
926        let enr = Enr::<SecretKey>::decode(&mut valid_record_buf).unwrap();
927        let pubkey = enr.public_key().encode();
928
929        // Byte array must be consumed after enr has finished decoding
930        assert!(valid_record_buf.is_empty());
931
932        assert_eq!(enr.ip4(), Some(Ipv4Addr::new(127, 0, 0, 1)));
933        assert_eq!(enr.id(), Some(String::from("v4")));
934        assert_eq!(enr.udp4(), Some(DEFAULT_DISCOVERY_PORT));
935        assert_eq!(enr.tcp4(), None);
936        assert_eq!(enr.signature(), &signature[..]);
937        assert_eq!(pubkey.to_vec(), expected_pubkey);
938        assert!(enr.verify());
939    }
940
941    // test for failing message decode
942    #[test]
943    fn decode_failing_packet() {
944        let packet = hex!(
945            "2467ab56952aedf4cfb8bb7830ddc8922d0f992185229919dad9de3841fe95d9b3a7b52459398235f6d3805644666d908b45edb3670414ed97f357afba51f71f7d35c1f45878ba732c3868b04ca42ff0ed347c99efcf3a5768afed68eb21ef960001db04c3808080c9840a480e8f82765f808466a9a06386019106833efe"
946        );
947
948        let _message = Message::decode(&packet[..]).unwrap();
949    }
950
951    // test for failing message decode
952    #[test]
953    fn decode_node() {
954        let packet = hex!("cb840000000082115c82115d");
955        let _message = NodeEndpoint::decode(&mut &packet[..]).unwrap();
956    }
957
958    // test vector from the enr library rlp encoding tests
959    // <https://github.com/sigp/enr/blob/e59dcb45ea07e423a7091d2a6ede4ad6d8ef2840/src/lib.rs#LL1206C35-L1206C35>
960    #[test]
961    fn encode_decode_enr_rlp() {
962        let key = SecretKey::new(&mut rand::rng());
963        let ip = Ipv4Addr::new(127, 0, 0, 1);
964        let tcp = 3000;
965
966        let enr = {
967            let mut builder = Enr::builder();
968            builder.ip(ip.into());
969            builder.tcp4(tcp);
970            builder.build(&key).unwrap()
971        };
972
973        let mut encoded_bytes = &alloy_rlp::encode(&enr)[..];
974        let decoded_enr = Enr::<SecretKey>::decode(&mut encoded_bytes).unwrap();
975
976        // Byte array must be consumed after enr has finished decoding
977        assert!(encoded_bytes.is_empty());
978
979        assert_eq!(decoded_enr, enr);
980        assert_eq!(decoded_enr.id(), Some("v4".into()));
981        assert_eq!(decoded_enr.ip4(), Some(ip));
982        assert_eq!(decoded_enr.tcp4(), Some(tcp));
983        assert_eq!(decoded_enr.public_key().encode(), key.public_key().encode());
984        assert!(decoded_enr.verify());
985    }
986
987    mod eip8 {
988        use super::*;
989        use alloy_primitives::b512;
990
991        fn junk_enr_request() -> Vec<u8> {
992            let mut buf = Vec::new();
993            // enr request is just an expiration
994            let expire: u64 = 123456;
995
996            // add some junk
997            let junk: u64 = 112233;
998
999            // rlp header encoding
1000            let payload_length = expire.length() + junk.length();
1001            alloy_rlp::Header { list: true, payload_length }.encode(&mut buf);
1002
1003            // fields
1004            expire.encode(&mut buf);
1005            junk.encode(&mut buf);
1006
1007            buf
1008        }
1009
1010        // checks that junk data at the end of the packet is discarded according to eip-8
1011        #[test]
1012        fn eip8_decode_enr_request() {
1013            let enr_request_with_junk = junk_enr_request();
1014
1015            let mut buf = enr_request_with_junk.as_slice();
1016            let decoded = EnrRequest::decode(&mut buf).unwrap();
1017            assert_eq!(decoded.expire, 123456);
1018        }
1019
1020        // checks that junk data at the end of the packet is discarded according to eip-8
1021        //
1022        // test vector from eip-8: https://eips.ethereum.org/EIPS/eip-8
1023        #[test]
1024        fn eip8_decode_findnode() {
1025            let findnode_with_junk = hex!(
1026                "c7c44041b9f7c7e41934417ebac9a8e1a4c6298f74553f2fcfdcae6ed6fe53163eb3d2b52e39fe91831b8a927bf4fc222c3902202027e5e9eb812195f95d20061ef5cd31d502e47ecb61183f74a504fe04c51e73df81f25c4d506b26db4517490103f84eb840ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31387574077f301b421bc84df7266c44e9e6d569fc56be00812904767bf5ccd1fc7f8443b9a35582999983999999280dc62cc8255c73471e0a61da0c89acdc0e035e260add7fc0c04ad9ebf3919644c91cb247affc82b69bd2ca235c71eab8e49737c937a2c396"
1027            );
1028
1029            let buf = findnode_with_junk.as_slice();
1030            let decoded = Message::decode(buf).unwrap();
1031
1032            let expected_id = hex!(
1033                "ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd31387574077f301b421bc84df7266c44e9e6d569fc56be00812904767bf5ccd1fc7f"
1034            );
1035            assert_matches!(decoded.msg, Message::FindNode(FindNode { id, expire: 1136239445 }) if id == expected_id);
1036        }
1037
1038        // checks that junk data at the end of the packet is discarded according to eip-8
1039        //
1040        // test vector from eip-8: https://eips.ethereum.org/EIPS/eip-8
1041        #[test]
1042        fn eip8_decode_neighbours() {
1043            let neighbours_with_junk = hex!(
1044                "c679fc8fe0b8b12f06577f2e802d34f6fa257e6137a995f6f4cbfc9ee50ed3710faf6e66f932c4c8d81d64343f429651328758b47d3dbc02c4042f0fff6946a50f4a49037a72bb550f3a7872363a83e1b9ee6469856c24eb4ef80b7535bcf99c0004f9015bf90150f84d846321163782115c82115db8403155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32f84984010203040101b840312c55512422cf9b8a4097e9a6ad79402e87a15ae909a4bfefa22398f03d20951933beea1e4dfa6f968212385e829f04c2d314fc2d4e255e0d3bc08792b069dbf8599020010db83c4d001500000000abcdef12820d05820d05b84038643200b172dcfef857492156971f0e6aa2c538d8b74010f8e140811d53b98c765dd2d96126051913f44582e8c199ad7c6d6819e9a56483f637feaac9448aacf8599020010db885a308d313198a2e037073488203e78203e8b8408dcab8618c3253b558d459da53bd8fa68935a719aff8b811197101a4b2b47dd2d47295286fc00cc081bb542d760717d1bdd6bec2c37cd72eca367d6dd3b9df738443b9a355010203b525a138aa34383fec3d2719a0"
1045            );
1046
1047            let buf = neighbours_with_junk.as_slice();
1048            let decoded = Message::decode(buf).unwrap();
1049
1050            let _ = NodeRecord {
1051                address: "99.33.22.55".parse().unwrap(),
1052                tcp_port: 4444,
1053                udp_port: 4445,
1054                id: b512!("3155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32"),
1055            }.length();
1056
1057            let expected_nodes: Vec<NodeRecord> = vec![
1058                NodeRecord {
1059                    address: "99.33.22.55".parse().unwrap(),
1060                    udp_port: 4444,
1061                    tcp_port: 4445,
1062                    id: b512!("3155e1427f85f10a5c9a7755877748041af1bcd8d474ec065eb33df57a97babf54bfd2103575fa829115d224c523596b401065a97f74010610fce76382c0bf32"),
1063                },
1064                NodeRecord {
1065                    address: "1.2.3.4".parse().unwrap(),
1066                    udp_port: 1,
1067                    tcp_port: 1,
1068                    id: b512!("312c55512422cf9b8a4097e9a6ad79402e87a15ae909a4bfefa22398f03d20951933beea1e4dfa6f968212385e829f04c2d314fc2d4e255e0d3bc08792b069db"),
1069                },
1070                NodeRecord {
1071                    address: "2001:db8:3c4d:15::abcd:ef12".parse().unwrap(),
1072                    udp_port: 3333,
1073                    tcp_port: 3333,
1074                    id: b512!("38643200b172dcfef857492156971f0e6aa2c538d8b74010f8e140811d53b98c765dd2d96126051913f44582e8c199ad7c6d6819e9a56483f637feaac9448aac"),
1075                },
1076                NodeRecord {
1077                    address: "2001:db8:85a3:8d3:1319:8a2e:370:7348".parse().unwrap(),
1078                    udp_port: 999,
1079                    tcp_port: 1000,
1080                    id: b512!("8dcab8618c3253b558d459da53bd8fa68935a719aff8b811197101a4b2b47dd2d47295286fc00cc081bb542d760717d1bdd6bec2c37cd72eca367d6dd3b9df73"),
1081                },
1082            ];
1083            assert_matches!(decoded.msg, Message::Neighbours(Neighbours { nodes, expire: 1136239445 }) if nodes == expected_nodes);
1084        }
1085
1086        // checks that junk data at the end of the packet is discarded according to eip-8, even
1087        // when the record and the junk together exceed the ENR size limit.
1088        #[test]
1089        fn eip8_decode_enr_response() {
1090            let key = SecretKey::new(&mut rand::rng());
1091            let enr = Enr::builder().build(&key).unwrap();
1092            let request_hash = B256::random();
1093            let junk = Bytes::from(vec![0xab; 300]);
1094
1095            let mut buf = Vec::new();
1096            let payload_length = request_hash.length() + enr.length() + junk.length();
1097            alloy_rlp::Header { list: true, payload_length }.encode(&mut buf);
1098            request_hash.encode(&mut buf);
1099            enr.encode(&mut buf);
1100            junk.encode(&mut buf);
1101
1102            let decoded = EnrResponse::decode(&mut buf.as_slice()).unwrap();
1103            assert_eq!(decoded, EnrResponse { request_hash, enr });
1104        }
1105    }
1106}