diff --git a/common/src/main/java/org/tron/core/config/Parameter.java b/common/src/main/java/org/tron/core/config/Parameter.java index 3b3d4e2ffd7..b98de0499fa 100644 --- a/common/src/main/java/org/tron/core/config/Parameter.java +++ b/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; 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 9dd950ae57b..67751230367 100644 --- a/framework/src/main/java/org/tron/core/net/P2pEventHandlerImpl.java +++ b/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 { 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 d95c5db87a7..0c4ba981494 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 @@ -128,7 +128,7 @@ public class PeerConnection { @Setter @Getter - private HelloMessage helloMessageReceive; + private volatile HelloMessage helloMessageReceive; @Setter @Getter diff --git a/framework/src/main/java/org/tron/core/net/peer/PeerStatusCheck.java b/framework/src/main/java/org/tron/core/net/peer/PeerStatusCheck.java index 04eac202484..3a770a5d2e6 100644 --- a/framework/src/main/java/org/tron/core/net/peer/PeerStatusCheck.java +++ b/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); diff --git a/framework/src/test/java/org/tron/core/net/Base58MessageLoggingTest.java b/framework/src/test/java/org/tron/core/net/Base58MessageLoggingTest.java index 8bc34312e26..dfb48660f7c 100644 --- a/framework/src/test/java/org/tron/core/net/Base58MessageLoggingTest.java +++ b/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(); diff --git a/framework/src/test/java/org/tron/core/net/P2pHelloGateTest.java b/framework/src/test/java/org/tron/core/net/P2pHelloGateTest.java new file mode 100644 index 00000000000..02c5e7747d2 --- /dev/null +++ b/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; + + @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 tronFactory = + Mockito.mockStatic(TronMessageFactory.class); + MockedStatic 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()); + } +} diff --git a/framework/src/test/java/org/tron/core/net/messagehandler/MessageHandlerTest.java b/framework/src/test/java/org/tron/core/net/messagehandler/MessageHandlerTest.java index c8205b6b721..615fb573962 100644 --- a/framework/src/test/java/org/tron/core/net/messagehandler/MessageHandlerTest.java +++ b/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()); diff --git a/framework/src/test/java/org/tron/core/net/peer/PeerStatusCheckMockTest.java b/framework/src/test/java/org/tron/core/net/peer/PeerStatusCheckMockTest.java index d2ee4be5b87..f201a087f3f 100644 --- a/framework/src/test/java/org/tron/core/net/peer/PeerStatusCheckMockTest.java +++ b/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(); + } + } } diff --git a/framework/src/test/java/org/tron/core/net/peer/PeerStatusCheckTest.java b/framework/src/test/java/org/tron/core/net/peer/PeerStatusCheckTest.java index 2d734f45215..7a4fb2fd403 100644 --- a/framework/src/test/java/org/tron/core/net/peer/PeerStatusCheckTest.java +++ b/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()