From fe6a2571801656ff1599ef87bdee20f519a5d1fe Mon Sep 17 00:00:00 2001
From: Nikita Koksharov <nkoksharov@redisson.pro>
Date: Thu, 1 Jun 2023 10:08:05 +0300
Subject: [PATCH] Feature - allowedClasses setting added to SerializationCodec
 https://github.com/redisson/redisson/security/code-scanning/4

---
 .../codec/CustomObjectInputStream.java        | 20 +++++++++-----
 .../redisson/codec/SerializationCodec.java    | 27 ++++++++++++-------
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/redisson/src/main/java/org/redisson/codec/CustomObjectInputStream.java b/redisson/src/main/java/org/redisson/codec/CustomObjectInputStream.java
index 827627968..14e093adf 100644
--- a/redisson/src/main/java/org/redisson/codec/CustomObjectInputStream.java
+++ b/redisson/src/main/java/org/redisson/codec/CustomObjectInputStream.java
@@ -15,13 +15,11 @@
  */
 package org.redisson.codec;
 
-import java.io.IOException;
-import java.io.InputStream;
-import java.io.ObjectInputStream;
-import java.io.ObjectStreamClass;
+import java.io.*;
 import java.lang.reflect.Proxy;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 
 /**
  * 
@@ -31,7 +29,14 @@ import java.util.List;
 public class CustomObjectInputStream extends ObjectInputStream {
 
     private final ClassLoader classLoader;
-    
+    private Set<String> allowedClasses;
+
+    public CustomObjectInputStream(ClassLoader classLoader, InputStream in,Set<String> allowedClasses) throws IOException {
+        super(in);
+        this.classLoader = classLoader;
+        this.allowedClasses = allowedClasses;
+    }
+
     public CustomObjectInputStream(ClassLoader classLoader, InputStream in) throws IOException {
         super(in);
         this.classLoader = classLoader;
@@ -41,6 +46,9 @@ public class CustomObjectInputStream extends ObjectInputStream {
     protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
         try {
             String name = desc.getName();
+            if (allowedClasses != null && !allowedClasses.contains(name)) {
+                throw new InvalidClassException("Class " + name + " isn't allowed");
+            }
             return Class.forName(name, false, classLoader);
         } catch (ClassNotFoundException e) {
             return super.resolveClass(desc);
@@ -56,7 +64,7 @@ public class CustomObjectInputStream extends ObjectInputStream {
             loadedClasses.add(clazz);
         }
         
-        return Proxy.getProxyClass(classLoader, loadedClasses.toArray(new Class[loadedClasses.size()]));
+        return Proxy.getProxyClass(classLoader, loadedClasses.toArray(new Class[0]));
     }
     
 }
diff --git a/redisson/src/main/java/org/redisson/codec/SerializationCodec.java b/redisson/src/main/java/org/redisson/codec/SerializationCodec.java
index 62374a4bf..35b1da431 100644
--- a/redisson/src/main/java/org/redisson/codec/SerializationCodec.java
+++ b/redisson/src/main/java/org/redisson/codec/SerializationCodec.java
@@ -15,19 +15,19 @@
  */
 package org.redisson.codec;
 
-import java.io.IOException;
-import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
-
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.ByteBufAllocator;
+import io.netty.buffer.ByteBufInputStream;
+import io.netty.buffer.ByteBufOutputStream;
 import org.redisson.client.codec.BaseCodec;
 import org.redisson.client.handler.State;
 import org.redisson.client.protocol.Decoder;
 import org.redisson.client.protocol.Encoder;
 
-import io.netty.buffer.ByteBuf;
-import io.netty.buffer.ByteBufAllocator;
-import io.netty.buffer.ByteBufInputStream;
-import io.netty.buffer.ByteBufOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.util.Set;
 
 /**
  * JDK's serialization codec.
@@ -51,7 +51,7 @@ public class SerializationCodec extends BaseCodec {
                     ObjectInputStream inputStream;
                     if (classLoader != null) {
                         Thread.currentThread().setContextClassLoader(classLoader);
-                        inputStream = new CustomObjectInputStream(classLoader, in);
+                        inputStream = new CustomObjectInputStream(classLoader, in, allowedClasses);
                     } else {
                         inputStream = new ObjectInputStream(in);
                     }
@@ -84,7 +84,8 @@ public class SerializationCodec extends BaseCodec {
             }
         }
     };
-    
+
+    private Set<String> allowedClasses;
     private final ClassLoader classLoader;
 
     public SerializationCodec() {
@@ -97,6 +98,12 @@ public class SerializationCodec extends BaseCodec {
 
     public SerializationCodec(ClassLoader classLoader, SerializationCodec codec) {
         this.classLoader = classLoader;
+        this.allowedClasses = codec.allowedClasses;
+    }
+
+    public SerializationCodec(ClassLoader classLoader, Set<String> allowedClasses) {
+        this.classLoader = classLoader;
+        this.allowedClasses = allowedClasses;
     }
     
     @Override