optimize: optimize raftsnapshot read (#6896)

apache · Sep 30, 2024 · 851e5de · 851e5de
1 parent 0db8c36
commit 851e5de
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 4 deletions.
diff --git a/...r/src/main/java/org/apache/seata/server/cluster/raft/snapshot/RaftSnapshotSerializer.java b/...r/src/main/java/org/apache/seata/server/cluster/raft/snapshot/RaftSnapshotSerializer.java
@@ -21,8 +21,13 @@
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import java.io.ObjectStreamClass;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 
+import org.apache.seata.common.exception.ErrorCode;
+import org.apache.seata.common.exception.SeataRuntimeException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -37,9 +42,19 @@ public class RaftSnapshotSerializer {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(RaftSnapshotSerializer.class);
 
+    private static final List<String> PERMITS = new ArrayList<>();
+    static {
+        PERMITS.add(RaftSnapshot.class.getName());
+        PERMITS.add(RaftSnapshot.SnapshotType.class.getName());
+        PERMITS.add(io.seata.server.cluster.raft.snapshot.RaftSnapshot.class.getName());
+        PERMITS.add(io.seata.server.cluster.raft.snapshot.RaftSnapshot.SnapshotType.class.getName());
+        PERMITS.add(java.lang.Enum.class.getName());
+        PERMITS.add("[B");
+    }
+
     public static byte[] encode(RaftSnapshot raftSnapshot) throws IOException {
         try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
-                ObjectOutputStream oos = new ObjectOutputStream(bos)) {
+             ObjectOutputStream oos = new ObjectOutputStream(bos)) {
             Serializer serializer =
                     EnhancedServiceLoader.load(Serializer.class, SerializerType.getByCode(raftSnapshot.getCodec()).name());
             Optional.ofNullable(raftSnapshot.getBody()).ifPresent(value -> raftSnapshot.setBody(
@@ -51,7 +66,7 @@ public static byte[] encode(RaftSnapshot raftSnapshot) throws IOException {
 
     public static byte[] encode(io.seata.server.cluster.raft.snapshot.RaftSnapshot raftSnapshot) throws IOException {
         try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
-                ObjectOutputStream oos = new ObjectOutputStream(bos)) {
+             ObjectOutputStream oos = new ObjectOutputStream(bos)) {
             Serializer serializer =
                     EnhancedServiceLoader.load(Serializer.class, SerializerType.getByCode(raftSnapshot.getCodec()).name());
             Optional.ofNullable(raftSnapshot.getBody()).ifPresent(value -> raftSnapshot.setBody(
@@ -63,7 +78,16 @@ public static byte[] encode(io.seata.server.cluster.raft.snapshot.RaftSnapshot r
 
     public static RaftSnapshot decode(byte[] raftSnapshotByte) throws IOException {
         try (ByteArrayInputStream bin = new ByteArrayInputStream(raftSnapshotByte);
-                ObjectInputStream ois = new ObjectInputStream(bin)) {
+             ObjectInputStream ois = new ObjectInputStream(bin) {
+                 @Override
+                 protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
+                     if (!PERMITS.contains(desc.getName())) {
+                         throw new SeataRuntimeException(ErrorCode.ERR_DESERIALIZATION_SECURITY,
+                                 "Failed to deserialize object: " + desc.getName() + " is not permitted");
+                     }
+                     return super.resolveClass(desc);
+                 }
+             }) {
             Object object = ois.readObject();
             RaftSnapshot raftSnapshot;
             if (object instanceof io.seata.server.cluster.raft.snapshot.RaftSnapshot) {
@@ -83,8 +107,11 @@ public static RaftSnapshot decode(byte[] raftSnapshotByte) throws IOException {
                     .ifPresent(value -> raftSnapshot.setBody(serializer.deserialize(CompressorFactory
                             .getCompressor(raftSnapshot.getCompressor()).decompress((byte[])raftSnapshot.getBody()))));
             return raftSnapshot;
-        } catch (ClassNotFoundException e) {
+        } catch (Exception e) {
             LOGGER.info("Failed to read raft snapshot: {}", e.getMessage(), e);
+            if (e instanceof SeataRuntimeException) {
+                throw (SeataRuntimeException)e;
+            }
             throw new IOException(e);
         }
     }

diff --git a/server/src/test/java/org/apache/seata/server/raft/RaftSyncMessageTest.java b/server/src/test/java/org/apache/seata/server/raft/RaftSyncMessageTest.java
@@ -116,6 +116,18 @@ public void testMsgSerializeCompatible() throws IOException {
         Assertions.assertEquals(1234, ((RaftBranchSessionSyncMsg) raftSyncMessageByBranch.getBody()).getBranchSession().getBranchId());
     }
 
+    @Test
+    public void testSecuritySnapshotSerialize() throws IOException {
+        TestSecurity testSecurity = new TestSecurity();
+        byte[] bytes;
+        try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
+             ObjectOutputStream oos = new ObjectOutputStream(bos)) {
+            oos.writeObject(testSecurity);
+            bytes =  bos.toByteArray();
+        }
+        Assertions.assertThrows(SeataRuntimeException.class,()->RaftSnapshotSerializer.decode(bytes));
+    }
+
     @Test
     public void testSnapshotSerialize() throws IOException, TransactionException {
         Map<String, GlobalSession> sessionMap = new HashMap<>();