feat(net):P2P message rate limit

Author: jwrct

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

@@ -178,9 +178,11 @@ private void processMessage(PeerConnection peer, byte[] data) {
           handshakeService.processHelloMessage(peer, (HelloMessage) msg);
           break;
         case P2P_DISCONNECT:
-          peer.getChannel().close();
-          peer.getNodeStatistics()
-                  .nodeDisconnectedRemote(((DisconnectMessage)msg).getReason());
+          if (peer.getP2pRateLimiter().tryAcquire(type.asByte())) {
+            peer.getChannel().close();
+            peer.getNodeStatistics()
+                .nodeDisconnectedRemote(((DisconnectMessage)msg).getReason());
+          }
           break;
         case SYNC_BLOCK_CHAIN:
           syncBlockChainMsgHandler.processMessage(peer, msg);

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

@@ -0,0 +1,30 @@
+package org.tron.core.net;
+
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.util.concurrent.RateLimiter;
+
+public class P2pRateLimiter {
+  private final Cache<Byte, RateLimiter> rateLimiters = CacheBuilder.newBuilder()
+      .maximumSize(256).build();
+
+  public void register (Byte type, double rate) {
+    rateLimiters.put(type, RateLimiter.create(rate));
+  }
+
+  public void acquire (Byte type) {
+    RateLimiter rateLimiter = rateLimiters.getIfPresent(type);
+    if (rateLimiter == null) {
+      return;
+    }
+    rateLimiter.acquire();
+  }
+
+  public boolean tryAcquire (Byte type) {
+    RateLimiter rateLimiter = rateLimiters.getIfPresent(type);
+    if (rateLimiter == null) {
+      return true;
+    }
+    return rateLimiter.tryAcquire();
+  }
+}

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

@@ -55,6 +55,13 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep
 
     FetchInvDataMessage fetchInvDataMsg = (FetchInvDataMessage) msg;
 
+    if (peer.isNeedSyncFromUs() && !peer.getP2pRateLimiter().tryAcquire(msg.getType().asByte())) {
+      // Discard messages that exceed the rate limit
+      logger.warn("{} message from peer {} exceeds the rate limit",
+          msg.getType(), peer.getInetSocketAddress());
+      return;
+    }
+
     check(peer, fetchInvDataMsg);
 
     InventoryType type = fetchInvDataMsg.getInventoryType();
@@ -156,6 +163,10 @@ 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) {
+          throw new P2pException(TypeEnum.BAD_MESSAGE, "fetch too more blocks, size:"
+              + fetchInvDataMsg.getHashList().size());
+        }
         for (Sha256Hash hash : fetchInvDataMsg.getHashList()) {
           long blockNum = new BlockId(hash).getNum();
           long minBlockNum =

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

@@ -31,6 +31,13 @@ public void processMessage(PeerConnection peer, TronMessage msg) throws P2pExcep
 
     SyncBlockChainMessage syncBlockChainMessage = (SyncBlockChainMessage) msg;
 
+    if (!peer.getP2pRateLimiter().tryAcquire(msg.getType().asByte())) {
+      // Discard messages that exceed the rate limit
+      logger.warn("{} message from peer {} exceeds the rate limit",
+          msg.getType(), peer.getInetSocketAddress());
+      return;
+    }
+
     if (!check(peer, syncBlockChainMessage)) {
       peer.disconnect(Protocol.ReasonCode.BAD_PROTOCOL);
       return;

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

@@ -1,5 +1,9 @@
 package org.tron.core.net.peer;
 
+import static org.tron.core.net.message.MessageTypes.FETCH_INV_DATA;
+import static org.tron.core.net.message.MessageTypes.P2P_DISCONNECT;
+import static org.tron.core.net.message.MessageTypes.SYNC_BLOCK_CHAIN;
+
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
 import com.google.protobuf.ByteString;
@@ -32,7 +36,9 @@
 import org.tron.core.config.args.Args;
 import org.tron.core.metrics.MetricsKey;
 import org.tron.core.metrics.MetricsUtil;
+import org.tron.core.net.P2pRateLimiter;
 import org.tron.core.net.TronNetDelegate;
+import org.tron.core.net.message.MessageTypes;
 import org.tron.core.net.message.adv.InventoryMessage;
 import org.tron.core.net.message.adv.TransactionsMessage;
 import org.tron.core.net.message.base.DisconnectMessage;
@@ -156,6 +162,8 @@ public class PeerConnection {
   @Setter
   @Getter
   private volatile boolean needSyncFromUs = true;
+  @Getter
+  private P2pRateLimiter p2pRateLimiter = new P2pRateLimiter();
 
   public void setChannel(Channel channel) {
     this.channel = channel;
@@ -164,6 +172,9 @@ public void setChannel(Channel channel) {
     }
     this.nodeStatistics = TronStatsManager.getNodeStatistics(channel.getInetAddress());
     lastInteractiveTime = System.currentTimeMillis();
+    p2pRateLimiter.register(SYNC_BLOCK_CHAIN.asByte(), 2);
+    p2pRateLimiter.register(FETCH_INV_DATA.asByte(), 1);
+    p2pRateLimiter.register(P2P_DISCONNECT.asByte(), 1);
   }
 
   public void setBlockBothHave(BlockId blockId) {