fix(net): reject unknown inventory type in P2P inv/fetch-inv handling

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;
   }

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);

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()) {

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();

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<Sha256Hash> blockIds = new LinkedList<>();

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);