feat(net): optimize sync block chain message check logic

Author: xxo1shine

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,15 +56,17 @@ 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<BlockId> blockIds = msg.getBlockIds();
     if (CollectionUtils.isEmpty(blockIds)) {
       throw new P2pException(TypeEnum.BAD_MESSAGE, "SyncBlockChain blockIds is empty");
     }
 
     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<BlockId> getLostBlockIds(List<BlockId> blockIds) throws P2pException {

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

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<BlockCapsule.BlockId> 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));
   }
 
 }