From 94f759761a078432dbf5026b5bec159c9d136a12 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Fri, 17 Feb 2023 12:04:22 +0100
Subject: [PATCH] fix(disc): enforce expiration timestamps by default (#1427)

---
 crates/net/discv4/src/config.rs | 16 +++++++++++++--
 crates/net/discv4/src/lib.rs    | 35 ++++++++++++++++++++++++++++-----
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/crates/net/discv4/src/config.rs b/crates/net/discv4/src/config.rs
index 0a8c14cbc..924a7b623 100644
--- a/crates/net/discv4/src/config.rs
+++ b/crates/net/discv4/src/config.rs
@@ -58,8 +58,10 @@ pub struct Discv4Config {
     pub enable_dht_random_walk: bool,
     /// Whether to automatically lookup peers.
     pub enable_lookup: bool,
-    /// Whether to enable EIP-868 extension
+    /// Whether to enforce EIP-868 extension
     pub enable_eip868: bool,
+    /// Whether to respect expiration timestamps in messages
+    pub enforce_expiration_timestamps: bool,
     /// Additional pairs to include in The [`Enr`](enr::Enr) if EIP-868 extension is enabled <https://eips.ethereum.org/EIPS/eip-868>
     pub additional_eip868_rlp_pairs: HashMap<Vec<u8>, Bytes>,
     /// If configured, try to resolve public ip
@@ -133,6 +135,7 @@ impl Default for Discv4Config {
             enable_dht_random_walk: true,
             enable_lookup: true,
             enable_eip868: true,
+            enforce_expiration_timestamps: true,
             additional_eip868_rlp_pairs: Default::default(),
             external_ip_resolver: Some(Default::default()),
             /// By default retry public IP using a 5min interval
@@ -209,12 +212,21 @@ impl Discv4ConfigBuilder {
         self
     }
 
-    /// Whether to enable EIP-868
+    /// Whether to enforce expiration timestamps in messages.
     pub fn enable_eip868(&mut self, enable_eip868: bool) -> &mut Self {
         self.config.enable_eip868 = enable_eip868;
         self
     }
 
+    /// Whether to enable EIP-868
+    pub fn enforce_expiration_timestamps(
+        &mut self,
+        enforce_expiration_timestamps: bool,
+    ) -> &mut Self {
+        self.config.enforce_expiration_timestamps = enforce_expiration_timestamps;
+        self
+    }
+
     /// Add another key value pair to include in the ENR
     pub fn add_eip868_pair(&mut self, key: impl AsRef<[u8]>, value: impl Encodable) -> &mut Self {
         let mut buf = BytesMut::new();
diff --git a/crates/net/discv4/src/lib.rs b/crates/net/discv4/src/lib.rs
index 1f9f474a8..ed52aef6a 100644
--- a/crates/net/discv4/src/lib.rs
+++ b/crates/net/discv4/src/lib.rs
@@ -353,8 +353,6 @@ pub struct Discv4Service {
     _tasks: JoinSet<()>,
     /// The routing table.
     kbuckets: KBucketsTable<NodeKey, NodeEntry>,
-    /// Whether to respect timestamps
-    check_timestamps: bool,
     /// Receiver for incoming messages
     ingress: IngressReceiver,
     /// Sender for sending outgoing messages
@@ -474,7 +472,6 @@ impl Discv4Service {
             pending_pings: Default::default(),
             pending_find_nodes: Default::default(),
             pending_enr_requests: Default::default(),
-            check_timestamps: false,
             commands_rx,
             update_listeners: Vec::with_capacity(1),
             lookup_interval: self_lookup_interval,
@@ -1313,7 +1310,7 @@ impl Discv4Service {
     /// Validate that given timestamp is not expired.
     fn ensure_not_expired(&self, expiration: u64) -> Result<(), ()> {
         let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
-        if self.check_timestamps && expiration < now {
+        if self.config.enforce_expiration_timestamps && expiration < now {
             debug!(target: "discv4", "Expired packet");
             return Err(())
         }
@@ -1979,7 +1976,7 @@ mod tests {
         let ping = Ping {
             from: rng_endpoint(&mut rng),
             to: rng_endpoint(&mut rng),
-            expire: 0,
+            expire: service.ping_expiration(),
             enr_sq: Some(rng.gen()),
         };
 
@@ -1997,6 +1994,34 @@ mod tests {
         };
     }
 
+    #[tokio::test]
+    async fn test_respect_ping_expiration() {
+        reth_tracing::init_test_tracing();
+        let mut rng = thread_rng();
+        let config = Discv4Config::builder().build();
+        let (_discv4, mut service) = create_discv4_with_config(config).await;
+
+        let v4: Ipv4Addr = "0.0.0.0".parse().unwrap();
+        let v6 = v4.to_ipv6_mapped();
+        let addr: SocketAddr = (v6, 30303).into();
+
+        let ping = Ping {
+            from: rng_endpoint(&mut rng),
+            to: rng_endpoint(&mut rng),
+            expire: SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs() - 1,
+            enr_sq: Some(rng.gen()),
+        };
+
+        let id = PeerId::random();
+        service.on_ping(ping, addr, id, H256::random());
+
+        let key = kad_key(id);
+        match service.kbuckets.entry(&key) {
+            kbucket::Entry::Absent(_) => {}
+            _ => unreachable!(),
+        };
+    }
+
     #[tokio::test]
     async fn test_single_lookups() {
         reth_tracing::init_test_tracing();