feat(net): optimize fetch inventory message check logic

tronprotocol · Jun 25, 2024 · 2c1cb1c · 2c1cb1c
1 parent d383177
commit 2c1cb1c
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 2 deletions.
diff --git a/framework/src/main/java/org/tron/core/net/messagehandler/FetchInvDataMsgHandler.java b/framework/src/main/java/org/tron/core/net/messagehandler/FetchInvDataMsgHandler.java
@@ -164,6 +164,10 @@ private void check(PeerConnection peer, FetchInvDataMessage fetchInvDataMsg) thr
             throw new P2pException(TypeEnum.BAD_MESSAGE,
                 "minBlockNum: " + minBlockNum + ", blockNum: " + blockNum);
           }
+          if (blockNum > peer.getLastSyncBlockId().getNum()) {
+            throw new P2pException(TypeEnum.BAD_MESSAGE,
+                "maxBlockNum: " + peer.getLastSyncBlockId().getNum() + ", blockNum: " + blockNum);
+          }
           if (peer.getSyncBlockIdCache().getIfPresent(hash) != null) {
             throw new P2pException(TypeEnum.BAD_MESSAGE,
                 new BlockId(hash).getString() + " is exist");

diff --git a/framework/src/test/java/org/tron/core/net/messagehandler/FetchInvDataMsgHandlerTest.java b/framework/src/test/java/org/tron/core/net/messagehandler/FetchInvDataMsgHandlerTest.java
@@ -21,8 +21,6 @@
 import org.tron.core.net.service.adv.AdvService;
 import org.tron.protos.Protocol;
 
-
-
 public class FetchInvDataMsgHandlerTest {
 
   @Test
@@ -62,4 +60,37 @@ public void testProcessMessage() throws Exception {
         new FetchInvDataMessage(blockIds, Protocol.Inventory.InventoryType.BLOCK));
     Assert.assertNotNull(syncBlockIdCache.getIfPresent(blockId));
   }
+
+  @Test
+  public void testSyncFetchCheck() {
+    BlockCapsule.BlockId blockId = new BlockCapsule.BlockId(Sha256Hash.ZERO_HASH, 10000L);
+    List<Sha256Hash> blockIds = new LinkedList<>();
+    blockIds.add(blockId);
+    FetchInvDataMessage msg =
+        new FetchInvDataMessage(blockIds, Protocol.Inventory.InventoryType.BLOCK);
+
+    PeerConnection peer = Mockito.mock(PeerConnection.class);
+    Mockito.when(peer.isNeedSyncFromUs()).thenReturn(true);
+    Cache<Item, Long> advInvSpread = CacheBuilder.newBuilder().maximumSize(100)
+        .expireAfterWrite(1, TimeUnit.HOURS).recordStats().build();
+    Mockito.when(peer.getAdvInvSpread()).thenReturn(advInvSpread);
+
+    FetchInvDataMsgHandler fetchInvDataMsgHandler = new FetchInvDataMsgHandler();
+
+    try {
+      Mockito.when(peer.getLastSyncBlockId())
+        .thenReturn(new BlockCapsule.BlockId(Sha256Hash.ZERO_HASH, 1000L));
+      fetchInvDataMsgHandler.processMessage(peer, msg);
+    } catch (Exception e) {
+      Assert.assertEquals(e.getMessage(), "maxBlockNum: 1000, blockNum: 10000");
+    }
+
+    try {
+      Mockito.when(peer.getLastSyncBlockId())
+        .thenReturn(new BlockCapsule.BlockId(Sha256Hash.ZERO_HASH, 20000L));
+      fetchInvDataMsgHandler.processMessage(peer, msg);
+    } catch (Exception e) {
+      Assert.assertEquals(e.getMessage(), "minBlockNum: 16000, blockNum: 10000");
+    }
+  }
 }