Prevent retransforming record with type annotation

To avoid JVM bug, we detect and prevent retransformation of record
with type annotation on record component.
see https://bugs.openjdk.org/browse/JDK-8376185

Author: jpbempel

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/ConfigurationUpdater.java

@@ -21,10 +21,17 @@
 import datadog.trace.bootstrap.debugger.ProbeRateLimiter;
 import datadog.trace.relocate.api.RatelimitedLogger;
 import datadog.trace.util.TagsHelper;
+import java.lang.annotation.Annotation;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Target;
 import java.lang.instrument.Instrumentation;
+import java.lang.reflect.AnnotatedType;
+import java.lang.reflect.Array;
 import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
 import java.lang.reflect.Parameter;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.EnumMap;
@@ -44,8 +51,29 @@
  * re-transformation of required classes
  */
 public class ConfigurationUpdater implements DebuggerContext.ProbeResolver, ConfigurationAcceptor {
-
+  private static final Logger LOGGER = LoggerFactory.getLogger(ConfigurationUpdater.class);
+  private static final int MINUTES_BETWEEN_ERROR_LOG = 5;
   private static final boolean JAVA_AT_LEAST_19 = JavaVirtualMachine.isJavaVersionAtLeast(19);
+  private static final boolean JAVA_AT_LEAST_16 = JavaVirtualMachine.isJavaVersionAtLeast(16);
+  private static final Method GET_RECORD_COMPONENTS_METHOD;
+  private static final Method GET_ANNOTATED_TYPES_METHOD;
+
+  static {
+    Method getRecordComponentsMethod = null;
+    Method getAnnotatedTypesMethod = null;
+    if (JAVA_AT_LEAST_16) {
+      try {
+        Class<?> recordClass = Class.forName("java.lang.Record");
+        getRecordComponentsMethod = recordClass.getClass().getDeclaredMethod("getRecordComponents");
+        Class<?> recordComponentClass = Class.forName("java.lang.reflect.RecordComponent");
+        getAnnotatedTypesMethod = recordComponentClass.getDeclaredMethod("getAnnotatedType");
+      } catch (Exception e) {
+        LOGGER.debug("Exception initializing reflection constants", e);
+      }
+    }
+    GET_RECORD_COMPONENTS_METHOD = getRecordComponentsMethod;
+    GET_ANNOTATED_TYPES_METHOD = getAnnotatedTypesMethod;
+  }
 
   public interface TransformerSupplier {
     DebuggerTransformer supply(
@@ -56,9 +84,6 @@ DebuggerTransformer supply(
         DebuggerSink debuggerSink);
   }
 
-  private static final Logger LOGGER = LoggerFactory.getLogger(ConfigurationUpdater.class);
-  private static final int MINUTES_BETWEEN_ERROR_LOG = 5;
-
   private final Instrumentation instrumentation;
   private final TransformerSupplier transformerSupplier;
   private final Lock configurationLock = new ReentrantLock();
@@ -185,6 +210,7 @@ private void handleProbesChanges(ConfigurationComparer changes, Configuration ne
     List<Class<?>> changedClasses =
         finder.getAllLoadedChangedClasses(instrumentation.getAllLoadedClasses(), changes);
     changedClasses = detectMethodParameters(changes, changedClasses);
+    changedClasses = detectRecordWithTypeAnnotation(changes, changedClasses);
     retransformClasses(changedClasses);
     // ensures that we have at least re-transformed 1 class
     if (changedClasses.size() > 0) {
@@ -248,6 +274,65 @@ private List<Class<?>> detectMethodParameters(
     return result;
   }
 
+  private List<Class<?>> detectRecordWithTypeAnnotation(
+      ConfigurationComparer changes, List<Class<?>> changedClasses) {
+    if (!JAVA_AT_LEAST_16) {
+      // records introduced in JDK 16 (final version)
+      return changedClasses;
+    }
+    List<Class<?>> result = new ArrayList<>();
+    for (Class<?> changedClass : changedClasses) {
+      boolean addClass = true;
+      try {
+        if (changedClass.getSuperclass().getTypeName().equals("java.lang.Record")
+            && Modifier.isFinal(changedClass.getModifiers())) {
+          if (hasTypeAnnotationOnRecordComponent(changedClass)) {
+            LOGGER.debug(
+                "Record with type annotation detected, instrumentation not supported for {}",
+                changedClass.getTypeName());
+            reportError(
+                changes,
+                "Record with type annotation detected, instrumentation not supported for "
+                    + changedClass.getTypeName());
+            addClass = false;
+          }
+        }
+      } catch (Exception e) {
+        LOGGER.debug("Exception detecting record with type annotation", e);
+      }
+      if (addClass) {
+        result.add(changedClass);
+      }
+    }
+    return result;
+  }
+
+  private boolean hasTypeAnnotationOnRecordComponent(Class<?> recordClass) {
+    if (GET_RECORD_COMPONENTS_METHOD == null || GET_ANNOTATED_TYPES_METHOD == null) {
+      return false;
+    }
+    try {
+      Object recordComponentsArray = GET_RECORD_COMPONENTS_METHOD.invoke(recordClass);
+      int len = Array.getLength(recordComponentsArray);
+      for (int i = 0; i < len; i++) {
+        Object recordComponent = Array.get(recordComponentsArray, i);
+        AnnotatedType annotatedType =
+            (AnnotatedType) GET_ANNOTATED_TYPES_METHOD.invoke(recordComponent);
+        for (Annotation annotation : annotatedType.getAnnotations()) {
+          Target annotationTarget = annotation.annotationType().getAnnotation(Target.class);
+          if (annotationTarget != null
+              && Arrays.stream(annotationTarget.value())
+                  .anyMatch(it -> it == ElementType.TYPE_USE)) {
+            return true;
+          }
+        }
+      }
+      return false;
+    } catch (Exception ex) {
+      return false;
+    }
+  }
+
   private void reportReceived(ConfigurationComparer changes) {
     for (ProbeDefinition def : changes.getAddedDefinitions()) {
       if (def instanceof ExceptionProbe) {

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/agent/DebuggerTransformer.java

@@ -67,6 +67,7 @@
 import org.objectweb.asm.Type;
 import org.objectweb.asm.commons.JSRInlinerAdapter;
 import org.objectweb.asm.tree.ClassNode;
+import org.objectweb.asm.tree.FieldNode;
 import org.objectweb.asm.tree.MethodNode;
 import org.objectweb.asm.tree.analysis.Analyzer;
 import org.objectweb.asm.tree.analysis.AnalyzerException;
@@ -275,6 +276,9 @@ public byte[] transform(
       }
       ClassNode classNode = parseClassFile(classFilePath, classfileBuffer);
       checkMethodParameters(classNode);
+      if (!checkRecordTypeAnnotation(classNode, definitions, fullyQualifiedClassName)) {
+        return null;
+      }
       boolean transformed =
           performInstrumentation(loader, fullyQualifiedClassName, definitions, classNode);
       if (transformed) {
@@ -333,6 +337,37 @@ private void checkMethodParameters(ClassNode classNode) {
     }
   }
 
+  /*
+   * Because of this bug (https://bugs.openjdk.org/browse/JDK-8376185), when a record using a type
+   * annotation is retransformed, the internal JVM representation of this record is corrupted
+   * and lead to exception in best cases but in JVM crashes in worst cases.
+   * Note: the bug happens only at retransform time and not instrumenting at load time. But the
+   * fact we have already instrumented the record at load time, will prevent us to remove the
+   * instrumentation because it needs a retransformation and will lead to corruption of the record
+   */
+  private boolean checkRecordTypeAnnotation(
+      ClassNode classNode, List<ProbeDefinition> definitions, String fullyQualifiedClassName) {
+    if (!ASMHelper.isRecord(classNode)) {
+      return true;
+    }
+    if (classNode.fields == null || classNode.fields.isEmpty()) {
+      return true;
+    }
+    for (FieldNode field : classNode.fields) {
+      if ((field.visibleTypeAnnotations != null && !field.visibleTypeAnnotations.isEmpty())
+          || (field.invisibleTypeAnnotations != null
+              && !field.invisibleTypeAnnotations.isEmpty())) {
+        reportInstrumentationFails(
+            definitions,
+            fullyQualifiedClassName,
+            "Instrumentation of a record with type annotation is not supported");
+        return false;
+      }
+    }
+    // no type annotation for components, not a problem
+    return true;
+  }
+
   private boolean skipInstrumentation(ClassLoader loader, String classFilePath) {
     if (definitionMatcher.isEmpty()) {
       LOGGER.debug("No debugger definitions present.");

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/CapturedSnapshotTest.java

@@ -3132,6 +3132,23 @@ public void noInstrumentationForAgentClasses() throws Exception {
     assertNull(result);
   }
 
+  @Test
+  @EnabledForJreRange(min = JRE.JAVA_17)
+  public void recordWithTypeAnnotation() throws IOException, URISyntaxException {
+    final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot33";
+    LogProbe probe1 = createMethodProbeAtExit(PROBE_ID1, CLASS_NAME, "parse", null);
+    TestSnapshotListener listener = installProbes(probe1);
+    Class<?> testClass = compileAndLoadClass(CLASS_NAME, "17");
+    Reflect.onClass(testClass).call("parse", "1").get();
+    ArgumentCaptor<ProbeId> probeIdCaptor = ArgumentCaptor.forClass(ProbeId.class);
+    ArgumentCaptor<String> strCaptor = ArgumentCaptor.forClass(String.class);
+    verify(probeStatusSink, times(1)).addError(probeIdCaptor.capture(), strCaptor.capture());
+    assertEquals(PROBE_ID1.getId(), probeIdCaptor.getAllValues().get(0).getId());
+    assertEquals(
+        "Instrumentation failed for com.datadog.debugger.CapturedSnapshot33: Instrumentation of a record with type annotation is not supported",
+        strCaptor.getAllValues().get(0));
+  }
+
   private TestSnapshotListener setupInstrumentTheWorldTransformer(
       String excludeFileName, String includeFileName) {
     Config config = mock(Config.class);

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/ConfigurationUpdaterTest.java

@@ -690,6 +690,25 @@ public void methodParametersAttributeRecord()
     assertEquals(testClass, allValues.get(0));
   }
 
+  @Test
+  @EnabledForJreRange(min = JRE.JAVA_17)
+  public void recordWithTypeAnnotation()
+      throws IOException, URISyntaxException, UnmodifiableClassException {
+    // make sure record method are not detected as having methodParameters attribute.
+    // /!\ record canonical constructor has the MethodParameters attribute,
+    // but not returned by Class::getDeclaredMethods()
+    final String CLASS_NAME = "com.datadog.debugger.CapturedSnapshot33";
+    Map<String, byte[]> buffers = compile(CLASS_NAME, SourceCompiler.DebugInfo.ALL, "17");
+    Class<?> testClass = loadClass(CLASS_NAME, buffers);
+    when(inst.getAllLoadedClasses()).thenReturn(new Class[] {testClass});
+    ConfigurationUpdater configurationUpdater = createConfigUpdater(debuggerSinkWithMockStatusSink);
+    configurationUpdater.accept(
+        REMOTE_CONFIG,
+        singletonList(LogProbe.builder().probeId(PROBE_ID).where(CLASS_NAME, "parse").build()));
+    verify(inst).getAllLoadedClasses();
+    verify(inst, times(0)).retransformClasses(any());
+  }
+
   private DebuggerTransformer createTransformer(
       Config tracerConfig,
       Configuration configuration,

dd-java-agent/agent-debugger/src/test/resources/com/datadog/debugger/CapturedSnapshot33.java

@@ -0,0 +1,24 @@
+package com.datadog.debugger;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@MyTypeAnnotation
+public record CapturedSnapshot33(@MyTypeUseAnnotation String strField) {
+  public static CapturedSnapshot33 parse(String arg) {
+    return new CapturedSnapshot33(arg);
+  }
+}
+
+
+@Target({ElementType.TYPE})
+@Retention(RetentionPolicy.RUNTIME)
+@interface MyTypeAnnotation {
+}
+
+@Target({ElementType.TYPE_USE})
+@Retention(RetentionPolicy.RUNTIME)
+@interface MyTypeUseAnnotation {
+}