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 5415ea435e3..8ee514b1bb7 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 @@ -55,7 +55,9 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep FetchInvDataMessage fetchInvDataMsg = (FetchInvDataMessage) msg; - check(peer, fetchInvDataMsg); + boolean isAdv = isAdvInv(peer, fetchInvDataMsg); + + check(peer, fetchInvDataMsg, isAdv); InventoryType type = fetchInvDataMsg.getInventoryType(); List transactions = Lists.newArrayList(); @@ -64,6 +66,10 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep for (Sha256Hash hash : fetchInvDataMsg.getHashList()) { Item item = new Item(hash, type); + if (isAdv) { + peer.getAdvInvSpread().invalidate(item); + } + Message message = advService.getMessage(item); if (message == null) { try { @@ -127,7 +133,21 @@ private void sendPbftCommitMessage(PeerConnection peer, BlockCapsule blockCapsul } } - private void check(PeerConnection peer, FetchInvDataMessage fetchInvDataMsg) throws P2pException { + public boolean isAdvInv(PeerConnection peer, FetchInvDataMessage msg) { + MessageTypes type = msg.getInvMessageType(); + if (type == MessageTypes.TRX) { + return true; + } + for (Sha256Hash hash : msg.getHashList()) { + if (peer.getAdvInvSpread().getIfPresent(new Item(hash, InventoryType.BLOCK)) == null) { + return false; + } + } + return true; + } + + private void check(PeerConnection peer, FetchInvDataMessage fetchInvDataMsg, + boolean isAdv) throws P2pException { MessageTypes type = fetchInvDataMsg.getInvMessageType(); if (type == MessageTypes.TRX) { @@ -144,38 +164,30 @@ private void check(PeerConnection peer, FetchInvDataMessage fetchInvDataMsg) thr + "maxCount: {}, fetchCount: {}, peer: {}", maxCount, fetchCount, peer.getInetAddress()); } - } else { - boolean isAdv = true; + } + + if (!isAdv) { + if (!peer.isNeedSyncFromUs()) { + throw new P2pException(TypeEnum.BAD_MESSAGE, "no need sync"); + } for (Sha256Hash hash : fetchInvDataMsg.getHashList()) { - if (peer.getAdvInvSpread().getIfPresent(new Item(hash, InventoryType.BLOCK)) == null) { - isAdv = false; - break; + long blockNum = new BlockId(hash).getNum(); + long minBlockNum = + peer.getLastSyncBlockId().getNum() - 2 * NetConstants.SYNC_FETCH_BATCH_NUM; + if (blockNum < minBlockNum) { + throw new P2pException(TypeEnum.BAD_MESSAGE, + "minBlockNum: " + minBlockNum + ", blockNum: " + blockNum); } - } - if (!isAdv) { - if (!peer.isNeedSyncFromUs()) { - throw new P2pException(TypeEnum.BAD_MESSAGE, "no need sync"); + if (blockNum > peer.getLastSyncBlockId().getNum()) { + throw new P2pException(TypeEnum.BAD_MESSAGE, + "maxBlockNum: " + peer.getLastSyncBlockId().getNum() + ", blockNum: " + blockNum); } - for (Sha256Hash hash : fetchInvDataMsg.getHashList()) { - long blockNum = new BlockId(hash).getNum(); - long minBlockNum = - peer.getLastSyncBlockId().getNum() - 2 * NetConstants.SYNC_FETCH_BATCH_NUM; - if (blockNum < minBlockNum) { - 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"); - } - peer.getSyncBlockIdCache().put(hash, System.currentTimeMillis()); + if (peer.getSyncBlockIdCache().getIfPresent(hash) != null) { + throw new P2pException(TypeEnum.BAD_MESSAGE, + new BlockId(hash).getString() + " is exist"); } + peer.getSyncBlockIdCache().put(hash, System.currentTimeMillis()); } } } - } 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 5fd6d6725ba..368532c7211 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 @@ -61,6 +61,31 @@ public void testProcessMessage() throws Exception { Assert.assertNotNull(syncBlockIdCache.getIfPresent(blockId)); } + @Test + public void testIsAdvInv() { + FetchInvDataMsgHandler fetchInvDataMsgHandler = new FetchInvDataMsgHandler(); + + List list = new LinkedList<>(); + list.add(Sha256Hash.ZERO_HASH); + FetchInvDataMessage msg = + new FetchInvDataMessage(list, Protocol.Inventory.InventoryType.TRX); + + boolean isAdv = fetchInvDataMsgHandler.isAdvInv(null, msg); + Assert.assertTrue(isAdv); + + PeerConnection peer = Mockito.mock(PeerConnection.class); + Cache advInvSpread = CacheBuilder.newBuilder().build(); + Mockito.when(peer.getAdvInvSpread()).thenReturn(advInvSpread); + + msg = new FetchInvDataMessage(list, Protocol.Inventory.InventoryType.BLOCK); + isAdv = fetchInvDataMsgHandler.isAdvInv(peer, msg); + Assert.assertTrue(!isAdv); + + advInvSpread.put(new Item(Sha256Hash.ZERO_HASH, Protocol.Inventory.InventoryType.BLOCK), 1L); + isAdv = fetchInvDataMsgHandler.isAdvInv(peer, msg); + Assert.assertTrue(isAdv); + } + @Test public void testSyncFetchCheck() { BlockCapsule.BlockId blockId = new BlockCapsule.BlockId(Sha256Hash.ZERO_HASH, 10000L);