From d069361202c10d6c64f7f3a6a5be9319831a5036 Mon Sep 17 00:00:00 2001
From: onebeastchris <github@onechris.mozmail.com>
Date: Fri, 6 Dec 2024 21:04:34 +0800
Subject: [PATCH] Add packet class to PacketErrorEvent, add missing EntityEvent

---
 .../network/event/session/PacketErrorEvent.java  | 16 ++++++++++++++--
 .../network/tcp/TcpPacketCodec.java              |  4 ++--
 .../protocol/data/game/entity/EntityEvent.java   |  3 ++-
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/protocol/src/main/java/org/geysermc/mcprotocollib/network/event/session/PacketErrorEvent.java b/protocol/src/main/java/org/geysermc/mcprotocollib/network/event/session/PacketErrorEvent.java
index e812ee59..c6f48096 100644
--- a/protocol/src/main/java/org/geysermc/mcprotocollib/network/event/session/PacketErrorEvent.java
+++ b/protocol/src/main/java/org/geysermc/mcprotocollib/network/event/session/PacketErrorEvent.java
@@ -1,13 +1,15 @@
 package org.geysermc.mcprotocollib.network.event.session;
 
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.geysermc.mcprotocollib.network.Session;
-
+import org.geysermc.mcprotocollib.network.packet.Packet;
 /**
  * Called when a session encounters an error while reading or writing packet data.
  */
 public class PacketErrorEvent implements SessionEvent {
     private final Session session;
     private final Throwable cause;
+    private final @Nullable Class<? extends Packet> packetClass;
     private boolean suppress = false;
 
     /**
@@ -16,9 +18,10 @@ public class PacketErrorEvent implements SessionEvent {
      * @param session Session that the error came from.
      * @param cause Cause of the error.
      */
-    public PacketErrorEvent(Session session, Throwable cause) {
+    public PacketErrorEvent(Session session, Throwable cause, @Nullable Class<? extends Packet> packetClass) {
         this.session = session;
         this.cause = cause;
+        this.packetClass = packetClass;
     }
 
     /**
@@ -39,6 +42,15 @@ public class PacketErrorEvent implements SessionEvent {
         return this.cause;
     }
 
+    /**
+     * Gets the packet class where the error occurred in, if it is known.
+     *
+     * @return the packet class related to this packet error
+     */
+    public Class<? extends Packet> getPacketClass() {
+        return this.packetClass;
+    }
+
     /**
      * Gets whether the error should be suppressed. If the error is not suppressed,
      * it will be passed on through internal error handling and disconnect the session.
diff --git a/protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpPacketCodec.java b/protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpPacketCodec.java
index 98a3e2f4..6ad9fd62 100644
--- a/protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpPacketCodec.java
+++ b/protocol/src/main/java/org/geysermc/mcprotocollib/network/tcp/TcpPacketCodec.java
@@ -57,7 +57,7 @@ public class TcpPacketCodec extends MessageToMessageCodec<ByteBuf, Packet> {
         } catch (Throwable t) {
             log.debug(marker, "Error encoding packet", t);
 
-            PacketErrorEvent e = new PacketErrorEvent(this.session, t);
+            PacketErrorEvent e = new PacketErrorEvent(this.session, t, packet.getClass());
             this.session.callEvent(e);
             if (!e.shouldSuppress()) {
                 throw new EncoderException(t);
@@ -104,7 +104,7 @@ public class TcpPacketCodec extends MessageToMessageCodec<ByteBuf, Packet> {
             // Advance buffer to end to make sure remaining data in this packet is skipped.
             buf.readerIndex(buf.readerIndex() + buf.readableBytes());
 
-            PacketErrorEvent e = new PacketErrorEvent(this.session, t);
+            PacketErrorEvent e = new PacketErrorEvent(this.session, t, packet.getClass());
             this.session.callEvent(e);
             if (!e.shouldSuppress()) {
                 throw new DecoderException(t);
diff --git a/protocol/src/main/java/org/geysermc/mcprotocollib/protocol/data/game/entity/EntityEvent.java b/protocol/src/main/java/org/geysermc/mcprotocollib/protocol/data/game/entity/EntityEvent.java
index c2e62a84..21d9488e 100644
--- a/protocol/src/main/java/org/geysermc/mcprotocollib/protocol/data/game/entity/EntityEvent.java
+++ b/protocol/src/main/java/org/geysermc/mcprotocollib/protocol/data/game/entity/EntityEvent.java
@@ -66,7 +66,8 @@ public enum EntityEvent {
     WARDEN_SONIC_BOOM,
     SNIFFER_MAKE_SOUND,
     ARMADILLO_PEEKING,
-    LIVING_EQUIPMENT_BREAK_BODY;
+    LIVING_EQUIPMENT_BREAK_BODY,
+    SHAKE;
 
     private static final EntityEvent[] VALUES = values();