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);
  }