Validate ObjectInputStream (#235)

GTNewHorizons · Jul 25, 2023 · 685934f · 685934f
1 parent 27c8086
commit 685934f
Showing 1 changed file with 28 additions and 1 deletion.
diff --git a/src/main/java/com/github/technus/tectech/mechanics/spark/RendererMessage.java b/src/main/java/com/github/technus/tectech/mechanics/spark/RendererMessage.java
@@ -8,6 +8,7 @@
 import java.io.InputStream;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import java.io.ObjectStreamClass;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Random;
@@ -16,6 +17,11 @@
 import net.minecraft.entity.player.EntityPlayer;
 import net.minecraft.world.World;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.apache.logging.log4j.Marker;
+import org.apache.logging.log4j.MarkerManager;
+
 import cpw.mods.fml.common.FMLCommonHandler;
 import cpw.mods.fml.common.network.simpleimpl.IMessage;
 import cpw.mods.fml.common.network.simpleimpl.MessageContext;
@@ -41,7 +47,7 @@ public void fromBytes(ByteBuf pBuffer) {
             byte[] boop = pBuffer.array();
             boop = Arrays.copyOfRange(boop, 1, boop.length);
             InputStream is = new ByteArrayInputStream(boop);
-            ObjectInputStream ois = new ObjectInputStream(is);
+            ObjectInputStream ois = new ValidatingObjectInputStream(is);
             Object data = ois.readObject();
             sparkList = (HashSet<ThaumSpark>) data;
         } catch (IOException | ClassNotFoundException ignored) {}
@@ -117,4 +123,25 @@ private static void thaumLightning(int tX, int tY, int tZ, int tXN, int tYN, int
             }
         }
     }
+
+    private static class ValidatingObjectInputStream extends ObjectInputStream {
+
+        private static final Logger logger = LogManager.getLogger();
+        private static final Marker securityMarker = MarkerManager.getMarker("SuspiciousPackets");
+
+        private ValidatingObjectInputStream(InputStream in) throws IOException {
+            super(in);
+        }
+
+        @Override
+        protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
+            String name = desc.getName();
+            if (!name.equals("java.util.HashSet")
+                    && !name.equals("com.github.technus.tectech.mechanics.spark.ThaumSpark")) {
+                logger.warn(securityMarker, "Received packet containing disallowed class: " + name);
+                throw new RuntimeException();
+            }
+            return super.resolveClass(desc);
+        }
+    }
 }