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 streamLiveNodes() { public CompletableFuture> 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 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); }