From 9f28eec44733e56ddb537cb3dfe968f6049a72f1 Mon Sep 17 00:00:00 2001
From: UnfamiliarLegacy <74633542+UnfamiliarLegacy@users.noreply.github.com>
Date: Mon, 24 Jun 2024 03:09:04 +0200
Subject: [PATCH] Add HPacket type checks to prevent mistakes

---
 .../java/gearth/extensions/Extension.java     | 13 ++++++--
 .../gearth/protocol/connection/HProxy.java    |  4 +++
 .../packetsafety/PacketTypeChecker.java       | 32 +++++++++++++++++++
 .../packetsafety/PacketTypeException.java     |  7 ++++
 4 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 G-Earth/src/main/java/gearth/protocol/connection/packetsafety/PacketTypeChecker.java
 create mode 100644 G-Earth/src/main/java/gearth/protocol/connection/packetsafety/PacketTypeException.java

diff --git a/G-Earth/src/main/java/gearth/extensions/Extension.java b/G-Earth/src/main/java/gearth/extensions/Extension.java
index 970a20b..82bb553 100644
--- a/G-Earth/src/main/java/gearth/extensions/Extension.java
+++ b/G-Earth/src/main/java/gearth/extensions/Extension.java
@@ -2,7 +2,7 @@ package gearth.extensions;
 
 import gearth.misc.HostInfo;
 import gearth.protocol.HPacketFormat;
-import gearth.services.extension_handler.extensions.implementations.network.NetworkExtensionMessage;
+import gearth.protocol.connection.packetsafety.PacketTypeChecker;
 import gearth.services.extension_handler.extensions.implementations.network.NetworkExtensionMessage.Incoming;
 import gearth.services.extension_handler.extensions.implementations.network.NetworkExtensionMessage.Outgoing;
 import gearth.services.packet_info.PacketInfoManager;
@@ -31,6 +31,11 @@ public abstract class Extension extends ExtensionBase {
 
     private OutputStream out = null;
 
+    /**
+     * Which client G-Earth is connected to.
+     */
+    private HClient clientType = null;
+
     private String getArgument(String[] args, String... arg) {
         for (int i = 0; i < args.length - 1; i++) {
             for (String str : arg) {
@@ -132,7 +137,7 @@ public abstract class Extension extends ExtensionBase {
                     int connectionPort = packet.readInteger();
                     String hotelVersion = packet.readString();
                     String clientIdentifier = packet.readString();
-                    HClient clientType = HClient.valueOf(packet.readString());
+                    clientType = HClient.valueOf(packet.readString());
                     setPacketInfoManager(PacketInfoManager.readFromPacket(packet));
 
                     Constants.UNITY_PACKETS = clientType == HClient.UNITY;
@@ -231,6 +236,10 @@ public abstract class Extension extends ExtensionBase {
         return send(packet, HMessage.Direction.TOSERVER);
     }
     private boolean send(HPacket packet, HMessage.Direction direction) {
+        if (clientType == null) return false;
+
+        PacketTypeChecker.ensureValid(clientType, direction, packet);
+
         if (packet.isCorrupted()) return false;
 
         if (!packet.isPacketComplete()) packet.completePacket(packetInfoManager);
diff --git a/G-Earth/src/main/java/gearth/protocol/connection/HProxy.java b/G-Earth/src/main/java/gearth/protocol/connection/HProxy.java
index 57d6d62..7c62ed1 100644
--- a/G-Earth/src/main/java/gearth/protocol/connection/HProxy.java
+++ b/G-Earth/src/main/java/gearth/protocol/connection/HProxy.java
@@ -1,7 +1,9 @@
 package gearth.protocol.connection;
 
+import gearth.protocol.HMessage;
 import gearth.protocol.HPacket;
 import gearth.protocol.connection.packetsafety.PacketSafetyManager;
+import gearth.protocol.connection.packetsafety.PacketTypeChecker;
 import gearth.protocol.connection.packetsafety.SafePacketsContainer;
 import gearth.services.packet_info.PacketInfo;
 import gearth.services.packet_info.PacketInfoManager;
@@ -57,6 +59,7 @@ public class HProxy {
 
     public boolean sendToServer(HPacket packet) {
         if (outHandler != null) {
+            PacketTypeChecker.ensureValid(hClient, HMessage.Direction.TOSERVER, packet);
             return outHandler.sendToStream(packet.toBytes());
         }
         return false;
@@ -64,6 +67,7 @@ public class HProxy {
 
     public boolean sendToClient(HPacket packet) {
         if (inHandler != null) {
+            PacketTypeChecker.ensureValid(hClient, HMessage.Direction.TOCLIENT, packet);
             return inHandler.sendToStream(packet.toBytes());
         }
         return false;
diff --git a/G-Earth/src/main/java/gearth/protocol/connection/packetsafety/PacketTypeChecker.java b/G-Earth/src/main/java/gearth/protocol/connection/packetsafety/PacketTypeChecker.java
new file mode 100644
index 0000000..b78a980
--- /dev/null
+++ b/G-Earth/src/main/java/gearth/protocol/connection/packetsafety/PacketTypeChecker.java
@@ -0,0 +1,32 @@
+package gearth.protocol.connection.packetsafety;
+
+import gearth.protocol.HMessage;
+import gearth.protocol.HPacket;
+import gearth.protocol.HPacketFormat;
+import gearth.protocol.connection.HClient;
+import gearth.protocol.packethandler.shockwave.packets.ShockPacketIncoming;
+import gearth.protocol.packethandler.shockwave.packets.ShockPacketOutgoing;
+
+public class PacketTypeChecker {
+    /**
+     * Ensure that the packet is of the correct type for the client
+     *
+     * @param clientType The client we are sending to.
+     * @param direction The direction of the packet.
+     * @param packet The packet to check.
+     * @throws PacketTypeException If the packet is not of the correct type.
+     */
+    public static void ensureValid(HClient clientType, HMessage.Direction direction, HPacket packet) {
+        if (clientType == HClient.SHOCKWAVE) {
+            if (direction == HMessage.Direction.TOSERVER && !(packet instanceof ShockPacketOutgoing)) {
+                throw new PacketTypeException(String.format("ShockPacketOutgoing expected for %s", clientType));
+            }
+
+            if (direction == HMessage.Direction.TOCLIENT && !(packet instanceof ShockPacketIncoming)) {
+                throw new PacketTypeException(String.format("ShockPacketIncoming expected for %s", clientType));
+            }
+        } else if (packet.getFormat() != HPacketFormat.EVA_WIRE) {
+            throw new PacketTypeException(String.format("Invalid packet, expected HPacket for %s", clientType));
+        }
+    }
+}
diff --git a/G-Earth/src/main/java/gearth/protocol/connection/packetsafety/PacketTypeException.java b/G-Earth/src/main/java/gearth/protocol/connection/packetsafety/PacketTypeException.java
new file mode 100644
index 0000000..edf0646
--- /dev/null
+++ b/G-Earth/src/main/java/gearth/protocol/connection/packetsafety/PacketTypeException.java
@@ -0,0 +1,7 @@
+package gearth.protocol.connection.packetsafety;
+
+public class PacketTypeException extends RuntimeException {
+    public PacketTypeException(String message) {
+        super(message);
+    }
+}