feat: store peer data

crates/net/src/discv4.rs

@@ -340,8 +340,8 @@ impl RLPDecode for PongMessage {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct NeighborsMessage {
     // nodes is the list of neighbors
-    nodes: Vec<Node>,
-    expiration: u64,
+    pub nodes: Vec<Node>,
+    pub expiration: u64,
 }
 
 impl NeighborsMessage {
@@ -352,10 +352,10 @@ impl NeighborsMessage {
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) struct Node {
-    ip: IpAddr,
-    udp_port: u16,
-    tcp_port: u16,
-    node_id: H512,
+    pub ip: IpAddr,
+    pub udp_port: u16,
+    pub tcp_port: u16,
+    pub node_id: H512,
 }
 
 impl RLPDecode for NeighborsMessage {

crates/net/src/kademlia.rs

@@ -0,0 +1,65 @@
+use crate::discv4::Node;
+use ethereum_rust_core::{H512, U256};
+use keccak_hash::keccak;
+use std::net::IpAddr;
+
+const MAX_NODES_PER_BUCKET: usize = 16;
+const NUMBER_OF_BUCKETS: usize = 256;
+
+#[derive(Debug)]
+pub struct KademliaTable {
+    local_node_id: H512,
+    buckets: Vec<Vec<PeerData>>,
+}
+
+impl KademliaTable {
+    pub fn new(local_node_id: H512) -> Self {
+        let buckets: Vec<Vec<PeerData>> = vec![vec![]; NUMBER_OF_BUCKETS];
+        Self {
+            local_node_id,
+            buckets,
+        }
+    }
+
+    pub fn insert(&mut self, peer: PeerData) {
+        let bucket_number = bucket_number(self.local_node_id, peer.node_id);
+        let bucket = &mut self.buckets[bucket_number];
+        if bucket.len() == MAX_NODES_PER_BUCKET {
+            // TODO: revalidate least recently seen node as described in
+            // <https://github.com/ethereum/devp2p/blob/master/discv4.md#kademlia-table>
+            bucket.pop();
+        }
+        bucket.push(peer);
+    }
+}
+
+/// Computes the distance between two nodes according to the discv4 protocol
+/// and returns the corresponding bucket number
+/// <https://github.com/ethereum/devp2p/blob/master/discv4.md#node-identities>
+pub fn bucket_number(node_id_1: H512, node_id_2: H512) -> usize {
+    let hash_1 = keccak(node_id_1);
+    let hash_2 = keccak(node_id_2);
+    let xor = hash_1 ^ hash_2;
+    let distance = U256::from_big_endian(xor.as_bytes());
+    distance.bits() - 1
+}
+
+#[derive(Clone, Debug)]
+#[allow(unused)]
+pub struct PeerData {
+    pub ip: IpAddr,
+    pub udp_port: u16,
+    pub tcp_port: u16,
+    pub node_id: H512,
+}
+
+impl From<Node> for PeerData {
+    fn from(node: Node) -> Self {
+        Self {
+            ip: node.ip,
+            udp_port: node.udp_port,
+            tcp_port: node.tcp_port,
+            node_id: node.node_id,
+        }
+    }
+}

crates/net/src/lib.rs

@@ -1,10 +1,16 @@
+pub mod bootnode;
 pub(crate) mod discv4;
+pub(crate) mod kademlia;
 use discv4::{Endpoint, FindNodeMessage, Message, Packet, PingMessage, PongMessage};
 use ethereum_rust_core::H512;
 use k256::elliptic_curve::sec1::ToEncodedPoint;
 use k256::elliptic_curve::PublicKey;
 use k256::{ecdsa::SigningKey, elliptic_curve::rand_core::OsRng};
 use keccak_hash::H256;
+
+use bootnode::BootNode;
+use kademlia::{KademliaTable, PeerData};
+use std::vec;
 use std::{
     net::SocketAddr,
     time::{Duration, SystemTime, UNIX_EPOCH},
@@ -14,8 +20,6 @@ use tokio::{
     try_join,
 };
 use tracing::info;
-use types::BootNode;
-pub mod types;
 
 const MAX_DISC_PACKET_SIZE: usize = 1280;
 
@@ -31,6 +35,10 @@ pub async fn start_network(udp_addr: SocketAddr, tcp_addr: SocketAddr, bootnodes
 async fn discover_peers(udp_addr: SocketAddr, bootnodes: Vec<BootNode>) {
     let udp_socket = UdpSocket::bind(udp_addr).await.unwrap();
     let signer = SigningKey::random(&mut OsRng);
+    let public_key = PublicKey::from(signer.verifying_key());
+    let encoded = public_key.to_encoded_point(false);
+    let local_node_id = H512::from_slice(&encoded.as_bytes()[1..]);
+
     let bootnode = match bootnodes.first() {
         Some(b) => b,
         None => {
@@ -41,16 +49,30 @@ async fn discover_peers(udp_addr: SocketAddr, bootnodes: Vec<BootNode>) {
     ping(&udp_socket, udp_addr, bootnode.socket_address, &signer).await;
 
     let mut buf = vec![0; MAX_DISC_PACKET_SIZE];
+    let mut table = KademliaTable::new(local_node_id);
     loop {
         let (read, from) = udp_socket.recv_from(&mut buf).await.unwrap();
         let packet = Packet::decode(&buf[..read]).unwrap();
         let msg = packet.get_message();
         info!("Received {read} bytes from {from}");
         info!("Message: {:?}", msg);
-        if let Message::Ping(_) = msg {
-            let ping_hash = packet.get_hash();
-            pong(&udp_socket, from, ping_hash, &signer).await;
-            find_node(&udp_socket, from, &signer).await;
+
+        match msg {
+            Message::Ping(_) => {
+                let ping_hash = packet.get_hash();
+                pong(&udp_socket, from, ping_hash, &signer).await;
+                find_node(&udp_socket, from, &signer).await;
+            }
+            Message::Neighbors(neighbors_msg) => {
+                let nodes = &neighbors_msg.nodes;
+                for node in nodes {
+                    let peer_data = PeerData::from(*node);
+                    table.insert(peer_data);
+                    let node_addr = SocketAddr::new(node.ip, node.udp_port);
+                    ping(&udp_socket, udp_addr, node_addr, &signer).await;
+                }
+            }
+            _ => {}
         }
     }
 }
@@ -81,7 +103,6 @@ async fn ping(
     };
 
     let ping: discv4::Message = discv4::Message::Ping(PingMessage::new(from, to, expiration));
-
     ping.encode_with_header(&mut buf, signer);
     socket.send_to(&buf, to_addr).await.unwrap();
 }
@@ -130,3 +151,18 @@ async fn serve_requests(tcp_addr: SocketAddr) {
     let tcp_socket = TcpSocket::new_v4().unwrap();
     tcp_socket.bind(tcp_addr).unwrap();
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use kademlia::bucket_number;
+    use std::str::FromStr;
+    #[test]
+    fn bucket_number_works_as_expected() {
+        let node_id_1 = H512::from_str("4dc429669029ceb17d6438a35c80c29e09ca2c25cc810d690f5ee690aa322274043a504b8d42740079c4f4cef50777c991010208b333b80bee7b9ae8e5f6b6f0").unwrap();
+        let node_id_2 = H512::from_str("034ee575a025a661e19f8cda2b6fd8b2fd4fe062f6f2f75f0ec3447e23c1bb59beb1e91b2337b264c7386150b24b621b8224180c9e4aaf3e00584402dc4a8386").unwrap();
+        let expected_bucket = 255;
+        let result = bucket_number(node_id_1, node_id_2);
+        assert_eq!(result, expected_bucket);
+    }
+}

ethereum_rust/src/cli.rs

@@ -1,5 +1,5 @@
 use clap::{Arg, ArgAction, Command};
-use ethereum_rust_net::types::BootNode;
+use ethereum_rust_net::bootnode::BootNode;
 
 pub fn cli() -> Command {
     Command::new("ethereum_rust")

ethereum_rust/src/main.rs

@@ -1,5 +1,5 @@
 use ethereum_rust_core::types::Genesis;
-use ethereum_rust_net::types::BootNode;
+use ethereum_rust_net::bootnode::BootNode;
 use std::{
     io::{self, BufReader},
     net::{SocketAddr, ToSocketAddrs},