feat(net):test for P2P message rate limit

Author: jwrct

framework/src/main/java/org/tron/core/net/messagehandler/FetchInvDataMsgHandler.java

@@ -41,6 +41,7 @@ public class FetchInvDataMsgHandler implements TronMsgHandler {
       .maximumSize(1000).expireAfterWrite(1, TimeUnit.HOURS).build();
 
   private static final int MAX_SIZE = 1_000_000;
+  private static final int MAX_FETCH_SIZE = 100;
   @Autowired
   private TronNetDelegate tronNetDelegate;
   @Autowired
@@ -163,7 +164,7 @@ private void check(PeerConnection peer, FetchInvDataMessage fetchInvDataMsg) thr
         if (!peer.isNeedSyncFromUs()) {
           throw new P2pException(TypeEnum.BAD_MESSAGE, "no need sync");
         }
-        if (fetchInvDataMsg.getHashList().size() > 100) {
+        if (fetchInvDataMsg.getHashList().size() > MAX_FETCH_SIZE) {
           throw new P2pException(TypeEnum.BAD_MESSAGE, "fetch too more blocks, size:"
               + fetchInvDataMsg.getHashList().size());
         }

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

@@ -0,0 +1,21 @@
+package org.tron.core.net;
+
+import static org.tron.core.net.message.MessageTypes.FETCH_INV_DATA;
+import static org.tron.core.net.message.MessageTypes.SYNC_BLOCK_CHAIN;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class P2pRateLimiterTest {
+  @Test
+  public void test() {
+    P2pRateLimiter limiter = new P2pRateLimiter();
+    limiter.register(SYNC_BLOCK_CHAIN.asByte(), 2);
+    limiter.acquire(SYNC_BLOCK_CHAIN.asByte());
+    limiter.tryAcquire(SYNC_BLOCK_CHAIN.asByte());
+    boolean ret = limiter.tryAcquire(SYNC_BLOCK_CHAIN.asByte());
+    Assert.assertFalse(ret);
+    ret = limiter.tryAcquire(FETCH_INV_DATA.asByte());
+    Assert.assertTrue(ret);
+  }
+}

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

@@ -13,6 +13,7 @@
 import org.tron.common.utils.Sha256Hash;
 import org.tron.core.capsule.BlockCapsule;
 import org.tron.core.config.Parameter;
+import org.tron.core.net.P2pRateLimiter;
 import org.tron.core.net.TronNetDelegate;
 import org.tron.core.net.message.adv.BlockMessage;
 import org.tron.core.net.message.adv.FetchInvDataMessage;
@@ -21,6 +22,8 @@
 import org.tron.core.net.service.adv.AdvService;
 import org.tron.protos.Protocol;
 
+import static org.tron.core.net.message.MessageTypes.FETCH_INV_DATA;
+
 public class FetchInvDataMsgHandlerTest {
 
   @Test
@@ -93,4 +96,35 @@ public void testSyncFetchCheck() {
       Assert.assertEquals(e.getMessage(), "minBlockNum: 16000, blockNum: 10000");
     }
   }
+
+  @Test
+  public void testRateLimiter() {
+    BlockCapsule.BlockId blockId = new BlockCapsule.BlockId(Sha256Hash.ZERO_HASH, 10000L);
+    List<Sha256Hash> blockIds = new LinkedList<>();
+    for (int i = 0; i <= 100; i++) {
+      blockIds.add(blockId);
+    }
+    FetchInvDataMessage msg =
+        new FetchInvDataMessage(blockIds, Protocol.Inventory.InventoryType.BLOCK);
+    PeerConnection peer = Mockito.mock(PeerConnection.class);
+    Mockito.when(peer.isNeedSyncFromUs()).thenReturn(true);
+    Cache<Item, Long> advInvSpread = CacheBuilder.newBuilder().maximumSize(100)
+        .expireAfterWrite(1, TimeUnit.HOURS).recordStats().build();
+    Mockito.when(peer.getAdvInvSpread()).thenReturn(advInvSpread);
+    P2pRateLimiter p2pRateLimiter = new P2pRateLimiter();
+    p2pRateLimiter.register(FETCH_INV_DATA.asByte(), 1);
+    Mockito.when(peer.getP2pRateLimiter()).thenReturn(p2pRateLimiter);
+    FetchInvDataMsgHandler fetchInvDataMsgHandler = new FetchInvDataMsgHandler();
+
+    try {
+      fetchInvDataMsgHandler.processMessage(peer, msg);
+    } catch (Exception e) {
+      Assert.assertEquals("fetch too more blocks, size:101", e.getMessage());
+    }
+    try {
+      fetchInvDataMsgHandler.processMessage(peer, msg);
+    } catch (Exception e) {
+      Assert.assertTrue(e.getMessage().endsWith("rate limit"));
+    }
+  }
 }