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 37d46f6a8f3..9027034ccc7 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 @@ -29,7 +29,10 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep SyncBlockChainMessage syncBlockChainMessage = (SyncBlockChainMessage) msg; - check(peer, syncBlockChainMessage); + if (!check(peer, syncBlockChainMessage)) { + peer.disconnect(Protocol.ReasonCode.BAD_PROTOCOL); + return; + } long remainNum = 0; @@ -53,7 +56,7 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep peer.sendMessage(new ChainInventoryMessage(blockIds, remainNum)); } - private void check(PeerConnection peer, SyncBlockChainMessage msg) throws P2pException { + private boolean check(PeerConnection peer, SyncBlockChainMessage msg) throws P2pException { List blockIds = msg.getBlockIds(); if (CollectionUtils.isEmpty(blockIds)) { throw new P2pException(TypeEnum.BAD_MESSAGE, "SyncBlockChain blockIds is empty"); @@ -61,7 +64,9 @@ private void check(PeerConnection peer, SyncBlockChainMessage msg) throws P2pExc BlockId firstId = blockIds.get(0); if (!tronNetDelegate.containBlockInMainChain(firstId)) { - throw new P2pException(TypeEnum.BAD_MESSAGE, "No first block:" + firstId.getString()); + logger.warn("Sync message from peer {} without the first block: {}", + peer.getInetSocketAddress(), firstId.getString()); + return false; } long headNum = tronNetDelegate.getHeadBlockId().getNum(); @@ -76,6 +81,8 @@ private void check(PeerConnection peer, SyncBlockChainMessage msg) throws P2pExc throw new P2pException(TypeEnum.BAD_MESSAGE, "lastSyncNum:" + lastSyncBlockId.getNum() + " gt lastNum:" + lastNum); } + + return true; } private LinkedList getLostBlockIds(List blockIds) throws P2pException { diff --git a/framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java b/framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java index ff6203ee870..d4d53da1748 100644 --- a/framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java +++ b/framework/src/test/java/org/tron/core/net/messagehandler/BlockMsgHandlerTest.java @@ -2,6 +2,7 @@ import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; +import java.io.File; import java.lang.reflect.Field; import java.net.InetSocketAddress; import java.util.List; @@ -10,6 +11,7 @@ import org.junit.Before; import org.junit.Test; import org.tron.common.application.TronApplicationContext; +import org.tron.common.utils.FileUtil; import org.tron.common.utils.Sha256Hash; import org.tron.core.Constant; import org.tron.core.capsule.BlockCapsule; @@ -26,23 +28,24 @@ public class BlockMsgHandlerTest { - protected TronApplicationContext context; + private TronApplicationContext context; private BlockMsgHandler handler; private PeerConnection peer; + private String dbPath = "output-block-message-handler-test"; /** * init context. */ @Before public void init() throws Exception { - Args.setParam(new String[]{"--output-directory", "output-directory", "--debug"}, + Args.setParam(new String[]{"--output-directory", dbPath, "--debug"}, Constant.TEST_CONF); context = new TronApplicationContext(DefaultConfig.class); handler = context.getBean(BlockMsgHandler.class); peer = context.getBean(PeerConnection.class); Channel c1 = new Channel(); InetSocketAddress a1 = new InetSocketAddress("100.1.1.1", 100); - Field field = c1.getClass().getDeclaredField("inetAddress"); + Field field = c1.getClass().getDeclaredField("inetAddress"); field.setAccessible(true); field.set(c1, a1.getAddress()); peer.setChannel(c1); @@ -120,5 +123,6 @@ public void testProcessMessage() { public void destroy() { Args.clearParam(); context.destroy(); + FileUtil.deleteDir(new File(dbPath)); } } 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 a0c37d246da..2dbad09c655 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,24 +1,75 @@ package org.tron.core.net.messagehandler; +import java.io.File; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.net.InetSocketAddress; import java.util.ArrayList; +import java.util.List; +import org.junit.After; import org.junit.Assert; +import org.junit.Before; import org.junit.Test; +import org.tron.common.application.TronApplicationContext; +import org.tron.common.utils.FileUtil; +import org.tron.core.Constant; +import org.tron.core.capsule.BlockCapsule; +import org.tron.core.config.DefaultConfig; +import org.tron.core.config.args.Args; import org.tron.core.exception.P2pException; import org.tron.core.net.message.sync.SyncBlockChainMessage; import org.tron.core.net.peer.PeerConnection; +import org.tron.p2p.connection.Channel; public class SyncBlockChainMsgHandlerTest { - private SyncBlockChainMsgHandler handler = new SyncBlockChainMsgHandler(); - private PeerConnection peer = new PeerConnection(); + private TronApplicationContext context; + private SyncBlockChainMsgHandler handler; + private PeerConnection peer; + private String dbPath = "output-sync-chain-test"; + + @Before + public void init() throws Exception { + Args.setParam(new String[]{"--output-directory", dbPath, "--debug"}, + Constant.TEST_CONF); + context = new TronApplicationContext(DefaultConfig.class); + handler = context.getBean(SyncBlockChainMsgHandler.class); + peer = context.getBean(PeerConnection.class); + Channel c1 = new Channel(); + InetSocketAddress a1 = new InetSocketAddress("100.1.1.1", 100); + Field field = c1.getClass().getDeclaredField("inetSocketAddress"); + field.setAccessible(true); + field.set(c1, a1); + + field = c1.getClass().getDeclaredField("inetAddress"); + field.setAccessible(true); + field.set(c1, a1.getAddress()); + + peer.setChannel(c1); + } @Test - public void testProcessMessage() { + public void testProcessMessage() throws Exception { try { handler.processMessage(peer, new SyncBlockChainMessage(new ArrayList<>())); } catch (P2pException e) { Assert.assertTrue(e.getMessage().equals("SyncBlockChain blockIds is empty")); } + + List blockIds = new ArrayList<>(); + blockIds.add(new BlockCapsule.BlockId()); + SyncBlockChainMessage message = new SyncBlockChainMessage(blockIds); + Method method = handler.getClass().getDeclaredMethod( + "check", PeerConnection.class, SyncBlockChainMessage.class); + method.setAccessible(true); + boolean f = (boolean)method.invoke(handler, peer, message); + Assert.assertTrue(!f); + } + + @After + public void destroy() { + Args.clearParam(); + FileUtil.deleteDir(new File(dbPath)); } }