From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Spottedleaf <spottedleaf@spottedleaf.dev>
Date: Fri, 30 Oct 2020 22:37:16 -0700
Subject: [PATCH] Add packet limiter config

Example config:
packet-limiter:
  kick-message: '&cSent too many packets'
  limits:
    all:
      interval: 7.0
      max-packet-rate: 500.0
    PacketPlayInAutoRecipe:
      interval: 4.0
      max-packet-rate: 5.0
      action: DROP

all section refers to all incoming packets, the action for all is
hard coded to KICK.

For specific limits, the section name is the class's name,
and an action can be defined: DROP or KICK

If interval or rate are less-than 0, the limit is ignored

diff --git a/src/main/java/com/tuinity/tuinity/config/TuinityConfig.java b/src/main/java/com/tuinity/tuinity/config/TuinityConfig.java
index 6473b951788a3ecbeee4811ea0288ff47b48355f..373b36fa36f98f64bb9ffa0dfaecefa3626729f4 100644
--- a/src/main/java/com/tuinity/tuinity/config/TuinityConfig.java
+++ b/src/main/java/com/tuinity/tuinity/config/TuinityConfig.java
@@ -1,6 +1,7 @@
 package com.tuinity.tuinity.config;
 
 import com.destroystokyo.paper.util.SneakyThrow;
+import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.level.TicketType;
 import org.bukkit.Bukkit;
 import org.bukkit.configuration.ConfigurationSection;
@@ -139,6 +140,98 @@ public final class TuinityConfig {
         lagCompensateBlockBreaking = TuinityConfig.getBoolean("lag-compensate-block-breaking", true);
     }
 
+    public static final class PacketLimit {
+        public final double packetLimitInterval;
+        public final double maxPacketRate;
+        public final ViolateAction violateAction;
+
+        public PacketLimit(final double packetLimitInterval, final double maxPacketRate, final ViolateAction violateAction) {
+            this.packetLimitInterval = packetLimitInterval;
+            this.maxPacketRate = maxPacketRate;
+            this.violateAction = violateAction;
+        }
+
+        public static enum ViolateAction {
+            KICK, DROP;
+        }
+    }
+
+    public static String kickMessage;
+    public static PacketLimit allPacketsLimit;
+    public static java.util.Map<Class<? extends net.minecraft.network.protocol.Packet<?>>, PacketLimit> packetSpecificLimits = new java.util.HashMap<>();
+
+    private static void packetLimiter() {
+        packetSpecificLimits.clear();
+        kickMessage = org.bukkit.ChatColor.translateAlternateColorCodes('&', TuinityConfig.getString("packet-limiter.kick-message", "&cSent too many packets"));
+        allPacketsLimit = new PacketLimit(
+                TuinityConfig.getDouble("packet-limiter.limits.all.interval", 7.0),
+                TuinityConfig.getDouble("packet-limiter.limits.all.max-packet-rate", 500.0),
+                PacketLimit.ViolateAction.KICK
+        );
+        if (allPacketsLimit.maxPacketRate <= 0.0 || allPacketsLimit.packetLimitInterval <= 0.0) {
+            allPacketsLimit = null;
+        }
+        final ConfigurationSection section = TuinityConfig.config.getConfigurationSection("packet-limiter.limits");
+
+        // add default packets
+
+        // auto recipe limiting
+        TuinityConfig.getDouble("packet-limiter.limits." +
+                net.minecraft.network.protocol.game.PacketPlayInAutoRecipe.class.getSimpleName() + ".interval", 4.0);
+        TuinityConfig.getDouble("packet-limiter.limits." +
+                net.minecraft.network.protocol.game.PacketPlayInAutoRecipe.class.getSimpleName() + ".max-packet-rate", 5.0);
+        TuinityConfig.getString("packet-limiter.limits." +
+                net.minecraft.network.protocol.game.PacketPlayInAutoRecipe.class.getSimpleName() + ".action", PacketLimit.ViolateAction.DROP.name());
+
+        final String canonicalName = MinecraftServer.class.getCanonicalName();
+        final String nmsPackage = canonicalName.substring(0, canonicalName.lastIndexOf("."));
+        for (final String packetClassName : section.getKeys(false)) {
+            if (packetClassName.equals("all")) {
+                continue;
+            }
+            Class<?> packetClazz = null;
+
+            try {
+                packetClazz = Class.forName(nmsPackage + "." + packetClassName);
+            } catch (final ClassNotFoundException ex) {
+                for (final String subpackage : java.util.Arrays.asList("game", "handshake", "login", "status")) {
+                    try {
+                        packetClazz = Class.forName("net.minecraft.network.protocol." + subpackage + "." + packetClassName);
+                    } catch (final ClassNotFoundException ignore) {}
+                }
+                if (packetClazz == null) {
+                    MinecraftServer.LOGGER.warn("Packet '" + packetClassName + "' does not exist, cannot limit it! Please update tuinity.yml");
+                    continue;
+                }
+            }
+
+            if (!net.minecraft.network.protocol.Packet.class.isAssignableFrom(packetClazz)) {
+                MinecraftServer.LOGGER.warn("Packet '" + packetClassName + "' does not exist, cannot limit it! Please update tuinity.yml");
+                continue;
+            }
+
+            if (!(section.get(packetClassName.concat(".interval")) instanceof Number) || !(section.get(packetClassName.concat(".max-packet-rate")) instanceof Number)) {
+                throw new RuntimeException("Packet limit setting " + packetClassName + " is missing interval or max-packet-rate!");
+            }
+
+            final String actionString = section.getString(packetClassName.concat(".action"), "KICK");
+            PacketLimit.ViolateAction action = PacketLimit.ViolateAction.KICK;
+            for (PacketLimit.ViolateAction test : PacketLimit.ViolateAction.values()) {
+                if (actionString.equalsIgnoreCase(test.name())) {
+                    action = test;
+                    break;
+                }
+            }
+
+            final double interval = section.getDouble(packetClassName.concat(".interval"));
+            final double rate = section.getDouble(packetClassName.concat(".max-packet-rate"));
+
+            if (interval > 0.0 && rate > 0.0) {
+                packetSpecificLimits.put((Class)packetClazz, new PacketLimit(interval, rate, action));
+            }
+        }
+    }
+
     public static final class WorldConfig {
 
         public final String worldName;
diff --git a/src/main/java/net/minecraft/network/NetworkManager.java b/src/main/java/net/minecraft/network/NetworkManager.java
index a9637772ead360ee476d59104b50a505e6d0ef4c..2abf64fda6e504a89b2e443af48894023d446cef 100644
--- a/src/main/java/net/minecraft/network/NetworkManager.java
+++ b/src/main/java/net/minecraft/network/NetworkManager.java
@@ -196,8 +196,64 @@ public class NetworkManager extends SimpleChannelInboundHandler<Packet<?>> {
         if (net.minecraft.server.MinecraftServer.getServer().isDebugging()) throwable.printStackTrace(); // Spigot
     }
 
+    // Tuinity start - packet limiter
+    protected final Object PACKET_LIMIT_LOCK = new Object();
+    protected final com.tuinity.tuinity.util.IntervalledCounter allPacketCounts = com.tuinity.tuinity.config.TuinityConfig.allPacketsLimit != null ? new com.tuinity.tuinity.util.IntervalledCounter(
+            (long)(com.tuinity.tuinity.config.TuinityConfig.allPacketsLimit.packetLimitInterval * 1.0e9)
+    ) : null;
+    protected final java.util.Map<Class<? extends net.minecraft.network.protocol.Packet<?>>, com.tuinity.tuinity.util.IntervalledCounter> packetSpecificLimits = new java.util.HashMap<>();
+
+    private boolean stopReadingPackets;
+    private void killForPacketSpam() {
+        this.sendPacket(new PacketPlayOutKickDisconnect(org.bukkit.craftbukkit.util.CraftChatMessage.fromString(com.tuinity.tuinity.config.TuinityConfig.kickMessage, true)[0]), (future) -> {
+            this.close(org.bukkit.craftbukkit.util.CraftChatMessage.fromString(com.tuinity.tuinity.config.TuinityConfig.kickMessage, true)[0]);
+        });
+        this.stopReading();
+        this.stopReadingPackets = true;
+    }
+    // Tuinity end - packet limiter
+
     protected void channelRead0(ChannelHandlerContext channelhandlercontext, Packet<?> packet) throws Exception {
         if (this.channel.isOpen()) {
+            // Tuinity start - packet limiter
+            if (this.stopReadingPackets) {
+                return;
+            }
+            if (this.allPacketCounts != null ||
+                    com.tuinity.tuinity.config.TuinityConfig.packetSpecificLimits.containsKey(packet.getClass())) {
+                long time = System.nanoTime();
+                synchronized (PACKET_LIMIT_LOCK) {
+                    if (this.allPacketCounts != null) {
+                        this.allPacketCounts.updateAndAdd(1, time);
+                        if (this.allPacketCounts.getRate() >= com.tuinity.tuinity.config.TuinityConfig.allPacketsLimit.maxPacketRate) {
+                            this.killForPacketSpam();
+                            return;
+                        }
+                    }
+
+                    for (Class<?> check = packet.getClass(); check != Object.class; check = check.getSuperclass()) {
+                        com.tuinity.tuinity.config.TuinityConfig.PacketLimit packetSpecificLimit =
+                                com.tuinity.tuinity.config.TuinityConfig.packetSpecificLimits.get(check);
+                        if (packetSpecificLimit == null) {
+                            continue;
+                        }
+                        com.tuinity.tuinity.util.IntervalledCounter counter = this.packetSpecificLimits.computeIfAbsent((Class)check, (clazz) -> {
+                            return new com.tuinity.tuinity.util.IntervalledCounter((long)(packetSpecificLimit.packetLimitInterval * 1.0e9));
+                        });
+                        counter.updateAndAdd(1, time);
+                        if (counter.getRate() >= packetSpecificLimit.maxPacketRate) {
+                            switch (packetSpecificLimit.violateAction) {
+                                case DROP:
+                                    return;
+                                case KICK:
+                                    this.killForPacketSpam();
+                                    return;
+                            }
+                        }
+                    }
+                }
+            }
+            // Tuinity end - packet limiter
             try {
                 a(packet, this.packetListener);
             } catch (CancelledPacketHandleException cancelledpackethandleexception) {
