From 7f6aac87bebb95856957f62c3e43498fd826b35c Mon Sep 17 00:00:00 2001
From: Dmitrii Shmatko <zilm13@users.noreply.github.com>
Date: Tue, 24 Sep 2024 15:00:44 +0400
Subject: [PATCH] Add `lookupNode` by nodeId in DiscoverySystem (#181)

---
 .../beacon/discovery/DiscoverySystem.java     | 17 ++++++++++
 .../integration/DiscoveryIntegrationTest.java | 32 +++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/src/main/java/org/ethereum/beacon/discovery/DiscoverySystem.java b/src/main/java/org/ethereum/beacon/discovery/DiscoverySystem.java
index 992345136..6c74d5be3 100644
--- a/src/main/java/org/ethereum/beacon/discovery/DiscoverySystem.java
+++ b/src/main/java/org/ethereum/beacon/discovery/DiscoverySystem.java
@@ -5,6 +5,7 @@
 
 import java.util.Collection;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 import org.apache.logging.log4j.LogManager;
@@ -114,4 +115,20 @@ public Stream<NodeRecord> streamLiveNodes() {
   public CompletableFuture<Collection<NodeRecord>> searchForNewPeers() {
     return taskManager.searchForNewPeers();
   }
+
+  /**
+   * Lookup node in locally stored KBuckets by its nodeId. Allows lookup of local node record.
+   *
+   * @param nodeId NodeId, big endian UInt256 Node ID in bytes
+   * @return NodeRecord if any found
+   */
+  public Optional<NodeRecord> lookupNode(final Bytes nodeId) {
+    if (nodeId.equals(getLocalNodeRecord().getNodeId())) {
+      return Optional.of(getLocalNodeRecord());
+    }
+    return buckets
+        .streamClosestNodes(nodeId)
+        .findFirst()
+        .filter(node -> node.getNodeId().equals(nodeId));
+  }
 }
diff --git a/src/test/java/org/ethereum/beacon/discovery/integration/DiscoveryIntegrationTest.java b/src/test/java/org/ethereum/beacon/discovery/integration/DiscoveryIntegrationTest.java
index aa69b8770..94b44d87c 100644
--- a/src/test/java/org/ethereum/beacon/discovery/integration/DiscoveryIntegrationTest.java
+++ b/src/test/java/org/ethereum/beacon/discovery/integration/DiscoveryIntegrationTest.java
@@ -414,6 +414,38 @@ public NodeRecord fromRlp(RLPReader reader) {
     assertFalse(bootnodePingResult.isCompletedExceptionally());
   }
 
+  @Test
+  public void shouldLookupNodes() throws Exception {
+    final DiscoverySystem bootnode = createDiscoveryClient();
+    final DiscoverySystem node1 = createDiscoveryClient(bootnode.getLocalNodeRecord());
+    final DiscoverySystem node2 = createDiscoveryClient(bootnode.getLocalNodeRecord());
+
+    waitFor(
+        () -> {
+          waitFor(node1.searchForNewPeers(), 10);
+          waitFor(node2.searchForNewPeers(), 10);
+          assertThat(bootnode.lookupNode(bootnode.getLocalNodeRecord().getNodeId()))
+              .contains(bootnode.getLocalNodeRecord());
+          assertThat(bootnode.lookupNode(node1.getLocalNodeRecord().getNodeId()))
+              .contains(node1.getLocalNodeRecord());
+          assertThat(bootnode.lookupNode(node2.getLocalNodeRecord().getNodeId()))
+              .contains(node2.getLocalNodeRecord());
+          assertThat(bootnode.lookupNode(node1.getLocalNodeRecord().getNodeId().not())).isEmpty();
+
+          assertThat(node1.lookupNode(bootnode.getLocalNodeRecord().getNodeId()))
+              .contains(bootnode.getLocalNodeRecord());
+          assertThat(node1.lookupNode(node2.getLocalNodeRecord().getNodeId()))
+              .contains(node2.getLocalNodeRecord());
+          assertThat(node1.lookupNode(node1.getLocalNodeRecord().getNodeId().not())).isEmpty();
+
+          assertThat(node2.lookupNode(node1.getLocalNodeRecord().getNodeId()))
+              .contains(node1.getLocalNodeRecord());
+          assertThat(node2.lookupNode(bootnode.getLocalNodeRecord().getNodeId()))
+              .contains(bootnode.getLocalNodeRecord());
+          assertThat(node2.lookupNode(node1.getLocalNodeRecord().getNodeId().not())).isEmpty();
+        });
+  }
+
   private DiscoverySystem createDiscoveryClient(final NodeRecord... bootnodes) throws Exception {
     return createDiscoveryClient(true, bootnodes);
   }