From 2dc2cd8df3f2e47026bed52c049f7f4d0f85b8ae Mon Sep 17 00:00:00 2001
From: Melledy <121644117+Melledy@users.noreply.github.com>
Date: Sun, 3 Dec 2023 02:32:28 -0800
Subject: [PATCH] Prevent packet spam

---
 .../server/game/GameServerPacketHandler.java   | 18 +++++++++++-------
 .../emu/lunarcore/server/game/GameSession.java |  6 +++++-
 .../lunarcore/server/packet/PacketHandler.java |  2 --
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/main/java/emu/lunarcore/server/game/GameServerPacketHandler.java b/src/main/java/emu/lunarcore/server/game/GameServerPacketHandler.java
index 5d801e4..27f8f01 100644
--- a/src/main/java/emu/lunarcore/server/game/GameServerPacketHandler.java
+++ b/src/main/java/emu/lunarcore/server/game/GameServerPacketHandler.java
@@ -5,10 +5,7 @@ import java.util.Set;
 import org.reflections.Reflections;
 
 import emu.lunarcore.LunarCore;
-import emu.lunarcore.server.packet.CmdId;
-import emu.lunarcore.server.packet.Opcodes;
-import emu.lunarcore.server.packet.PacketHandler;
-import emu.lunarcore.server.packet.SessionState;
+import emu.lunarcore.server.packet.*;
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 
@@ -51,6 +48,15 @@ public class GameServerPacketHandler {
         PacketHandler handler = this.handlers.get(cmdId);
 
         if (handler != null) {
+            // Check cooldown to prevent packet spam
+            long timestamp = System.currentTimeMillis();
+            if (session.getPacketCooldown().get(cmdId) >= timestamp) {
+                //LunarCore.getLogger().warn("Dropped a packet " + CmdIdUtils.getCmdIdName(cmdId));
+                return;
+            } else {
+                session.getPacketCooldown().put(cmdId, timestamp + 5);
+            }
+
             try {
                 // Make sure session is ready for packets
                 SessionState state = session.getState();
@@ -70,14 +76,12 @@ public class GameServerPacketHandler {
                         return;
                     }
                 }
-
+                
                 // Handle packet
                 handler.handle(session, data);
             } catch (Exception ex) {
                 ex.printStackTrace();
             }
-            
-            return; // Packet successfully handled
         }
 
         // Log unhandled packets
diff --git a/src/main/java/emu/lunarcore/server/game/GameSession.java b/src/main/java/emu/lunarcore/server/game/GameSession.java
index 1eeb945..e8dcf57 100644
--- a/src/main/java/emu/lunarcore/server/game/GameSession.java
+++ b/src/main/java/emu/lunarcore/server/game/GameSession.java
@@ -11,6 +11,8 @@ import emu.lunarcore.server.packet.SessionState;
 import emu.lunarcore.util.Utils;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
+import it.unimi.dsi.fastutil.ints.Int2LongMap;
+import it.unimi.dsi.fastutil.ints.Int2LongOpenHashMap;
 import kcp.highway.Ukcp;
 import lombok.AccessLevel;
 import lombok.Getter;
@@ -18,6 +20,7 @@ import lombok.Getter;
 @Getter
 public class GameSession {
     private final GameServer server;
+    private final Int2LongMap packetCooldown;
     private InetSocketAddress address;
 
     private Account account;
@@ -25,13 +28,14 @@ public class GameSession {
 
     // Network
     @Getter(AccessLevel.PRIVATE) private Ukcp ukcp;
-
+    
     // Flags
     private SessionState state = SessionState.WAITING_FOR_TOKEN;
     private boolean useSecretKey;
 
     private GameSession(GameServer server) {
         this.server = server;
+        this.packetCooldown = new Int2LongOpenHashMap();
     }
 
     public GameSession(GameServer server, Ukcp ukcp) {
diff --git a/src/main/java/emu/lunarcore/server/packet/PacketHandler.java b/src/main/java/emu/lunarcore/server/packet/PacketHandler.java
index c593373..ebe2b4b 100644
--- a/src/main/java/emu/lunarcore/server/packet/PacketHandler.java
+++ b/src/main/java/emu/lunarcore/server/packet/PacketHandler.java
@@ -3,7 +3,5 @@ package emu.lunarcore.server.packet;
 import emu.lunarcore.server.game.GameSession;
 
 public abstract class PacketHandler {
-    protected static final byte[] EMPTY_BYTE_ARRAY = new byte[0];
-
     public abstract void handle(GameSession session, byte[] data) throws Exception;
 }