Filter classes accepted by TransformUpgrader deserialization (apache#40438)

Author: chamikaramj

.github/trigger_files/beam_PostCommit_TransformService_Direct.json

@@ -1,5 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "revision": 5,
+  "revision": 2,
   "https://github.com/apache/beam/pull/39990": "removing dead code from FnApiDoFnRunner"
 }

sdks/java/core/src/main/java/org/apache/beam/sdk/util/construction/TransformUpgrader.java

@@ -24,6 +24,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InvalidClassException;
+import java.io.ObjectInputFilter;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.net.ServerSocket;
@@ -69,6 +70,29 @@ public class TransformUpgrader implements AutoCloseable {
   private static final Logger LOG = LoggerFactory.getLogger(TransformUpgrader.class);
   private static final String UPGRADE_NAMESPACE = "transform:upgrade:";
 
+  public static final String DESERIALIZATION_FILTER_PROPERTY =
+      "beam.transformUpgrader.deserializationFilter";
+
+  private static final String DEFAULT_DESERIALIZATION_FILTER_PATTERN =
+      "maxdepth=100;maxrefs=10000;maxbytes=10485760;maxarray=10000;"
+          + "org.apache.beam.**;"
+          + "com.google.**;"
+          + "org.apache.kafka.**;"
+          + "org.apache.avro.**;"
+          + "org.joda.time.**;"
+          + "!java.lang.ProcessBuilder;"
+          + "!java.lang.Runtime;"
+          + "!java.lang.Thread;"
+          + "java.lang.**;"
+          + "java.util.**;"
+          + "java.io.**;"
+          + "java.math.**;"
+          + "java.time.**;"
+          + "!*";
+
+  private static final ObjectInputFilter DEFAULT_DESERIALIZATION_FILTER =
+      ObjectInputFilter.Config.createFilter(DEFAULT_DESERIALIZATION_FILTER_PATTERN);
+
   @VisibleForTesting static final String UPGRADE_KEY = "upgraded_to_version";
 
   private ExpansionServiceClientFactory clientFactory;
@@ -463,12 +487,21 @@ public static byte[] toByteArray(Object object) {
    * A utility method that converts a byte array obtained by invoking {@link #toByteArray(Object)}
    * back to a Java object.
    *
+   * <p>Applies a deserialization filter with resource limits and a package allowlist. The filter
+   * pattern can be overridden via the {@link #DESERIALIZATION_FILTER_PROPERTY} system property.
+   *
    * @param bytes an array of bytes generated by invoking the {@link #toByteArray(Object)} method.
    * @return re-generated object.
    */
   public static Object fromByteArray(byte[] bytes) throws InvalidClassException {
     try (ByteArrayInputStream bis = new ByteArrayInputStream(bytes);
         ObjectInputStream in = new ObjectInputStream(bis)) {
+      String customPattern = System.getProperty(DESERIALIZATION_FILTER_PROPERTY);
+      ObjectInputFilter filter =
+          (customPattern != null && !customPattern.isEmpty())
+              ? ObjectInputFilter.Config.createFilter(customPattern)
+              : DEFAULT_DESERIALIZATION_FILTER;
+      in.setObjectInputFilter(filter);
       return in.readObject();
     } catch (InvalidClassException e) {
       LOG.info(

sdks/java/core/src/test/java/org/apache/beam/sdk/util/construction/TransformUpgraderTest.java

@@ -18,13 +18,16 @@
 package org.apache.beam.sdk.util.construction;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import com.google.auto.service.AutoService;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
+import java.io.InvalidClassException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import java.net.URI;
 import java.nio.charset.Charset;
 import java.util.Collections;
 import java.util.List;
@@ -487,4 +490,40 @@ public void testVersionComparison() throws Exception {
     assertTrue(TransformUpgrader.compareVersions("2.55.0-SNAPSHOT", "2.53.0") > 0);
     assertTrue(TransformUpgrader.compareVersions("2.55.0.dev", "2.53.0") > 0);
   }
+
+  @Test
+  public void testFromByteArrayAllowedAndBlockedClasses() throws Exception {
+    byte[] stringBytes = TransformUpgrader.toByteArray("hello");
+    assertEquals("hello", TransformUpgrader.fromByteArray(stringBytes));
+
+    byte[] beamObjectBytes = TransformUpgrader.toByteArray(new TestTransform(5));
+    TestTransform deserializedTransform =
+        (TestTransform) TransformUpgrader.fromByteArray(beamObjectBytes);
+    assertEquals(Integer.valueOf(5), deserializedTransform.getTestParam());
+
+    byte[] disallowedPackageBytes = TransformUpgrader.toByteArray(URI.create("http://example.com"));
+    assertThrows(
+        InvalidClassException.class, () -> TransformUpgrader.fromByteArray(disallowedPackageBytes));
+
+    byte[] disallowedClassBytes = TransformUpgrader.toByteArray(ProcessBuilder.class);
+    assertThrows(
+        InvalidClassException.class, () -> TransformUpgrader.fromByteArray(disallowedClassBytes));
+  }
+
+  @Test
+  public void testFromByteArrayCustomFilterProperty() throws Exception {
+    URI uri = URI.create("http://example.com");
+    byte[] uriBytes = TransformUpgrader.toByteArray(uri);
+    byte[] largeArrayBytes = TransformUpgrader.toByteArray(new byte[100]);
+
+    try {
+      System.setProperty(
+          TransformUpgrader.DESERIALIZATION_FILTER_PROPERTY, "maxarray=10;java.net.**;!*");
+      assertEquals(uri, TransformUpgrader.fromByteArray(uriBytes));
+      assertThrows(
+          InvalidClassException.class, () -> TransformUpgrader.fromByteArray(largeArrayBytes));
+    } finally {
+      System.clearProperty(TransformUpgrader.DESERIALIZATION_FILTER_PROPERTY);
+    }
+  }
 }