feat: implement Ping message decoding

crates/core/src/rlp/decode.rs

@@ -1,10 +1,9 @@
-use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
-
 use super::{
     constants::{RLP_EMPTY_LIST, RLP_NULL},
     error::RLPDecodeError,
 };
 use bytes::{Bytes, BytesMut};
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 
 /// Trait for decoding RLP encoded slices of data.
 /// See <https://ethereum.org/en/developers/docs/data-structures-and-encoding/rlp/#rlp-decoding> for more information.
@@ -313,12 +312,34 @@ impl<T1: RLPDecode, T2: RLPDecode, T3: RLPDecode> RLPDecode for (T1, T2, T3) {
     }
 }
 
+impl<T1: RLPDecode, T2: RLPDecode, T3: RLPDecode, T4: RLPDecode> RLPDecode for (T1, T2, T3, T4) {
+    fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> {
+        if rlp.is_empty() {
+            return Err(RLPDecodeError::InvalidLength);
+        }
+        let (is_list, payload, input_rest) = decode_rlp_item(rlp)?;
+        if !is_list {
+            return Err(RLPDecodeError::MalformedData);
+        }
+        let (first, first_rest) = T1::decode_unfinished(payload)?;
+        let (second, second_rest) = T2::decode_unfinished(first_rest)?;
+        let (third, third_rest) = T3::decode_unfinished(second_rest)?;
+        let (fourth, fourth_rest) = T4::decode_unfinished(third_rest)?;
+        // check that there is no more data to decode after the third element.
+        if !fourth_rest.is_empty() {
+            return Err(RLPDecodeError::MalformedData);
+        }
+
+        Ok(((first, second, third, fourth), input_rest))
+    }
+}
+
 /// Decodes an RLP item from a slice of bytes.
 /// It returns a 3-element tuple with the following elements:
 /// - A boolean indicating if the item is a list or not.
 /// - The payload of the item, without its prefix.
 /// - The remaining bytes after the item.
-pub(crate) fn decode_rlp_item(data: &[u8]) -> Result<(bool, &[u8], &[u8]), RLPDecodeError> {
+pub fn decode_rlp_item(data: &[u8]) -> Result<(bool, &[u8], &[u8]), RLPDecodeError> {
     if data.is_empty() {
         return Err(RLPDecodeError::InvalidLength);
     }

crates/core/src/rlp/structs.rs

@@ -66,6 +66,20 @@ impl<'a> Decoder<'a> {
         Ok((field, updated_self))
     }
 
+    /// Returns Some(field) if there's some field to decode, otherwise returns None
+    pub fn decode_optional_field<T: RLPDecode>(self) -> (Option<T>, Self) {
+        match <T as RLPDecode>::decode_unfinished(self.payload) {
+            Ok((field, rest)) => {
+                let updated_self = Self {
+                    payload: rest,
+                    ..self
+                };
+                (Some(field), updated_self)
+            }
+            Err(_) => (None, self),
+        }
+    }
+
     pub fn finish(self) -> Result<&'a [u8], RLPDecodeError> {
         if self.payload.is_empty() {
             Ok(self.remaining)

crates/net/src/discv4.rs

@@ -1,18 +1,23 @@
-use std::net::IpAddr;
-
 use bytes::BufMut;
-use ethrex_core::rlp::{encode::RLPEncode, structs};
+use ethrex_core::rlp::{
+    decode::RLPDecode,
+    encode::RLPEncode,
+    error::RLPDecodeError,
+    structs::{self, Decoder},
+};
+use ethrex_core::H256;
 use k256::ecdsa::{signature::Signer, SigningKey};
+use std::net::IpAddr;
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 // TODO: remove when all variants are used
 // NOTE: All messages could have more fields than specified by the spec.
 // Those additional fields should be ignored, and the message must be accepted.
 #[allow(dead_code)]
 pub(crate) enum Message {
     /// A ping message. Should be responded to with a Pong message.
     Ping(PingMessage),
-    Pong(()),
+    Pong(PongMessage),
     FindNode(()),
     Neighbors(()),
     ENRRequest(()),
@@ -53,9 +58,28 @@ impl Message {
             Message::ENRResponse(_) => 0x06,
         }
     }
+    #[allow(unused)]
+    pub fn decode_with_header(encoded_msg: &[u8]) -> Result<Message, RLPDecodeError> {
+        let signature_len = 65;
+        let hash_len = 32;
+        let packet_index = signature_len + hash_len;
+        let packet_type = encoded_msg[packet_index];
+        let msg = &encoded_msg[packet_index + 1..];
+        match packet_type {
+            0x01 => {
+                let ping = PingMessage::decode(msg)?;
+                Ok(Message::Ping(ping))
+            }
+            0x02 => {
+                let pong = PongMessage::decode(msg)?;
+                Ok(Message::Pong(pong))
+            }
+            _ => todo!(),
+        }
+    }
 }
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) struct Endpoint {
     pub ip: IpAddr,
     pub udp_port: u16,
@@ -72,7 +96,23 @@ impl RLPEncode for Endpoint {
     }
 }
 
-#[derive(Debug, Clone, Copy)]
+impl RLPDecode for Endpoint {
+    fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> {
+        let decoder = Decoder::new(rlp)?;
+        let (ip, decoder) = decoder.decode_field("ip")?;
+        let (udp_port, decoder) = decoder.decode_field("udp_port")?;
+        let (tcp_port, decoder) = decoder.decode_field("tcp_port")?;
+        let remaining = decoder.finish()?;
+        let endpoint = Endpoint {
+            ip,
+            udp_port,
+            tcp_port,
+        };
+        Ok((endpoint, remaining))
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) struct PingMessage {
     /// The Ping message version. Should be set to 4, but mustn't be enforced.
     version: u8,
@@ -120,12 +160,82 @@ impl RLPEncode for PingMessage {
     }
 }
 
+impl RLPDecode for PingMessage {
+    fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> {
+        let decoder = Decoder::new(rlp)?;
+        let (_version, decoder): (u8, Decoder) = decoder.decode_field("version")?;
+        let (from, decoder) = decoder.decode_field("from")?;
+        let (to, decoder) = decoder.decode_field("to")?;
+        let (expiration, decoder) = decoder.decode_field("expiration")?;
+        let (enr_seq, decoder) = decoder.decode_optional_field();
+
+        let ping = PingMessage::new(from, to, expiration);
+
+        let ping = if let Some(seq) = enr_seq {
+            ping.with_enr_seq(seq)
+        } else {
+            ping
+        };
+
+        let remaining = decoder.finish()?;
+        Ok((ping, remaining))
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct PongMessage {
+    /// The endpoint of the receiver.
+    to: Endpoint,
+    /// The hash of the corresponding ping packet.
+    ping_hash: H256,
+    /// The expiration time of the message. If the message is older than this time,
+    /// it shouldn't be responded to.
+    expiration: u64,
+    /// The ENR sequence number of the sender. This field is optional.
+    enr_seq: Option<u64>,
+}
+
+impl PongMessage {
+    #[allow(unused)]
+    pub fn new(to: Endpoint, ping_hash: H256, expiration: u64) -> Self {
+        Self {
+            to,
+            ping_hash,
+            expiration,
+            enr_seq: None,
+        }
+    }
+
+    // TODO: remove when used
+    #[allow(unused)]
+    pub fn with_enr_seq(self, enr_seq: u64) -> Self {
+        Self {
+            enr_seq: Some(enr_seq),
+            ..self
+        }
+    }
+}
+
+impl RLPDecode for PongMessage {
+    fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> {
+        let decoder = Decoder::new(rlp)?;
+        let (to, decoder) = decoder.decode_field("to")?;
+        let (ping_hash, decoder) = decoder.decode_field("ping_hash")?;
+        let (expiration, decoder) = decoder.decode_field("expiration")?;
+        let (enr_seq, decoder) = decoder.decode_field("enr_seq")?;
+        let remaining = decoder.finish()?;
+        let pong = PongMessage::new(to, ping_hash, expiration).with_enr_seq(enr_seq);
+
+        Ok((pong, remaining))
+    }
+}
+
 #[cfg(test)]
 mod tests {
-    use std::{fmt::Write, str::FromStr};
-
     use super::*;
     use keccak_hash::H256;
+    use std::num::ParseIntError;
+    use std::{fmt::Write, str::FromStr};
 
     fn to_hex(bytes: &[u8]) -> String {
         bytes.iter().fold(String::new(), |mut buf, b| {
@@ -175,7 +285,107 @@ mod tests {
         let signature = "34f486e4e92f2fdf592912aa77ad51db532dd7f9b426092384c9c2e9919414fd480d57f4f3b2b1964ed6eb1c94b1e4b9f6bfe9b44b1d1ac3d94c38c4cce915bc01";
         let pkt_type = "02";
         let msg = "f7c984bebfbc3982765f80a03e1bf98f025f98d54ed2f61bbef63b6b46f50e12d7b937d6bdea19afd640be2384667d9af086018cf3c3bcdd";
-        let _encoded_packet = [hash, signature, pkt_type, msg].concat();
-        // TODO
+        let encoded_packet = [hash, signature, pkt_type, msg].concat();
+
+        let decoded = Message::decode_with_header(
+            &decode_hex(&encoded_packet).expect("Failed while parsing encoded_packet"),
+        )
+        .unwrap();
+        let to = Endpoint {
+            ip: IpAddr::from_str("190.191.188.57").unwrap(),
+            udp_port: 30303,
+            tcp_port: 0,
+        };
+        let expiration: u64 = 1719507696;
+        let ping_hash: H256 =
+            H256::from_str("3e1bf98f025f98d54ed2f61bbef63b6b46f50e12d7b937d6bdea19afd640be23")
+                .unwrap();
+        let enr_seq = 1704896740573;
+        let expected =
+            Message::Pong(PongMessage::new(to, ping_hash, expiration).with_enr_seq(enr_seq));
+        assert_eq!(decoded, expected);
+    }
+
+    pub fn decode_hex(s: &str) -> Result<Vec<u8>, ParseIntError> {
+        (0..s.len())
+            .step_by(2)
+            .map(|i| u8::from_str_radix(&s[i..i + 2], 16))
+            .collect()
+    }
+
+    #[test]
+    fn test_decode_ping_message() {
+        let expiration: u64 = 17195043770;
+
+        let from = Endpoint {
+            ip: IpAddr::from_str("1.2.3.4").unwrap(),
+            udp_port: 1613,
+            tcp_port: 6363,
+        };
+        let to = Endpoint {
+            ip: IpAddr::from_str("255.255.2.5").unwrap(),
+            udp_port: 3063,
+            tcp_port: 0,
+        };
+
+        let msg = Message::Ping(PingMessage::new(from, to, expiration));
+
+        let key_bytes =
+            H256::from_str("577d8278cc7748fad214b5378669b420f8221afb45ce930b7f22da49cbc545f3")
+                .unwrap();
+        let signer = SigningKey::from_slice(key_bytes.as_bytes()).unwrap();
+
+        let mut buf = Vec::new();
+
+        msg.encode_with_header(&mut buf, signer.clone());
+        let result = Message::decode_with_header(&buf).expect("Failed decoding PingMessage");
+        assert_eq!(result, msg);
+    }
+
+    #[test]
+    fn test_decode_ping_message_with_enr_seq() {
+        let expiration: u64 = 17195043770;
+
+        let from = Endpoint {
+            ip: IpAddr::from_str("1.2.3.4").unwrap(),
+            udp_port: 1613,
+            tcp_port: 6363,
+        };
+        let to = Endpoint {
+            ip: IpAddr::from_str("255.255.2.5").unwrap(),
+            udp_port: 3063,
+            tcp_port: 0,
+        };
+
+        let enr_seq = 1704896740573;
+        let msg = Message::Ping(PingMessage::new(from, to, expiration).with_enr_seq(enr_seq));
+
+        let key_bytes =
+            H256::from_str("577d8278cc7748fad214b5378669b420f8221afb45ce930b7f22da49cbc545f3")
+                .unwrap();
+        let signer = SigningKey::from_slice(key_bytes.as_bytes()).unwrap();
+
+        let mut buf = Vec::new();
+
+        msg.encode_with_header(&mut buf, signer.clone());
+        let result = Message::decode_with_header(&buf).expect("Failed decoding PingMessage");
+        assert_eq!(result, msg);
+    }
+
+    #[test]
+    fn test_decode_endpoint() {
+        let endpoint = Endpoint {
+            ip: IpAddr::from_str("255.255.2.5").unwrap(),
+            udp_port: 3063,
+            tcp_port: 0,
+        };
+
+        let encoded = {
+            let mut buf = vec![];
+            (endpoint).encode(&mut buf);
+            buf
+        };
+        let decoded = Endpoint::decode(&encoded).expect("Failed decoding Endpoint");
+        assert_eq!(endpoint, decoded);
     }
 }

crates/net/src/lib.rs

@@ -1,12 +1,11 @@
 pub(crate) mod discv4;
 
 use std::{
-    fmt::Write,
     net::SocketAddr,
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
 
-use discv4::{Endpoint, PingMessage};
+use discv4::{Endpoint, Message, PingMessage};
 use k256::{ecdsa::SigningKey, elliptic_curve::rand_core::OsRng};
 use tokio::{
     net::{TcpSocket, UdpSocket},
@@ -36,15 +35,8 @@ async fn discover_peers(udp_addr: SocketAddr) {
 
     let (read, from) = udp_socket.recv_from(&mut buf).await.unwrap();
     info!("Received {read} bytes from {from}");
-    info!("Message: {}", to_hex(&buf[..read]));
-}
-
-// TODO: maybe remove this
-fn to_hex(bytes: &[u8]) -> String {
-    bytes.iter().fold(String::new(), |mut buf, b| {
-        let _ = write!(&mut buf, "{b:02x}");
-        buf
-    })
+    let msg = Message::decode_with_header(&buf[..read]);
+    info!("Message: {:?}", msg);
 }
 
 async fn ping(socket: &UdpSocket, local_addr: SocketAddr, to_addr: SocketAddr) {