From ee226fc889ce97c48c373a8c33374577e66c9b2f Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Sun, 26 Mar 2023 06:42:08 -0400
Subject: [PATCH] feat: use dnsdisc peers for discv4 bootstrap (#1983)

---
 crates/net/discv4/src/lib.rs        | 10 ++++++++++
 crates/net/network/src/discovery.rs |  8 ++++++++
 2 files changed, 18 insertions(+)

diff --git a/crates/net/discv4/src/lib.rs b/crates/net/discv4/src/lib.rs
index 4e05ec4dd..ac4c28ad5 100644
--- a/crates/net/discv4/src/lib.rs
+++ b/crates/net/discv4/src/lib.rs
@@ -259,6 +259,12 @@ impl Discv4 {
         self.safe_send_to_service(cmd);
     }
 
+    /// Adds the node to the table, if it is not already present.
+    pub fn add_node(&self, node_record: NodeRecord) {
+        let cmd = Discv4Command::Add(node_record);
+        self.safe_send_to_service(cmd);
+    }
+
     /// Adds the peer and id to the ban list.
     ///
     /// This will prevent any future inclusion in the table
@@ -1413,6 +1419,9 @@ impl Discv4Service {
                 while let Poll::Ready(cmd) = rx.poll_recv(cx) {
                     if let Some(cmd) = cmd {
                         match cmd {
+                            Discv4Command::Add(enr) => {
+                                self.add_node(enr);
+                            }
                             Discv4Command::Lookup { node_id, tx } => {
                                 let node_id = node_id.unwrap_or(self.local_node_record.id);
                                 self.lookup_with(node_id, tx);
@@ -1606,6 +1615,7 @@ pub(crate) async fn receive_loop(udp: Arc<UdpSocket>, tx: IngressSender, local_i
 
 /// The commands sent from the frontend to the service
 enum Discv4Command {
+    Add(NodeRecord),
     SetTcpPort(u16),
     SetEIP868RLPPair { key: Vec<u8>, rlp: Bytes },
     Ban(PeerId, IpAddr),
diff --git a/crates/net/network/src/discovery.rs b/crates/net/network/src/discovery.rs
index 33287e54d..7842c8ab9 100644
--- a/crates/net/network/src/discovery.rs
+++ b/crates/net/network/src/discovery.rs
@@ -123,6 +123,13 @@ impl Discovery {
         self.local_enr.id
     }
 
+    /// Add a node to the discv4 table.
+    pub(crate) fn add_discv4_node(&self, node: NodeRecord) {
+        if let Some(discv4) = &self.discv4 {
+            discv4.add_node(node);
+        }
+    }
+
     /// Processes an incoming [NodeRecord] update from a discovery service
     fn on_node_record_update(&mut self, record: NodeRecord, fork_id: Option<ForkId>) {
         let id = record.id;
@@ -179,6 +186,7 @@ impl Discovery {
             while let Some(Poll::Ready(Some(update))) =
                 self.dns_discovery_updates.as_mut().map(|updates| updates.poll_next_unpin(cx))
             {
+                self.add_discv4_node(update.node_record);
                 self.on_node_record_update(update.node_record, update.fork_id);
             }