feat: implement rlpx handshake

crates/core/src/rlp/encode.rs

@@ -5,6 +5,14 @@ use tinyvec::ArrayVec;
 
 use super::constants::RLP_NULL;
 
+/// Function for encoding a value to RLP.
+/// For encoding the value into a buffer directly, use [`RLPEncode::encode`].
+pub fn encode<T: RLPEncode>(value: T) -> Vec<u8> {
+    let mut buf = Vec::new();
+    value.encode(&mut buf);
+    buf
+}
+
 pub trait RLPEncode {
     fn encode(&self, buf: &mut dyn BufMut);
 

crates/net/Cargo.toml

@@ -11,5 +11,16 @@ ethereum_rust-core.workspace = true
 tracing.workspace = true
 tokio.workspace = true
 bytes.workspace = true
-k256 = "0.13.3"
+
+k256 = { version = "0.13.3", features = ["ecdh"] }
 keccak-hash = "0.10.0"
+
+# RLPx
+concat-kdf = "0.1.0"
+hmac = "0.12.1"
+aes = "0.8.4"
+ctr = "0.9.2"
+rand = "0.8.5"
+
+[dev-dependencies]
+hex-literal = "0.4.1"

crates/net/src/discv4.rs

@@ -7,9 +7,10 @@ use ethereum_rust_core::rlp::{
 };
 use ethereum_rust_core::{H256, H512, H520};
 use k256::ecdsa::SigningKey;
-use std::net::IpAddr;
+use std::net::{IpAddr, SocketAddr};
 
 #[allow(unused)]
+#[derive(Debug)]
 pub struct Packet {
     hash: H256,
     signature: H520,
@@ -145,6 +146,12 @@ pub(crate) struct Endpoint {
     pub tcp_port: u16,
 }
 
+impl Endpoint {
+    pub fn tcp_address(&self) -> Option<SocketAddr> {
+        (self.tcp_port != 0).then_some(SocketAddr::new(self.ip, self.tcp_port))
+    }
+}
+
 impl RLPEncode for Endpoint {
     fn encode(&self, buf: &mut dyn BufMut) {
         structs::Encoder::new(buf)
@@ -178,12 +185,12 @@ pub(crate) struct PingMessage {
     /// The endpoint of the sender.
     pub from: Endpoint,
     /// The endpoint of the receiver.
-    to: Endpoint,
+    pub to: Endpoint,
     /// The expiration time of the message. If the message is older than this time,
     /// it shouldn't be responded to.
-    expiration: u64,
+    pub expiration: u64,
     /// The ENR sequence number of the sender. This field is optional.
-    enr_seq: Option<u64>,
+    pub enr_seq: Option<u64>,
 }
 
 impl PingMessage {
@@ -222,10 +229,10 @@ impl RLPEncode for PingMessage {
 #[derive(Debug, PartialEq, Eq)]
 pub(crate) struct FindNodeMessage {
     /// The target is a 64-byte secp256k1 public key.
-    target: H512,
+    pub target: H512,
     /// The expiration time of the message. If the message is older than this time,
     /// it shouldn't be responded to.
-    expiration: u64,
+    pub expiration: u64,
 }
 
 impl FindNodeMessage {
@@ -280,14 +287,14 @@ impl RLPDecode for PingMessage {
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(crate) struct PongMessage {
     /// The endpoint of the receiver.
-    to: Endpoint,
+    pub to: Endpoint,
     /// The hash of the corresponding ping packet.
-    ping_hash: H256,
+    pub 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,
+    pub expiration: u64,
     /// The ENR sequence number of the sender. This field is optional.
-    enr_seq: Option<u64>,
+    pub enr_seq: Option<u64>,
 }
 
 impl PongMessage {

crates/net/src/lib.rs

@@ -1,40 +1,46 @@
-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},
 };
+
+use bootnode::BootNode;
+use discv4::{Endpoint, FindNodeMessage, Message, Packet, PingMessage, PongMessage};
+use ethereum_rust_core::H512;
+use k256::{
+    ecdsa::{RecoveryId, Signature, SigningKey, VerifyingKey},
+    elliptic_curve::{rand_core::OsRng, sec1::ToEncodedPoint, PublicKey},
+    SecretKey,
+};
+use kademlia::{KademliaTable, PeerData};
+use keccak_hash::H256;
+use rlpx::ecies::RLPxConnection;
 use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
     net::{TcpSocket, UdpSocket},
     try_join,
 };
-use tracing::info;
+use tracing::{info, warn};
+
+pub mod bootnode;
+pub(crate) mod discv4;
+pub(crate) mod kademlia;
+pub mod rlpx;
 
 const MAX_DISC_PACKET_SIZE: usize = 1280;
 
 pub async fn start_network(udp_addr: SocketAddr, tcp_addr: SocketAddr, bootnodes: Vec<BootNode>) {
     info!("Starting discovery service at {udp_addr}");
     info!("Listening for requests at {tcp_addr}");
+    let signer = SigningKey::random(&mut OsRng);
 
-    let discovery_handle = tokio::spawn(discover_peers(udp_addr, bootnodes));
-    let server_handle = tokio::spawn(serve_requests(tcp_addr));
+    let discovery_handle = tokio::spawn(discover_peers(udp_addr, signer.clone(), bootnodes));
+    let server_handle = tokio::spawn(serve_requests(tcp_addr, signer));
     try_join!(discovery_handle, server_handle).unwrap();
 }
 
-async fn discover_peers(udp_addr: SocketAddr, bootnodes: Vec<BootNode>) {
+async fn discover_peers(udp_addr: SocketAddr, signer: SigningKey, 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..]);
@@ -147,9 +153,73 @@ async fn pong(socket: &UdpSocket, to_addr: SocketAddr, ping_hash: H256, signer:
     socket.send_to(&buf, to_addr).await.unwrap();
 }
 
-async fn serve_requests(tcp_addr: SocketAddr) {
+async fn serve_requests(tcp_addr: SocketAddr, signer: SigningKey) {
+    let secret_key: SecretKey = signer.clone().into();
     let tcp_socket = TcpSocket::new_v4().unwrap();
     tcp_socket.bind(tcp_addr).unwrap();
+
+    let mut udp_addr = tcp_addr;
+    udp_addr.set_port(tcp_addr.port() + 1);
+    let udp_socket = UdpSocket::bind(udp_addr).await.unwrap();
+
+    // Try contacting a known peer
+    // TODO: this is just an example, and we should do this dynamically
+    let str_udp_addr = "127.0.0.1:51311";
+
+    let udp_addr: SocketAddr = str_udp_addr.parse().unwrap();
+
+    let mut buf = vec![0; MAX_DISC_PACKET_SIZE];
+
+    let (mut msg, sig_bytes, endpoint) = loop {
+        ping(&udp_socket, tcp_addr, udp_addr, &signer).await;
+
+        let (read, from) = udp_socket.recv_from(&mut buf).await.unwrap();
+        info!("Received {read} bytes from {from}");
+        let packet = Packet::decode(&buf[..read]).unwrap();
+        info!("Message: {:?}", packet);
+
+        match packet.get_message() {
+            Message::Pong(pong) => {
+                break (&buf[32 + 65..read], &buf[32..32 + 65], pong.to);
+            }
+            Message::Ping(ping) => {
+                break (&buf[32 + 65..read], &buf[32..32 + 65], ping.from);
+            }
+            _ => {
+                warn!("Unexpected message type");
+            }
+        };
+    };
+
+    let digest = keccak_hash::keccak_buffer(&mut msg).unwrap();
+    let signature = &Signature::from_bytes(sig_bytes[..64].into()).unwrap();
+    let rid = RecoveryId::from_byte(sig_bytes[64]).unwrap();
+
+    let peer_pk = VerifyingKey::recover_from_prehash(&digest.0, signature, rid).unwrap();
+
+    let mut conn = RLPxConnection::random();
+    let mut auth_message = vec![];
+    conn.encode_auth_message(&secret_key, &peer_pk.into(), &mut auth_message);
+
+    let tcp_addr = "127.0.0.1:58617";
+    // NOTE: for some reason kurtosis peers don't publish their active TCP port
+    let tcp_addr = endpoint.tcp_address().unwrap_or(tcp_addr.parse().unwrap());
+
+    let mut stream = TcpSocket::new_v4()
+        .unwrap()
+        .connect(tcp_addr)
+        .await
+        .unwrap();
+
+    stream.write_all(&auth_message).await.unwrap();
+    info!("Sent auth message correctly!");
+    stream.read_exact(&mut buf[..2]).await.unwrap();
+    let auth_data = buf[..2].try_into().unwrap();
+    let msg_size = u16::from_be_bytes(auth_data) as usize;
+    let msg = &mut buf[..msg_size];
+    stream.read_exact(msg).await.unwrap();
+    conn.decode_ack_message(&secret_key, msg, auth_data);
+    info!("Completed handshake!");
 }
 
 #[cfg(test)]

crates/net/src/rlpx.rs

@@ -0,0 +1 @@
+pub mod ecies;