Avoid final field mutation in UnsafeUtils

dd-java-agent/build.gradle

@@ -67,6 +67,7 @@ dependencies {
 def generalShadowJarConfig(ShadowJar shadowJarTask) {
   shadowJarTask.with {
     mergeServiceFiles()
+    zip64 = true
 
     duplicatesStrategy = DuplicatesStrategy.FAIL
 

dd-java-agent/instrumentation/junit/junit-5/junit-5.3/src/main/java/datadog/trace/instrumentation/junit5/execution/TestDescriptorHandle.java

@@ -2,10 +2,14 @@
 
 import datadog.trace.agent.tooling.muzzle.Reference;
 import datadog.trace.util.MethodHandles;
-import datadog.trace.util.UnsafeUtils;
 import java.lang.invoke.MethodHandle;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import org.junit.platform.commons.util.ClassLoaderUtils;
 import org.junit.platform.engine.TestDescriptor;
@@ -40,7 +44,7 @@ public TestDescriptorHandle(TestDescriptor testDescriptor) {
      * Cloning has to be done before each test retry to
      * compensate for the state modifications.
      */
-    this.testDescriptor = UnsafeUtils.tryShallowClone(testDescriptor);
+    this.testDescriptor = cloneIfNeeded(testDescriptor);
   }
 
   public TestDescriptor withIdSuffix(Map<String, Object> suffices) {
@@ -49,8 +53,99 @@ public TestDescriptor withIdSuffix(Map<String, Object> suffices) {
       updatedId = updatedId.append(e.getKey(), String.valueOf(e.getValue()));
     }
 
-    TestDescriptor descriptorClone = UnsafeUtils.tryShallowClone(testDescriptor);
+    TestDescriptor descriptorClone = cloneIfNeeded(testDescriptor);
     METHOD_HANDLES.invoke(UNIQUE_ID_SETTER, descriptorClone, updatedId);
     return descriptorClone;
   }
+
+  private TestDescriptor cloneIfNeeded(TestDescriptor original) {
+    Class<?> clazz = original.getClass();
+    String name = clazz.getName();
+
+    if (name.endsWith(".TestMethodTestDescriptor")) {
+      // No need to clone for TestMethodTestDescriptor because no states are modified
+      return original;
+    }
+
+    if (name.endsWith(".TestTemplateInvocationTestDescriptor")) {
+      return copyConstructor(
+          original,
+          clazz,
+          "uniqueId",
+          "testClass",
+          "testMethod",
+          "invocationContext",
+          "index",
+          "configuration");
+    }
+
+    if (name.endsWith(".DynamicTestTestDescriptor")) {
+      return copyConstructor(
+          original, clazz, "uniqueId", "index", "dynamicTest", "source", "configuration");
+    }
+
+    throw new IllegalStateException("Unexpected class of TestDescriptor: " + name);
+  }
+
+  private TestDescriptor copyConstructor(TestDescriptor original, Class<?> clazz, String... names) {
+    try {
+      Map<String, Object> fields = getFields(clazz, original);
+      for (Constructor<?> ctr : clazz.getDeclaredConstructors()) {
+        Object[] args = match(ctr, fields, original, names);
+        if (args != null) {
+          ctr.setAccessible(true);
+          return (TestDescriptor) ctr.newInstance(args);
+        }
+      }
+      throw new IllegalStateException(
+          "Failed to find appropriate constructor to clone: " + clazz.getName());
+    } catch (Throwable e) {
+      throw new IllegalStateException("Failed to clone via constructor", e);
+    }
+  }
+
+  private Map<String, Object> getFields(Class<?> clazz, Object original)
+      throws IllegalAccessException {
+    Map<String, Object> fields = new HashMap<>();
+    while (clazz != null && clazz != Object.class) {
+      for (Field field : clazz.getDeclaredFields()) {
+        field.setAccessible(true);
+        Object value = field.get(original);
+
+        // Constructor may need `testClass` and `testMethod` parameters, but they are wrapped by
+        // `methodInfo`.
+        // Expand `methodInfo` to be used in constructor later.
+        if (field.getName().equals("methodInfo")) {
+          Map<String, Object> methodInfoFields = getFields(field.getType(), value);
+          fields.putAll(methodInfoFields);
+        }
+
+        fields.put(field.getName(), value);
+      }
+      clazz = clazz.getSuperclass();
+    }
+    return fields;
+  }
+
+  private Object[] match(
+      Constructor<?> ctr, Map<String, Object> fields, Object original, String... names)
+      throws IllegalAccessException {
+    int cnt = ctr.getParameterCount();
+
+    List<Object> args = new ArrayList<>();
+
+    for (String name : names) {
+      for (Map.Entry<String, Object> field : fields.entrySet()) {
+        String k = field.getKey();
+        Object v = field.getValue();
+
+        if (k.equals(name)) {
+          args.add(v);
+          break;
+        }
+      }
+    }
+
+    return args.size() == cnt ? args.toArray() : null;
+  }
 }

internal-api/build.gradle.kts

@@ -270,6 +270,8 @@ dependencies {
   // it contains annotations that are also present in the instrumented application classes
   api("com.datadoghq:dd-javac-plugin-client:0.2.2")
 
+  implementation(libs.bytebuddy)
+
   testImplementation("org.snakeyaml:snakeyaml-engine:2.9")
   testImplementation(project(":utils:test-utils"))
   testImplementation(libs.bundles.junit5)

internal-api/src/main/java/datadog/trace/util/UnsafeUtils.java

@@ -1,8 +1,16 @@
 package datadog.trace.util;
 
+import static net.bytebuddy.matcher.ElementMatchers.isFinal;
+
 import de.thetaphi.forbiddenapis.SuppressForbidden;
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.asm.ModifierAdjustment;
+import net.bytebuddy.description.modifier.FieldManifestation;
+import net.bytebuddy.dynamic.DynamicType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import sun.misc.Unsafe;
@@ -13,6 +21,8 @@ public abstract class UnsafeUtils {
 
   private static final Unsafe UNSAFE = getUnsafe();
 
+  private static final Map<String, Class<?>> CLASS_CACHE = new ConcurrentHashMap<>();
+
   private static Unsafe getUnsafe() {
     try {
       Field f = Unsafe.class.getDeclaredField("theUnsafe");
@@ -37,7 +47,7 @@ private static Unsafe getUnsafe() {
    * @param <T> Type of the object being cloned
    */
   @SuppressWarnings("unchecked")
-  public static <T> T tryShallowClone(T original) {
+  public static <T> T originalTryShallowClone(T original) {
     if (UNSAFE == null) {
       log.debug("Unsafe is unavailable, {} will not be cloned", original);
       return original;
@@ -58,7 +68,45 @@ public static <T> T tryShallowClone(T original) {
     }
   }
 
-  // TODO: JEP 500 - avoid mutating final fields
+  public static <T> T tryShallowClone(T original) {
+    if (UNSAFE == null) {
+      log.debug("Unsafe is unavailable, {} will not be cloned", original);
+      return original;
+    }
+    try {
+      Class<?> clazz = original.getClass();
+      if (!CLASS_CACHE.containsKey(clazz.getName())) {
+        CLASS_CACHE.put(clazz.getName(), createNonFinalSubclass(clazz));
+      }
+      Class<?> nonFinalSubclass = CLASS_CACHE.get(clazz.getName());
+
+      T clone = (T) UNSAFE.allocateInstance(nonFinalSubclass);
+
+      while (clazz != Object.class) {
+        cloneFields(clazz, original, clone);
+        clazz = clazz.getSuperclass();
+      }
+      return clone;
+
+    } catch (Throwable t) {
+      log.debug("Error while cloning {}: {}", original, t);
+      t.printStackTrace();
+      return original;
+    }
+  }
+
+  private static Class<?> createNonFinalSubclass(Class<?> clazz) throws Exception {
+    DynamicType.Unloaded<?> dynamicType =
+        new ByteBuddy()
+            .subclass(clazz)
+            .visit(new ModifierAdjustment().withFieldModifiers(isFinal(), FieldManifestation.PLAIN))
+            .make();
+    return dynamicType.load(clazz.getClassLoader()).getLoaded();
+  }
+
+  // Field::set() is forbidden because it may be used to mutate final fields, disallowed by
+  // https://openjdk.org/jeps/500.
+  // However, in this case we skip final fields, so it is safe.
   @SuppressForbidden
   private static void cloneFields(Class<?> clazz, Object original, Object clone) throws Exception {
     for (Field field : clazz.getDeclaredFields()) {

internal-api/src/test/groovy/datadog/trace/util/UnsafeUtilsTest.groovy

@@ -4,7 +4,7 @@ import spock.lang.Specification
 
 class UnsafeUtilsTest extends Specification {
 
-  def "test try shallow clone"() {
+  def "test try shallow clone does not clone final fields"() {
     given:
     def inner = new MyClass("a", false, [], "b", 2, null, null)
     def instance = new MyClass("aaa", true, [ 4, 5, 6, ], "ddd", 1, new int[] {
@@ -16,13 +16,27 @@ class UnsafeUtilsTest extends Specification {
 
     then:
     clone !== instance
-    clone.a === instance.a
-    clone.b == instance.b
-    clone.c === instance.c
-    clone.d === instance.d
-    clone.e == instance.e
-    clone.f === instance.f
-    clone.g === instance.g
+    clone.a == null
+    clone.b == false
+    clone.c == null
+    clone.d == null
+    clone.e == 0
+    clone.f == null
+    clone.g == null
+  }
+
+  def "test try shallow clone clones non-final fields"() {
+    given:
+    def instance = new MyNonFinalClass("a", 1, [2, 3, 4])
+
+    when:
+    def clone = UnsafeUtils.tryShallowClone(instance)
+
+    then:
+    clone !== instance
+    clone.h == instance.h
+    clone.j == instance.j
+    clone.k === instance.k
   }
 
   private static class MyParentClass {
@@ -65,4 +79,16 @@ class UnsafeUtilsTest extends Specification {
       this.g = g
     }
   }
+
+  private static class MyNonFinalClass {
+    private String h
+    private int j
+    private List<Integer> k
+
+    MyNonFinalClass(String h, int j, List<Integer> k) {
+      this.h = h
+      this.j = j
+      this.k = k
+    }
+  }
 }