From 0779c0091cac1b184b6c5755ebd844b63661497a Mon Sep 17 00:00:00 2001 From: zeusoo001 Date: Tue, 25 Jun 2024 18:53:38 +0800 Subject: [PATCH] feat(net): optimize fetch inventory message check logic --- .../FetchInvDataMsgHandler.java | 4 +++ .../FetchInvDataMsgHandlerTest.java | 35 +++++++++++++++++-- 2 files changed, 37 insertions(+), 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 index 5e797c084b3..5415ea435e3 100644 --- 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 index 404d275276a..5fd6d6725ba 100644 --- 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 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 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"); + } + } }