diff --git a/framework/src/main/java/org/tron/core/net/messagehandler/ChainInventoryMsgHandler.java b/framework/src/main/java/org/tron/core/net/messagehandler/ChainInventoryMsgHandler.java index bd2e428418c..ecff89fcf67 100644 --- a/framework/src/main/java/org/tron/core/net/messagehandler/ChainInventoryMsgHandler.java +++ b/framework/src/main/java/org/tron/core/net/messagehandler/ChainInventoryMsgHandler.java @@ -108,6 +108,11 @@ private void check(PeerConnection peer, ChainInventoryMessage msg) throws P2pExc throw new P2pException(TypeEnum.BAD_MESSAGE, "blockIds is empty"); } + if (msg.getRemainNum() < 0) { + throw new P2pException(TypeEnum.BAD_MESSAGE, + "remainNum is negative: " + msg.getRemainNum()); + } + if (blockIds.size() > NetConstants.SYNC_FETCH_BATCH_NUM + 1) { throw new P2pException(TypeEnum.BAD_MESSAGE, "big blockIds size: " + blockIds.size()); } @@ -137,9 +142,12 @@ private void check(PeerConnection peer, ChainInventoryMessage msg) throws P2pExc long maxFutureNum = maxRemainTime / BLOCK_PRODUCED_INTERVAL + tronNetDelegate.getSolidBlockId().getNum(); long lastNum = blockIds.get(blockIds.size() - 1).getNum(); - if (lastNum + msg.getRemainNum() > maxFutureNum) { - throw new P2pException(TypeEnum.BAD_MESSAGE, "lastNum: " + lastNum + " + remainNum: " - + msg.getRemainNum() + " > futureMaxNum: " + maxFutureNum); + long declaredHighestNum = lastNum + msg.getRemainNum(); + if (declaredHighestNum < 0 || declaredHighestNum > maxFutureNum) { + throw new P2pException(TypeEnum.BAD_MESSAGE, + "Invalid declared highest block number: " + declaredHighestNum + + ", lastNum: " + lastNum + ", remainNum: " + msg.getRemainNum() + + ", futureMaxNum: " + maxFutureNum); } } } diff --git a/framework/src/main/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandler.java b/framework/src/main/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandler.java index 5c18e014978..dd9d5e47ebe 100644 --- a/framework/src/main/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandler.java +++ b/framework/src/main/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandler.java @@ -58,8 +58,7 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep } private boolean check(PeerConnection peer, SyncBlockChainMessage msg) throws P2pException { - if (peer.getRemainNum() > 0 - && !peer.getP2pRateLimiter().tryAcquire(msg.getType().asByte())) { + if (!peer.getP2pRateLimiter().tryAcquire(msg.getType().asByte())) { // Discard messages that exceed the rate limit logger.warn("{} message from peer {} exceeds the rate limit", msg.getType(), peer.getInetSocketAddress()); diff --git a/framework/src/main/java/org/tron/core/net/peer/PeerConnection.java b/framework/src/main/java/org/tron/core/net/peer/PeerConnection.java index 7d7457cf2fc..5e2115bb90d 100644 --- a/framework/src/main/java/org/tron/core/net/peer/PeerConnection.java +++ b/framework/src/main/java/org/tron/core/net/peer/PeerConnection.java @@ -146,7 +146,7 @@ public class PeerConnection { private volatile long remainNum; @Getter private Cache syncBlockIdCache = CacheBuilder.newBuilder() - .maximumSize(2 * NetConstants.SYNC_FETCH_BATCH_NUM).recordStats().build(); + .maximumSize(2 * NetConstants.SYNC_FETCH_BATCH_NUM + 1).recordStats().build(); @Setter @Getter private Deque syncBlockToFetch = new ConcurrentLinkedDeque<>(); diff --git a/framework/src/test/java/org/tron/core/net/messagehandler/ChainInventoryMsgHandlerTest.java b/framework/src/test/java/org/tron/core/net/messagehandler/ChainInventoryMsgHandlerTest.java index 56853c3dbb7..aab20891f75 100644 --- a/framework/src/test/java/org/tron/core/net/messagehandler/ChainInventoryMsgHandlerTest.java +++ b/framework/src/test/java/org/tron/core/net/messagehandler/ChainInventoryMsgHandlerTest.java @@ -1,5 +1,7 @@ package org.tron.core.net.messagehandler; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; import java.util.ArrayList; import java.util.LinkedList; import java.util.List; @@ -7,13 +9,16 @@ import org.junit.Assert; import org.junit.BeforeClass; import org.junit.Test; +import org.mockito.Mockito; import org.tron.common.TestConstants; import org.tron.common.utils.Pair; +import org.tron.common.utils.ReflectUtils; +import org.tron.common.utils.Sha256Hash; import org.tron.core.capsule.BlockCapsule.BlockId; import org.tron.core.config.Parameter.NetConstants; import org.tron.core.config.args.Args; import org.tron.core.exception.P2pException; -import org.tron.core.net.message.keepalive.PingMessage; +import org.tron.core.net.TronNetDelegate; import org.tron.core.net.message.sync.ChainInventoryMessage; import org.tron.core.net.peer.PeerConnection; @@ -78,4 +83,52 @@ public void testProcessMessage() throws Exception { Assert.assertNull(msg.getAnswerMessage()); } + @Test + public void testNegativeRemainNumRejected() throws Exception { + assertCheckRejects(createContinuousBlockIds(0L), -1L); + } + + @Test + public void testRemainNumOverflowRejected() throws Exception { + assertCheckRejects(createContinuousBlockIds( + Long.MAX_VALUE - NetConstants.SYNC_FETCH_BATCH_NUM + 1), 1L); + } + + private void assertCheckRejects(List ids, long remainNum) throws Exception { + ChainInventoryMsgHandler messageHandler = new ChainInventoryMsgHandler(); + TronNetDelegate tronNetDelegate = Mockito.mock(TronNetDelegate.class); + ReflectUtils.setFieldValue(messageHandler, "tronNetDelegate", tronNetDelegate); + Mockito.when(tronNetDelegate.getHeadBlockId()).thenReturn( + new BlockId(Sha256Hash.ZERO_HASH, 1L)); + Mockito.when(tronNetDelegate.getSolidBlockId()).thenReturn( + new BlockId(Sha256Hash.ZERO_HASH, 1L)); + Mockito.when(tronNetDelegate.getBlockTime(Mockito.any())).thenReturn(0L); + + PeerConnection connection = Mockito.mock(PeerConnection.class); + LinkedList requestedIds = new LinkedList<>(); + requestedIds.add(ids.get(0)); + Mockito.when(connection.getSyncChainRequested()).thenReturn( + new Pair<>(requestedIds, System.currentTimeMillis())); + + Method check = ChainInventoryMsgHandler.class.getDeclaredMethod( + "check", PeerConnection.class, ChainInventoryMessage.class); + check.setAccessible(true); + try { + check.invoke(messageHandler, connection, new ChainInventoryMessage(ids, remainNum)); + Assert.fail("Expected invalid remainNum to be rejected"); + } catch (InvocationTargetException e) { + Assert.assertTrue(e.getCause() instanceof P2pException); + Assert.assertEquals(P2pException.TypeEnum.BAD_MESSAGE, + ((P2pException) e.getCause()).getType()); + } + } + + private List createContinuousBlockIds(long firstNum) { + List ids = new ArrayList<>(); + for (int i = 0; i < NetConstants.SYNC_FETCH_BATCH_NUM; i++) { + ids.add(new BlockId(Sha256Hash.ZERO_HASH, firstNum + i)); + } + return ids; + } + } 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 7ea6337443d..52e3f1bec8e 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 @@ -68,6 +68,29 @@ public void testProcessMessage() throws Exception { Assert.assertNotNull(syncBlockIdCache.getIfPresent(blockId)); } + @Test + public void testSyncBlockIdCacheRetainsOldestHashInValidWindow() { + PeerConnection peer = new PeerConnection(); + Sha256Hash firstHash = createHash(0); + int windowSize = 2 * (int) Parameter.NetConstants.SYNC_FETCH_BATCH_NUM + 1; + + for (int i = 0; i < windowSize; i++) { + peer.getSyncBlockIdCache().put(createHash(i), (long) i); + } + peer.getSyncBlockIdCache().cleanUp(); + + Assert.assertNotNull(peer.getSyncBlockIdCache().getIfPresent(firstHash)); + } + + private Sha256Hash createHash(int value) { + byte[] bytes = new byte[Sha256Hash.LENGTH]; + bytes[28] = (byte) (value >>> 24); + bytes[29] = (byte) (value >>> 16); + bytes[30] = (byte) (value >>> 8); + bytes[31] = (byte) value; + return Sha256Hash.wrap(bytes); + } + @Test public void testIsAdvInv() { FetchInvDataMsgHandler fetchInvDataMsgHandler = new FetchInvDataMsgHandler(); diff --git a/framework/src/test/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandlerTest.java b/framework/src/test/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandlerTest.java index 08c5484880f..f8edf5f7bc8 100644 --- a/framework/src/test/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandlerTest.java +++ b/framework/src/test/java/org/tron/core/net/messagehandler/SyncBlockChainMsgHandlerTest.java @@ -1,5 +1,7 @@ package org.tron.core.net.messagehandler; +import static org.tron.core.net.message.MessageTypes.SYNC_BLOCK_CHAIN; + import java.io.IOException; import java.lang.reflect.Field; import java.lang.reflect.InvocationTargetException; @@ -14,6 +16,7 @@ import org.junit.ClassRule; import org.junit.Test; import org.junit.rules.TemporaryFolder; +import org.mockito.Mockito; import org.tron.common.TestConstants; import org.tron.common.application.TronApplicationContext; import org.tron.common.utils.Sha256Hash; @@ -22,6 +25,7 @@ import org.tron.core.config.DefaultConfig; import org.tron.core.config.args.Args; import org.tron.core.exception.P2pException; +import org.tron.core.net.P2pRateLimiter; import org.tron.core.net.TronNetDelegate; import org.tron.core.net.message.sync.BlockInventoryMessage; import org.tron.core.net.message.sync.SyncBlockChainMessage; @@ -160,6 +164,24 @@ public void testBlockIdsAtLimit() throws Exception { } } + @Test + public void testRemainNumZeroStillConsumesSyncBlockChainRateLimit() throws Exception { + PeerConnection rateLimitedPeer = Mockito.mock(PeerConnection.class); + P2pRateLimiter rateLimiter = new P2pRateLimiter(); + rateLimiter.register(SYNC_BLOCK_CHAIN.asByte(), 0.0001D); + Mockito.when(rateLimitedPeer.getP2pRateLimiter()).thenReturn(rateLimiter); + + BlockId genesis = context.getBean(TronNetDelegate.class).getGenesisBlockId(); + SyncBlockChainMessage message = new SyncBlockChainMessage( + java.util.Collections.singletonList(genesis)); + Method checkMethod = SyncBlockChainMsgHandler.class + .getDeclaredMethod("check", PeerConnection.class, SyncBlockChainMessage.class); + checkMethod.setAccessible(true); + + Assert.assertTrue((boolean) checkMethod.invoke(handler, rateLimitedPeer, message)); + Assert.assertFalse((boolean) checkMethod.invoke(handler, rateLimitedPeer, message)); + } + @AfterClass public static void destroy() { for (PeerConnection p : PeerManager.getPeers()) {