Add lookupNode by nodeId in DiscoverySystem (#181)

Consensys · Sep 24, 2024 · 7f6aac8 · 7f6aac8
1 parent c66f3d7
commit 7f6aac8
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 0 deletions.
diff --git 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
@@ -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);
   }