PacketLimiter: Make queues for each packet type and add global limit

Author: IceTank

src/main/kotlin/com/lambda/module/modules/network/PacketLimiter.kt

@@ -33,44 +33,53 @@ import net.minecraft.network.packet.c2s.play.TeleportConfirmC2SPacket
 
 // ToDo: HUD info
 object PacketLimiter : Module(
-    name = "PacketLimiter",
-    description = "Limits the amount of packets sent to the server",
-    tag = ModuleTag.NETWORK,
+	name = "PacketLimiter",
+	description = "Limits the amount of packets sent to the server",
+	tag = ModuleTag.NETWORK,
 ) {
-    private var packetQueue = LimitedDecayQueue<PacketEvent.Send.Pre>(99, 1000)
-    private val limit by setting("Limit", 99, 1..100, 1, "The maximum amount of packets to send per given time interval", unit = " packets")
-        .onValueChange { _, to -> packetQueue.setSizeLimit(to) }
+	private var packetQueueMap = mutableMapOf<String, LimitedDecayQueue<PacketEvent.Send.Pre>>()
+	private val globalQueue = LimitedDecayQueue<PacketEvent.Send.Pre>(1, 1)
 
-    private val interval by setting("Duration", 4000L, 1L..10000L, 50L, "The interval / duration in milliseconds to limit packets for", unit = " ms")
-        .onValueChange { _, to -> packetQueue.setDecayTime(to) }
+	private val limit by setting("Limit per packet", 99, 1..100, 1, "The maximum amount of packets to send per given time interval", unit = " packets")
+		.onValueChange { _, to -> packetQueueMap.forEach { (t, u) -> u.setSizeLimit(to) } }
+	private val globalLimit by setting("Global Limit", 1000, 100..5000, 50, "The maximum amount of packets to send overall per given time interval", unit = " packets")
+		.onValueChange { _, to -> globalQueue.setSizeLimit(to) }
 
-    private val defaultIgnorePackets = setOf(
-        nameOf<CommonPongC2SPacket>(),
-        nameOf<PositionAndOnGround>(),
-        nameOf<Full>(),
-        nameOf<LookAndOnGround>(),
-        nameOf<OnGroundOnly>(),
-        nameOf<TeleportConfirmC2SPacket>()
-    )
+	private val interval by setting("Duration", 4000L, 1L..10000L, 50L, "The interval / duration in milliseconds to limit packets for", unit = " ms")
+		.onValueChange { _, to -> packetQueueMap.forEach { (t, u) -> u.setDecayTime(to) }; globalQueue.setDecayTime(to) }
 
-    inline fun <reified T> nameOf(): String = T::class.className
+	private val defaultIgnorePackets = setOf(
+		nameOf<CommonPongC2SPacket>(),
+		nameOf<PositionAndOnGround>(),
+		nameOf<Full>(),
+		nameOf<LookAndOnGround>(),
+		nameOf<OnGroundOnly>(),
+		nameOf<TeleportConfirmC2SPacket>()
+	)
 
-    // ToDo: Find a way to have a list of serverbound packets
-    private val ignorePackets by setting("Ignore Packets", defaultIgnorePackets, description = "Packets to ignore when limiting")
+	inline fun <reified T> nameOf(): String = T::class.className
 
-    init {
-        onEnable {
-            packetQueue = LimitedDecayQueue(limit, interval)
-        }
+	// ToDo: Find a way to have a list of serverbound packets
+	private val ignorePackets by setting("Ignore Packets", defaultIgnorePackets, description = "Packets to ignore when limiting")
 
-        listen<PacketEvent.Send.Pre>(Int.MAX_VALUE) {
-            if (it.packet::class.java.name in ignorePackets) return@listen
+	init {
+		onEnable {
+			packetQueueMap.clear()
+			globalQueue.setDecayTime(interval)
+			globalQueue.setSizeLimit(globalLimit)
+		}
 
-//            this@PacketLimiter.info("Packet sent: ${it.packet::class.simpleName} (${packetQueue.size} / $limit) ${Instant.now()}")
-            if (packetQueue.add(it)) return@listen
+		listen<PacketEvent.Send.Pre>(Int.MAX_VALUE) {
+			if (it.packet::class.java.name in ignorePackets) return@listen
 
-            it.cancel()
-            this@PacketLimiter.info("Packet limit reached, dropping packet: ${it.packet::class.simpleName} (${packetQueue.size} / $limit)")
-        }
-    }
+			//            this@PacketLimiter.info("Packet sent: ${it.packet::class.simpleName} (${packetQueue.size} / $limit) ${Instant.now()}")
+			val queue = packetQueueMap.getOrPut(it.packet::class.java.name) {
+				LimitedDecayQueue(limit, interval)
+			}
+			if (queue.add(it)) return@listen
+
+			it.cancel()
+			this@PacketLimiter.info("Packet limit reached, dropping packet: ${it.packet::class.simpleName} (${queue.size} / $limit)")
+		}
+	}
 }