fix(net): gate peer traffic on hello validation and enforce timeout

common/src/main/java/org/tron/core/config/Parameter.java

@@ -100,6 +100,7 @@ public class NodeConstant {
   public class NetConstants {
 
     public static final long SYNC_FETCH_BATCH_NUM = 2000;
+    public static final long HELLO_TIME_OUT = 10000L;
     public static final long ADV_TIME_OUT = 20000L;
     public static final long SYNC_TIME_OUT = 5000L;
     public static final long NET_MAX_TRX_PER_SECOND = 700L;

framework/src/main/java/org/tron/core/net/P2pEventHandlerImpl.java

@@ -127,6 +127,13 @@ public void onMessage(Channel c, byte[] data) {
       return;
     }
 
+    if (peerConnection.getHelloMessageReceive() == null
+        && data[0] != MessageTypes.P2P_HELLO.asByte()
+        && data[0] != MessageTypes.P2P_DISCONNECT.asByte()) {
+      peerConnection.disconnect(Protocol.ReasonCode.BAD_PROTOCOL);
+      return;
+    }
+
     if (MessageTypes.PBFT_MSG.asByte() == data[0]) {
       PbftMessage message = null;
       try {

framework/src/main/java/org/tron/core/net/peer/PeerConnection.java

@@ -128,7 +128,7 @@ public class PeerConnection {
 
   @Setter
   @Getter
-  private HelloMessage helloMessageReceive;
+  private volatile HelloMessage helloMessageReceive;
 
   @Setter
   @Getter

framework/src/main/java/org/tron/core/net/peer/PeerStatusCheck.java

@@ -56,6 +56,14 @@ public void statusCheck() {
         isDisconnected = true;
       }
 
+      if (!isDisconnected) {
+        isDisconnected = peer.getHelloMessageReceive() == null
+            && peer.getChannel().getStartTime() <= now - NetConstants.HELLO_TIME_OUT;
+        if (isDisconnected) {
+          logger.warn("Peer {} hello message timeout", peer.getInetAddress());
+        }
+      }
+
       if (!isDisconnected) {
         isDisconnected = peer.getAdvInvRequest().values().stream()
             .anyMatch(time -> time < now - NetConstants.ADV_TIME_OUT);

framework/src/test/java/org/tron/core/net/Base58MessageLoggingTest.java

@@ -139,6 +139,8 @@ public void pbftSignatureFailureLogIsSafeAndDisabledPbftStillReturnsEarly() thro
         .addSrAddress(bytes(16384)).build().toByteString());
     for (boolean enabled : new boolean[]{false, true}) {
       PeerConnection peer = peer();
+      // PBFT messages reach this handler only after the HELLO handshake.
+      when(peer.getHelloMessageReceive()).thenReturn(mock(HelloMessage.class));
       Channel channel = peer.getChannel();
       P2pEventHandlerImpl handler = new P2pEventHandlerImpl();
       PbftMsgHandler pbftHandler = new PbftMsgHandler();

framework/src/test/java/org/tron/core/net/P2pHelloGateTest.java

@@ -0,0 +1,104 @@
+package org.tron.core.net;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+import org.tron.common.utils.ReflectUtils;
+import org.tron.core.ChainBaseManager;
+import org.tron.core.ChainBaseManager.NodeType;
+import org.tron.core.capsule.BlockCapsule.BlockId;
+import org.tron.core.net.message.MessageTypes;
+import org.tron.core.net.message.PbftMessageFactory;
+import org.tron.core.net.message.TronMessageFactory;
+import org.tron.core.net.message.base.DisconnectMessage;
+import org.tron.core.net.message.handshake.HelloMessage;
+import org.tron.core.net.peer.PeerConnection;
+import org.tron.core.net.peer.PeerManager;
+import org.tron.core.net.service.handshake.HandshakeService;
+import org.tron.core.net.service.statistics.NodeStatistics;
+import org.tron.core.net.service.statistics.PeerStatistics;
+import org.tron.p2p.connection.Channel;
+import org.tron.p2p.discover.Node;
+import org.tron.protos.Protocol.ReasonCode;
+
+public class P2pHelloGateTest {
+
+  private final P2pEventHandlerImpl handler = new P2pEventHandlerImpl();
+  private final Channel channel = mock(Channel.class);
+  private final PeerConnection peer = mock(PeerConnection.class);
+  private final HandshakeService handshakeService = mock(HandshakeService.class);
+  private MockedStatic<PeerManager> peerManager;
+
+  @Before
+  public void setUp() {
+    ReflectUtils.setFieldValue(handler, "handshakeService", handshakeService);
+    peerManager = Mockito.mockStatic(PeerManager.class);
+    peerManager.when(() -> PeerManager.getPeerConnection(channel)).thenReturn(peer);
+  }
+
+  @After
+  public void tearDown() {
+    peerManager.close();
+  }
+
+  @Test
+  public void testRejectMessagesBeforeHelloWithoutParsing() {
+    try (MockedStatic<TronMessageFactory> tronFactory =
+             Mockito.mockStatic(TronMessageFactory.class);
+         MockedStatic<PbftMessageFactory> pbftFactory =
+             Mockito.mockStatic(PbftMessageFactory.class)) {
+      int rejected = 0;
+      for (MessageTypes type : MessageTypes.values()) {
+        if (type != MessageTypes.P2P_HELLO && type != MessageTypes.P2P_DISCONNECT) {
+          handler.onMessage(channel, new byte[] {type.asByte()});
+          rejected++;
+        }
+      }
+      verify(peer, times(rejected)).disconnect(ReasonCode.BAD_PROTOCOL);
+      tronFactory.verifyNoInteractions();
+      pbftFactory.verifyNoInteractions();
+    }
+  }
+
+  @Test
+  public void testHelloIsDispatchedBeforeHandshake() throws Exception {
+    ChainBaseManager chain = mock(ChainBaseManager.class);
+    when(chain.getGenesisBlockId()).thenReturn(new BlockId());
+    when(chain.getSolidBlockId()).thenReturn(new BlockId());
+    when(chain.getHeadBlockId()).thenReturn(new BlockId());
+    when(chain.getNodeType()).thenReturn(NodeType.FULL);
+    HelloMessage hello = new HelloMessage(
+        new Node(new byte[64], "127.0.0.1", null, 18888), 0, chain);
+    when(peer.getPeerStatistics()).thenReturn(new PeerStatistics());
+
+    handler.onMessage(channel, hello.getSendBytes());
+
+    verify(handshakeService).processHelloMessage(eq(peer), any(HelloMessage.class));
+    verify(peer, never()).disconnect(any());
+  }
+
+  @Test
+  public void testDisconnectIsDispatchedBeforeHandshake() {
+    NodeStatistics stats = mock(NodeStatistics.class);
+    when(peer.getPeerStatistics()).thenReturn(new PeerStatistics());
+    when(peer.getP2pRateLimiter()).thenReturn(new P2pRateLimiter());
+    when(peer.getChannel()).thenReturn(channel);
+    when(peer.getNodeStatistics()).thenReturn(stats);
+
+    handler.onMessage(channel, new DisconnectMessage(ReasonCode.PEER_QUITING).getSendBytes());
+
+    verify(channel).close();
+    verify(stats).nodeDisconnectedRemote(ReasonCode.PEER_QUITING);
+    verify(peer, never()).disconnect(any());
+  }
+}

framework/src/test/java/org/tron/core/net/messagehandler/MessageHandlerTest.java

@@ -24,6 +24,7 @@
 import org.tron.core.config.args.Args;
 import org.tron.core.net.P2pEventHandlerImpl;
 import org.tron.core.net.TronNetService;
+import org.tron.core.net.message.handshake.HelloMessage;
 import org.tron.core.net.message.keepalive.PingMessage;
 import org.tron.core.net.peer.PeerConnection;
 import org.tron.core.net.peer.PeerManager;
@@ -82,6 +83,7 @@ public void testPbft() {
     Assert.assertFalse(c1.isDisconnect());
 
     peer = PeerManager.getPeers().get(0);
+    peer.setHelloMessageReceive(mock(HelloMessage.class));
     BlockCapsule blockCapsule = new BlockCapsule(1, Sha256Hash.ZERO_HASH,
         System.currentTimeMillis(), ByteString.EMPTY);
     PbftMessage pbftMessage = PbftMessage.fullNodePrePrepareBlockMsg(blockCapsule, 0L);
@@ -102,7 +104,7 @@ public void testPing() {
     Channel c1 = mock(Channel.class);
     Mockito.when(c1.getInetSocketAddress()).thenReturn(a1);
     Mockito.when(c1.getInetAddress()).thenReturn(a1.getAddress());
-    PeerManager.add(ctx, c1);
+    PeerManager.add(ctx, c1).setHelloMessageReceive(mock(HelloMessage.class));
 
     PingMessage pingMessage = new PingMessage();
     p2pEventHandler.onMessage(c1, pingMessage.getSendBytes());

framework/src/test/java/org/tron/core/net/peer/PeerStatusCheckMockTest.java

@@ -2,16 +2,26 @@
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import java.util.Collections;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import org.junit.After;
 import org.junit.Test;
 import org.mockito.Mockito;
 import org.tron.common.utils.ReflectUtils;
+import org.tron.core.config.Parameter.NetConstants;
+import org.tron.core.net.TronNetDelegate;
+import org.tron.core.net.message.handshake.HelloMessage;
+import org.tron.p2p.connection.Channel;
+import org.tron.protos.Protocol.ReasonCode;
 
 public class PeerStatusCheckMockTest {
   @After
@@ -41,4 +51,33 @@ public void testInitException() {
     Mockito.verify(peerStatusCheck).statusCheck();
   }
 
+  @Test
+  public void testHelloTimeout() {
+    PeerStatusCheck check = new PeerStatusCheck();
+    PeerConnection peer = mock(PeerConnection.class);
+    Channel channel = mock(Channel.class);
+    TronNetDelegate delegate = mock(TronNetDelegate.class);
+    when(delegate.getActivePeer()).thenReturn(Collections.singletonList(peer));
+    when(peer.getChannel()).thenReturn(channel);
+    when(peer.getAdvInvRequest()).thenReturn(Collections.emptyMap());
+    when(peer.getSyncBlockRequested()).thenReturn(Collections.emptyMap());
+    ReflectUtils.setFieldValue(check, "tronNetDelegate", delegate);
+    try {
+      long now = System.currentTimeMillis();
+      when(channel.getStartTime()).thenReturn(now);
+      check.statusCheck();
+      verify(peer, never()).disconnect(any());
+
+      when(channel.getStartTime()).thenReturn(now - NetConstants.HELLO_TIME_OUT - 1);
+      check.statusCheck();
+      verify(peer).disconnect(ReasonCode.TIME_OUT);
+
+      clearInvocations(peer);
+      when(peer.getHelloMessageReceive()).thenReturn(mock(HelloMessage.class));
+      check.statusCheck();
+      verify(peer, never()).disconnect(any());
+    } finally {
+      check.close();
+    }
+  }
 }

framework/src/test/java/org/tron/core/net/peer/PeerStatusCheckTest.java

@@ -13,6 +13,7 @@
 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.net.message.handshake.HelloMessage;
 import org.tron.p2p.connection.Channel;
 
 
@@ -44,7 +45,7 @@ public void testCheck() {
       ReflectUtils.setFieldValue(c1, "ctx", spy(ChannelHandlerContext.class));
       Mockito.doNothing().when(c1).send((byte[]) any());
 
-      PeerManager.add(context, c1);
+      PeerManager.add(context, c1).setHelloMessageReceive(Mockito.mock(HelloMessage.class));
     }
 
     PeerManager.getPeers().get(0).getSyncBlockRequested()