From d40d9e4d15f9fa2a5fb0db964d938797989b794e Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Sat, 11 Feb 2023 22:09:48 +0100
Subject: [PATCH] test(disc): add wrong to test (#1286)

---
 crates/net/discv4/src/lib.rs        | 47 ++++++++++++++++++++++++++---
 crates/net/discv4/src/test_utils.rs |  3 +-
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/crates/net/discv4/src/lib.rs b/crates/net/discv4/src/lib.rs
index 92e2bbd510..7be022d3fb 100644
--- a/crates/net/discv4/src/lib.rs
+++ b/crates/net/discv4/src/lib.rs
@@ -881,7 +881,8 @@ impl Discv4Service {
         // send the pong first, but the PONG and optionally PING don't need to be send in a
         // particular order
         let msg = Message::Pong(Pong {
-            to: ping.from,
+            // we use the actual address of the peer
+            to: record.into(),
             echo: hash,
             expire: ping.expire,
             enr_sq: self.enr_seq(),
@@ -2088,18 +2089,54 @@ mod tests {
         let _ = discv4.lookup_self().await;
     }
 
+    // sends a PING packet with wrong 'to' field and expects a PONG response.
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_check_wrong_to() {
+        reth_tracing::init_test_tracing();
+
+        let config = Discv4Config::builder().external_ip_resolver(None).build();
+        let (_discv4, mut service_1) = create_discv4_with_config(config.clone()).await;
+        let (_discv4, mut service_2) = create_discv4_with_config(config).await;
+
+        // ping node 2 with wrong to field
+        let mut ping = Ping {
+            from: service_1.local_node_record.into(),
+            to: service_2.local_node_record.into(),
+            expire: service_1.ping_expiration(),
+            enr_sq: service_1.enr_seq(),
+        };
+        ping.to.address = "192.0.2.0".parse().unwrap();
+
+        let echo_hash = service_1.send_packet(Message::Ping(ping), service_2.local_addr());
+        let ping_request = PingRequest {
+            sent_at: Instant::now(),
+            node: service_2.local_node_record,
+            echo_hash,
+            reason: PingReason::Initial,
+        };
+        service_1.pending_pings.insert(*service_2.local_peer_id(), ping_request);
+
+        // wait for the processed ping
+        let event = poll_fn(|cx| service_2.poll(cx)).await;
+        assert_eq!(event, Discv4Event::Ping);
+
+        // we now wait for PONG
+        let event = poll_fn(|cx| service_1.poll(cx)).await;
+        assert_eq!(event, Discv4Event::Pong);
+        // followed by a ping
+        let event = poll_fn(|cx| service_1.poll(cx)).await;
+        assert_eq!(event, Discv4Event::Ping);
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn test_check_ping_pong() {
         reth_tracing::init_test_tracing();
 
-        let config = Discv4Config::builder().build();
+        let config = Discv4Config::builder().external_ip_resolver(None).build();
         let (_discv4, mut service_1) = create_discv4_with_config(config.clone()).await;
         let (_discv4, mut service_2) = create_discv4_with_config(config).await;
 
-        service_1.local_enr_mut().address = IpAddr::V4(Ipv4Addr::UNSPECIFIED);
-        service_2.local_enr_mut().address = IpAddr::V4(Ipv4Addr::UNSPECIFIED);
         // send ping from 1 -> 2
-
         service_1.add_node(service_2.local_node_record);
 
         // wait for the processed ping
diff --git a/crates/net/discv4/src/test_utils.rs b/crates/net/discv4/src/test_utils.rs
index 839246a78d..0657ec98ab 100644
--- a/crates/net/discv4/src/test_utils.rs
+++ b/crates/net/discv4/src/test_utils.rs
@@ -216,9 +216,8 @@ pub async fn create_discv4_with_config(config: Discv4Config) -> (Discv4, Discv4S
     let socket = SocketAddr::from_str("0.0.0.0:0").unwrap();
     let (secret_key, pk) = SECP256K1.generate_keypair(&mut rng);
     let id = PeerId::from_slice(&pk.serialize_uncompressed()[1..]);
-    let external_addr = reth_net_nat::external_ip().await.unwrap_or_else(|| socket.ip());
     let local_enr =
-        NodeRecord { address: external_addr, tcp_port: socket.port(), udp_port: socket.port(), id };
+        NodeRecord { address: socket.ip(), tcp_port: socket.port(), udp_port: socket.port(), id };
     Discv4::bind(socket, local_enr, secret_key, config).await.unwrap()
 }