diff --git a/src/main/java/org/tron/p2p/discover/protocol/kad/NodeHandler.java b/src/main/java/org/tron/p2p/discover/protocol/kad/NodeHandler.java index de42c5f4..e4bbafba 100644 --- a/src/main/java/org/tron/p2p/discover/protocol/kad/NodeHandler.java +++ b/src/main/java/org/tron/p2p/discover/protocol/kad/NodeHandler.java @@ -12,6 +12,7 @@ import org.tron.p2p.discover.message.kad.NeighborsMessage; import org.tron.p2p.discover.message.kad.PingMessage; import org.tron.p2p.discover.message.kad.PongMessage; +import org.tron.p2p.discover.protocol.kad.table.KademliaOptions; import org.tron.p2p.discover.socket.UdpEvent; @Slf4j(topic = "net") @@ -24,6 +25,7 @@ public class NodeHandler { private AtomicInteger pingTrials = new AtomicInteger(3); private volatile boolean waitForPong = false; private volatile boolean waitForNeighbors = false; + private volatile long rcvFindNodeMsgTime; public NodeHandler(Node node, KadService kadService) { this.node = node; @@ -139,6 +141,11 @@ public void handleNeighbours(NeighborsMessage msg, InetSocketAddress sender) { } public void handleFindNode(FindNodeMessage msg) { + long now = System.currentTimeMillis(); + if (now - rcvFindNodeMsgTime < KademliaOptions.DISCOVER_CYCLE / 2) { + return; + } + rcvFindNodeMsgTime = now; List closest = kadService.getTable().getClosestNodes(msg.getTargetId()); sendNeighbours(closest, msg.getTimestamp()); } diff --git a/src/test/java/org/tron/p2p/discover/protocol/kad/NodeHandlerTest.java b/src/test/java/org/tron/p2p/discover/protocol/kad/NodeHandlerTest.java index 215e4192..2851b1d2 100644 --- a/src/test/java/org/tron/p2p/discover/protocol/kad/NodeHandlerTest.java +++ b/src/test/java/org/tron/p2p/discover/protocol/kad/NodeHandlerTest.java @@ -13,6 +13,8 @@ import java.lang.reflect.Constructor; import java.lang.reflect.Field; import java.net.InetSocketAddress; +import java.util.ArrayList; +import java.util.Collections; public class NodeHandlerTest { @@ -81,6 +83,29 @@ public void testChangeState() throws Exception { Assert.assertTrue(kadService.getTable().contains(replaceNode)); } + @Test + public void testHandleFindNode() throws Exception { + NodeHandler nodeHandler = new NodeHandler(new Node(new byte[1], "", "", 0), null); + Field field = nodeHandler.getClass().getDeclaredField("rcvFindNodeMsgTime"); + field.setAccessible(true); + long time = (Long) field.get(nodeHandler); + Assert.assertEquals(time, 0); + + try { + nodeHandler.handleFindNode(null); + } catch (Exception e) {} + + time = (Long) field.get(nodeHandler); + Assert.assertTrue(time > 0); + + try { + nodeHandler.handleFindNode(null); + } catch (Exception e) {} + + long time2 = (Long) field.get(nodeHandler); + Assert.assertEquals(time, time2); + } + @AfterClass public static void destroy() { kadService.close();