diff --git a/framework/src/main/java/org/tron/core/net/P2pEventHandlerImpl.java b/framework/src/main/java/org/tron/core/net/P2pEventHandlerImpl.java index f703779c61..9dd950ae57 100644 --- a/framework/src/main/java/org/tron/core/net/P2pEventHandlerImpl.java +++ b/framework/src/main/java/org/tron/core/net/P2pEventHandlerImpl.java @@ -215,7 +215,8 @@ private void processMessage(PeerConnection peer, byte[] data) { } } - private boolean checkInvRateLimit(PeerConnection peer, InventoryMessage msg) { + private boolean checkInvRateLimit(PeerConnection peer, InventoryMessage msg) + throws P2pException { InventoryType invType = msg.getInventoryType(); int currentSize = msg.getInventory().getIdsCount(); MessageStatistics stats = peer.getPeerStatistics().messageStatistics; @@ -237,6 +238,9 @@ private boolean checkInvRateLimit(PeerConnection peer, InventoryMessage msg) { peer.getInetAddress(), count, currentSize, maxBlockInvIn10s); return false; } + } else { + throw new P2pException(P2pException.TypeEnum.BAD_MESSAGE, + "unknown inventory type: " + msg.getInventory().getTypeValue()); } return true; } 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 b1f2646808..2d6d47d1dd 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 @@ -160,9 +160,13 @@ private void check(PeerConnection peer, FetchInvDataMessage fetchInvDataMsg, "FetchInvData contains duplicate hashes, size: " + hashList.size()); } - MessageTypes type = fetchInvDataMsg.getInvMessageType(); + InventoryType invType = fetchInvDataMsg.getInventoryType(); + if (invType != InventoryType.TRX && invType != InventoryType.BLOCK) { + throw new P2pException(TypeEnum.BAD_MESSAGE, + "unknown inventory type: " + fetchInvDataMsg.getInventory().getTypeValue()); + } - if (type == MessageTypes.TRX) { + if (invType == InventoryType.TRX) { for (Sha256Hash hash : fetchInvDataMsg.getHashList()) { if (peer.getAdvInvSpread().getIfPresent(new Item(hash, InventoryType.TRX)) == null) { throw new P2pException(TypeEnum.BAD_MESSAGE, "not spread inv: " + hash); diff --git a/framework/src/main/java/org/tron/core/net/messagehandler/InventoryMsgHandler.java b/framework/src/main/java/org/tron/core/net/messagehandler/InventoryMsgHandler.java index 59232a8d25..f96f7f0b0f 100644 --- a/framework/src/main/java/org/tron/core/net/messagehandler/InventoryMsgHandler.java +++ b/framework/src/main/java/org/tron/core/net/messagehandler/InventoryMsgHandler.java @@ -63,6 +63,10 @@ private boolean check(PeerConnection peer, InventoryMessage inventoryMessage) } InventoryType type = inventoryMessage.getInventoryType(); + if (type != InventoryType.TRX && type != InventoryType.BLOCK) { + throw new P2pException(TypeEnum.BAD_MESSAGE, + "unknown inventory type: " + inventoryMessage.getInventory().getTypeValue()); + } int size = hashList.size(); if (peer.isNeedSyncFromPeer() || peer.isNeedSyncFromUs()) { diff --git a/framework/src/test/java/org/tron/core/net/P2pEventHandlerImplTest.java b/framework/src/test/java/org/tron/core/net/P2pEventHandlerImplTest.java index 93b84450f7..52cdfa9c82 100644 --- a/framework/src/test/java/org/tron/core/net/P2pEventHandlerImplTest.java +++ b/framework/src/test/java/org/tron/core/net/P2pEventHandlerImplTest.java @@ -3,6 +3,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import com.google.protobuf.ByteString; import java.lang.reflect.Method; import java.net.InetSocketAddress; import java.util.ArrayList; @@ -212,6 +213,35 @@ public void testCheckInvRateLimitBlockBoundary() throws Exception { peer.getPeerStatistics().messageStatistics.tronInBlockInventoryElement.getCount(10)); } + @Test + public void testCheckInvRateLimitUnknownTypeRejected() throws Exception { + CommonParameter parameter = CommonParameter.getInstance(); + parameter.setMaxTps(10); + parameter.setMaxBlockInvPerSecond(10); + + PeerStatistics peerStatistics = new PeerStatistics(); + PeerConnection peer = mock(PeerConnection.class); + Mockito.when(peer.getPeerStatistics()).thenReturn(peerStatistics); + + P2pEventHandlerImpl handler = new P2pEventHandlerImpl(); + Method method = handler.getClass() + .getDeclaredMethod("processMessage", PeerConnection.class, byte[].class); + method.setAccessible(true); + + // craft an Inventory whose enum type holds an undefined value (2) + Protocol.Inventory inv = Protocol.Inventory.newBuilder() + .setTypeValue(2) + .addIds(ByteString.copyFrom(new byte[32])) + .build(); + InventoryMessage msg = new InventoryMessage(inv.toByteArray()); + Assert.assertEquals(InventoryType.UNRECOGNIZED, msg.getInventoryType()); + + method.invoke(handler, peer, msg.getSendBytes()); + + // checkInvRateLimit throws BAD_MESSAGE -> processException -> disconnect(BAD_PROTOCOL) + verify(peer).disconnect(Protocol.ReasonCode.BAD_PROTOCOL); + } + @Test public void testUpdateLastInteractiveTime() throws Exception { PeerConnection peer = new PeerConnection(); 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 e8ec425781..7ea6337443 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 @@ -156,6 +156,27 @@ public void testDuplicateHashRejected() throws Exception { } } + @Test + public void testUnknownInventoryTypeRejected() throws Exception { + FetchInvDataMsgHandler handler = new FetchInvDataMsgHandler(); + PeerConnection peer = Mockito.mock(PeerConnection.class); + + // craft a FetchInvData whose enum type holds an undefined value (2) + Protocol.Inventory inv = Protocol.Inventory.newBuilder() + .setTypeValue(2) + .addIds(com.google.protobuf.ByteString.copyFrom(new byte[32])) + .build(); + FetchInvDataMessage msg = new FetchInvDataMessage(inv.toByteArray()); + Assert.assertEquals(Protocol.Inventory.InventoryType.UNRECOGNIZED, msg.getInventoryType()); + + try { + handler.processMessage(peer, msg); + Assert.fail("Expected P2pException for unknown inventory type"); + } catch (P2pException e) { + Assert.assertEquals(P2pException.TypeEnum.BAD_MESSAGE, e.getType()); + } + } + @Test public void testRateLimiter() { List blockIds = new LinkedList<>(); diff --git a/framework/src/test/java/org/tron/core/net/messagehandler/InventoryMsgHandlerTest.java b/framework/src/test/java/org/tron/core/net/messagehandler/InventoryMsgHandlerTest.java index 3d24ff2a4b..9e2523dff2 100644 --- a/framework/src/test/java/org/tron/core/net/messagehandler/InventoryMsgHandlerTest.java +++ b/framework/src/test/java/org/tron/core/net/messagehandler/InventoryMsgHandlerTest.java @@ -2,6 +2,7 @@ import static org.mockito.Mockito.mock; +import com.google.protobuf.ByteString; import java.lang.reflect.Field; import java.net.InetAddress; import java.net.InetSocketAddress; @@ -18,6 +19,7 @@ import org.tron.core.net.message.adv.InventoryMessage; import org.tron.core.net.peer.PeerConnection; import org.tron.p2p.connection.Channel; +import org.tron.protos.Protocol.Inventory; import org.tron.protos.Protocol.Inventory.InventoryType; public class InventoryMsgHandlerTest { @@ -74,6 +76,30 @@ public void testDuplicateHashesRejected() throws Exception { } } + @Test + public void testUnknownInventoryTypeRejected() throws Exception { + InventoryMsgHandler handler = new InventoryMsgHandler(); + Args.setParam(new String[] {}, TestConstants.TEST_CONF); + + // craft an Inventory whose enum type holds an undefined value (2) + Inventory inv = Inventory.newBuilder() + .setTypeValue(2) + .addIds(ByteString.copyFrom(new byte[32])) + .build(); + InventoryMessage msg = new InventoryMessage(inv.toByteArray()); + Assert.assertEquals(InventoryType.UNRECOGNIZED, msg.getInventoryType()); + + PeerConnection peer = new PeerConnection(); + peer.setChannel(getChannel("1.0.0.5", 1000)); + + try { + handler.processMessage(peer, msg); + Assert.fail("Expected P2pException for unknown inventory type"); + } catch (P2pException e) { + Assert.assertEquals(P2pException.TypeEnum.BAD_MESSAGE, e.getType()); + } + } + private Channel getChannel(String host, int port) throws Exception { Channel channel = new Channel(); InetSocketAddress inetSocketAddress = new InetSocketAddress(host, port);