Fix other ois implementations

core/src/main/java/io/dogboy/serializationisbad/core/ClassFilteringObjectInputStream.java

@@ -24,7 +24,7 @@ public ClassFilteringObjectInputStream(InputStream in, PatchModule patchModule)
         this(in, patchModule, null);
     }
 
-    private boolean isClassAllowed(String className) {
+    private static boolean isClassAllowed(String className, PatchModule patchModule) {
         // strip all array dimensions, just get the base type
         while (className.startsWith("[")) {
             className = className.substring(1);
@@ -35,12 +35,12 @@ private boolean isClassAllowed(String className) {
         }
 
         if (SerializationIsBad.getInstance().getConfig().getClassAllowlist().contains(className)
-                || this.patchModule.getClassAllowlist().contains(className)) {
+                || patchModule.getClassAllowlist().contains(className)) {
             return true;
         }
 
         Set<String> allowedPackages = new HashSet<>(SerializationIsBad.getInstance().getConfig().getPackageAllowlist());
-        allowedPackages.addAll(this.patchModule.getPackageAllowlist());
+        allowedPackages.addAll(patchModule.getPackageAllowlist());
 
         for (String allowedPackage : allowedPackages) {
             if (className.startsWith(allowedPackage + ".")) {
@@ -53,13 +53,7 @@ private boolean isClassAllowed(String className) {
 
     @Override
     protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
-        SerializationIsBad.logger.debug("Resolving class " + desc.getName());
-
-        if (!this.isClassAllowed(desc.getName())) {
-            SerializationIsBad.logger.warn("Tried to resolve class " + desc.getName() + ", which is not allowed to be deserialized");
-            if (SerializationIsBad.getInstance().getConfig().isExecuteBlocking())
-                throw new ClassNotFoundException("Class " + desc.getName() + " is not allowed to be deserialized");
-        }
+        ClassFilteringObjectInputStream.resolveClassPrecheck(desc, this.patchModule);
 
         if (this.parentClassLoader == null) {
             return super.resolveClass(desc);
@@ -78,6 +72,16 @@ protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, Clas
         }
     }
 
+    public static void resolveClassPrecheck(ObjectStreamClass desc, PatchModule patchModule) throws ClassNotFoundException {
+        SerializationIsBad.logger.debug("Resolving class " + desc.getName());
+
+        if (!ClassFilteringObjectInputStream.isClassAllowed(desc.getName(), patchModule)) {
+            SerializationIsBad.logger.warn("Tried to resolve class " + desc.getName() + ", which is not allowed to be deserialized");
+            if (SerializationIsBad.getInstance().getConfig().isExecuteBlocking())
+                throw new ClassNotFoundException("Class " + desc.getName() + " is not allowed to be deserialized");
+        }
+    }
+
     private static final HashMap<String, Class<?>> primClasses = new HashMap<>(8, 1.0F);
     static {
         ClassFilteringObjectInputStream.primClasses.put("boolean", boolean.class);

core/src/main/java/io/dogboy/serializationisbad/core/Patches.java

@@ -20,7 +20,8 @@ public class Patches {
 
     public static PatchModule getPatchModuleForClass(String className) {
         for (PatchModule patchModule : SerializationIsBad.getInstance().getConfig().getPatchModules()) {
-            if (patchModule.getClassesToPatch().contains(className)) {
+            if (patchModule.getClassesToPatch().contains(className)
+                    || patchModule.getCustomOISClasses().contains(className)) {
                 return patchModule;
             }
         }
@@ -43,6 +44,33 @@ private static byte[] writeClassNode(ClassNode classNode) {
 
     private static void applyPatches(String className, ClassNode classNode, boolean passClassLoader) {
         SerializationIsBad.logger.info("Applying patches to " + className);
+        PatchModule patchModule = Patches.getPatchModuleForClass(className);
+        if (patchModule == null) {
+            SerializationIsBad.logger.info("  No patches to apply");
+            return;
+        }
+
+        if (patchModule.getCustomOISClasses().contains(className) && "java/io/ObjectInputStream".equals(classNode.superName)) {
+            for (MethodNode methodNode : classNode.methods) {
+                if (!"resolveClass".equals(methodNode.name)) continue;
+
+                InsnList additionalInstructions = new InsnList();
+                additionalInstructions.add(new VarInsnNode(Opcodes.ALOAD, 1)); // Class Descriptor
+                additionalInstructions.add(new LdcInsnNode(className));
+                additionalInstructions.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "io/dogboy/serializationisbad/core/Patches",
+                        "getPatchModuleForClass", "(Ljava/lang/String;)Lio/dogboy/serializationisbad/core/config/PatchModule;", false));
+                additionalInstructions.add(new MethodInsnNode(Opcodes.INVOKESTATIC, "io/dogboy/serializationisbad/core/ClassFilteringObjectInputStream",
+                        "resolveClassPrecheck", "(Ljava/io/ObjectStreamClass;Lio/dogboy/serializationisbad/core/config/PatchModule;)V", false));
+
+                methodNode.instructions.insertBefore(methodNode.instructions.getFirst(), additionalInstructions);
+
+                SerializationIsBad.logger.info("  Injecting resolveClass precheck in method " + methodNode.name);
+
+                break;
+            }
+
+            return;
+        }
 
         for (MethodNode methodNode : classNode.methods) {
             InsnList instructions = methodNode.instructions;

core/src/main/java/io/dogboy/serializationisbad/core/config/PatchModule.java

@@ -5,11 +5,13 @@
 
 public class PatchModule {
     private Set<String> classesToPatch;
+    private Set<String> customOISClasses;
     private Set<String> classAllowlist;
     private Set<String> packageAllowlist;
 
     public PatchModule() {
         this.classesToPatch = new HashSet<>();
+        this.customOISClasses = new HashSet<>();
         this.classAllowlist = new HashSet<>();
         this.packageAllowlist = new HashSet<>();
     }
@@ -22,6 +24,15 @@ public void setClassesToPatch(Set<String> classesToPatch) {
         this.classesToPatch = classesToPatch;
     }
 
+
+    public Set<String> getCustomOISClasses() {
+        return this.customOISClasses;
+    }
+
+    public void setCustomOISClasses(Set<String> customOISClasses) {
+        this.customOISClasses = customOISClasses;
+    }
+
     public Set<String> getClassAllowlist() {
         return this.classAllowlist;
     }
@@ -37,4 +48,5 @@ public Set<String> getPackageAllowlist() {
     public void setPackageAllowlist(Set<String> packageAllowlist) {
         this.packageAllowlist = packageAllowlist;
     }
+
 }

modlauncher/src/main/java/io/dogboy/serializationisbad/modlauncher/SIBTransformer.java

@@ -11,6 +11,7 @@
 
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class SIBTransformer implements ITransformer<ClassNode> {
     private final PatchModule patchModule;
@@ -36,7 +37,8 @@ public TransformerVoteResult castVote(ITransformerVotingContext context) {
 
     @Override
     public Set<Target> targets() {
-        return this.patchModule.getClassesToPatch().stream()
+        return Stream.concat(this.patchModule.getClassesToPatch().stream(),
+                        this.patchModule.getCustomOISClasses().stream())
                 .map(Target::targetClass)
                 .collect(Collectors.toSet());
     }