Optimize the SourceFile tracking

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/util/ClassFileHelper.java

@@ -1,12 +1,20 @@
 package com.datadog.debugger.util;
 
+import java.nio.charset.StandardCharsets;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.tree.ClassNode;
 
 /** Helper class for extracting information of a class file */
 public class ClassFileHelper {
+  private static final int CONSTANT_POOL_COUNT_OFFSET = 8;
+  private static final int CONSTANT_POOL_BASE_OFFSET = 10;
+
   public static String extractSourceFile(byte[] classFileBuffer) {
-    // TODO maybe by scanning the byte array directly we can avoid doing an expensive parsing
+    return extractSourceFileOffsetVersion(classFileBuffer);
+  }
+
+  // Version using ASM library
+  private static String extractSourceFileASM(byte[] classFileBuffer) {
     ClassReader classReader = new ClassReader(classFileBuffer);
     ClassNode classNode = new ClassNode();
     classReader.accept(classNode, ClassReader.SKIP_FRAMES);
@@ -33,4 +41,126 @@ public static String stripPackagePath(String classPath) {
     }
     return classPath;
   }
+
+  // Based on JVM spec: https://docs.oracle.com/javase/specs/jvms/se23/html/jvms-4.html
+  // Extracts the SourceFile attribute from a Java class file byte array with minimal parsing.
+  // This method is based on the JVM spec and does not use any external libraries.
+  // We are scanning the constant pool to keep file offsets for later fetching of the SourceFile
+  // attribute value. As the constant pool is a variable length structure, we need to scan them
+  // and based on the tag, we can calculate the length of the entry to skip to the next one.
+  private static String extractSourceFileOffsetVersion(byte[] classFileBytes) {
+    // Quick validation of minimum class file size and magic number
+    if (classFileBytes == null
+        || classFileBytes.length < 10
+        || classFileBytes[0] != (byte) 0xCA
+        || classFileBytes[1] != (byte) 0xFE
+        || classFileBytes[2] != (byte) 0xBA
+        || classFileBytes[3] != (byte) 0xBE) {
+      return null;
+    }
+    int constantPoolCount = readUnsignedShort(classFileBytes, CONSTANT_POOL_COUNT_OFFSET);
+    int[] constantPoolOffsets = new int[constantPoolCount];
+    int currentOffset = CONSTANT_POOL_BASE_OFFSET;
+    // based on the JVM spec, constant pool starts from index 1 until constantPoolCount - 1
+    for (int i = 0; i < constantPoolCount - 1; i++) {
+      constantPoolOffsets[i] = currentOffset;
+      int tag = classFileBytes[constantPoolOffsets[i]];
+      switch (tag) {
+        case 1: // CONSTANT_Utf8
+          int length = readUnsignedShort(classFileBytes, constantPoolOffsets[i] + 1);
+          currentOffset += 3 + length;
+          break;
+        case 7: // CONSTANT_Class
+        case 8: // CONSTANT_String
+        case 16: // CONSTANT_MethodType
+        case 19: // CONSTANT_Module
+        case 20: // CONSTANT_Package
+          currentOffset += 3;
+          break;
+        case 15: // CONSTANT_MethodHandle
+          currentOffset += 4;
+          break;
+        case 3: // CONSTANT_Integer
+        case 4: // CONSTANT_Float
+        case 9: // CONSTANT_Fieldref
+        case 10: // CONSTANT_Methodref
+        case 11: // CONSTANT_InterfaceMethodref
+        case 12: // CONSTANT_NameAndType
+        case 17: // CONSTANT_Dynamic
+        case 18: // CONSTANT_InvokeDynamic
+          currentOffset += 5;
+          break;
+        case 5: // CONSTANT_Long
+        case 6: // CONSTANT_Double
+          currentOffset += 9;
+          i++; // Double slot
+          break;
+        default:
+          throw new IllegalArgumentException("Unknown constant pool tag: " + tag);
+      }
+    }
+    currentOffset += 2; // Skip access flags
+    currentOffset += 2; // Skip this class
+    currentOffset += 2; // Skip super class
+    int interfacesCount = readUnsignedShort(classFileBytes, currentOffset);
+    currentOffset += 2 + interfacesCount * 2; // Skip interfaces
+    // skip fields
+    currentOffset = skipFieldsOrMethods(classFileBytes, currentOffset);
+    // skip Methods
+    currentOffset = skipFieldsOrMethods(classFileBytes, currentOffset);
+    int attributesCount = readUnsignedShort(classFileBytes, currentOffset);
+    currentOffset += 2; // Skip attributes count
+    for (int i = 0; i < attributesCount; i++) {
+      int attributeNameIndex = readUnsignedShort(classFileBytes, currentOffset);
+      currentOffset += 2; // Skip attribute name index
+      int attributeLength = (int) readUnsignedInt(classFileBytes, currentOffset);
+      currentOffset += 4; // Skip attribute length
+      if (attributeNameIndex == 0) {
+        continue;
+      }
+      // read attribute name
+      int utf8Offset = constantPoolOffsets[attributeNameIndex - 1];
+      int utf8Len = readUnsignedShort(classFileBytes, utf8Offset + 1);
+      String utf8 = new String(classFileBytes, utf8Offset + 3, utf8Len, StandardCharsets.UTF_8);
+      if ("SourceFile".equals(utf8)) {
+        // read SourceFile attribute
+        int sourceFileIndex = readUnsignedShort(classFileBytes, currentOffset);
+        int sourceFileOffset = constantPoolOffsets[sourceFileIndex - 1];
+        int sourceFileLen = readUnsignedShort(classFileBytes, sourceFileOffset + 1);
+        return new String(
+            classFileBytes, sourceFileOffset + 3, sourceFileLen, StandardCharsets.UTF_8);
+      }
+      currentOffset += attributeLength; // Skip attribute data
+    }
+    return null;
+  }
+
+  private static int skipFieldsOrMethods(byte[] classFileBytes, int currentOffset) {
+    int fieldsCount = readUnsignedShort(classFileBytes, currentOffset);
+    currentOffset += 2; // Skip count
+    for (int i = 0; i < fieldsCount; i++) {
+      currentOffset += 6; // Skip access flags, name index, descriptor index
+      int attributesCount = readUnsignedShort(classFileBytes, currentOffset);
+      currentOffset += 2; // Skip attributes count
+      for (int j = 0; j < attributesCount; j++) {
+        currentOffset += 2; // Skip attribute name index
+        int attributeLength = (int) readUnsignedInt(classFileBytes, currentOffset);
+        currentOffset += 4 + attributeLength; // Skip attribute length and data
+      }
+    }
+    return currentOffset;
+  }
+
+  // read unsigned short from byte array
+  private static int readUnsignedShort(byte[] bytes, int offset) {
+    return ((bytes[offset] & 0xFF) << 8) | (bytes[offset + 1] & 0xFF);
+  }
+
+  // read unsigned int from byte array
+  private static long readUnsignedInt(byte[] bytes, int offset) {
+    return ((long) (bytes[offset] & 0xFF) << 24)
+        + ((bytes[offset + 1] & 0xFF) << 16)
+        + ((bytes[offset + 2] & 0xFF) << 8)
+        + (bytes[offset + 3] & 0xFF);
+  }
 }

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

@@ -7,7 +7,6 @@
 import static com.datadog.debugger.util.MoshiSnapshotHelper.REDACTED_IDENT_REASON;
 import static com.datadog.debugger.util.MoshiSnapshotHelper.REDACTED_TYPE_REASON;
 import static com.datadog.debugger.util.MoshiSnapshotTestHelper.VALUE_ADAPTER;
-import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
 import static datadog.trace.bootstrap.debugger.util.Redaction.REDACTED_VALUE;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -27,6 +26,7 @@
 import static utils.InstrumentationTestHelper.getLineForLineProbe;
 import static utils.InstrumentationTestHelper.loadClass;
 import static utils.TestHelper.getFixtureContent;
+import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.el.DSL;
 import com.datadog.debugger.el.ProbeCondition;

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

@@ -2,11 +2,11 @@
 
 import static com.datadog.debugger.util.MoshiSnapshotHelper.NOT_CAPTURED_REASON;
 import static com.datadog.debugger.util.MoshiSnapshotTestHelper.VALUE_ADAPTER;
-import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
+import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.instrumentation.InstrumentationResult;
 import com.datadog.debugger.probe.LogProbe;

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

@@ -1,7 +1,5 @@
 package com.datadog.debugger.agent;
 
-import static com.datadog.debugger.util.TestHelper.setEnvVar;
-import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
@@ -14,6 +12,8 @@
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static utils.TestHelper.setEnvVar;
+import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.util.RemoteConfigHelper;
 import datadog.common.container.ContainerInfo;

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

@@ -1,6 +1,5 @@
 package com.datadog.debugger.agent;
 
-import static com.datadog.debugger.util.ClassFileHelperTest.getClassFileBytes;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.mockito.ArgumentMatchers.any;
@@ -10,6 +9,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static utils.TestClassFileHelper.getClassFileBytes;
 
 import com.datadog.debugger.instrumentation.DiagnosticMessage;
 import com.datadog.debugger.instrumentation.InstrumentationResult;

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

@@ -1,7 +1,7 @@
 package com.datadog.debugger.agent;
 
-import static com.datadog.debugger.util.ClassFileHelperTest.getClassFileBytes;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static utils.TestClassFileHelper.getClassFileBytes;
 
 import com.datadog.debugger.probe.LogProbe;
 import java.lang.instrument.IllegalClassFormatException;

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

@@ -1,8 +1,8 @@
 package com.datadog.debugger.agent;
 
-import static com.datadog.debugger.util.ClassFileHelperTest.getClassFileBytes;
 import static java.util.Arrays.asList;
 import static org.junit.jupiter.api.Assertions.*;
+import static utils.TestClassFileHelper.getClassFileBytes;
 
 import com.datadog.debugger.probe.LogProbe;
 import com.datadog.debugger.probe.MetricProbe;

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/DefaultExceptionDebuggerTest.java

@@ -1,7 +1,6 @@
 package com.datadog.debugger.exception;
 
 import static com.datadog.debugger.exception.DefaultExceptionDebugger.SNAPSHOT_ID_TAG_FMT;
-import static com.datadog.debugger.util.TestHelper.assertWithTimeout;
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 import static java.util.stream.Collectors.toList;
@@ -16,6 +15,7 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static utils.TestHelper.assertWithTimeout;
 
 import com.datadog.debugger.agent.ConfigurationAcceptor;
 import com.datadog.debugger.agent.ConfigurationUpdater;

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/ExceptionProbeInstrumentationTest.java

@@ -6,13 +6,13 @@
 import static com.datadog.debugger.exception.DefaultExceptionDebugger.ERROR_DEBUG_INFO_CAPTURED;
 import static com.datadog.debugger.exception.DefaultExceptionDebugger.SNAPSHOT_ID_TAG_FMT;
 import static com.datadog.debugger.util.MoshiSnapshotTestHelper.getValue;
-import static com.datadog.debugger.util.TestHelper.assertWithTimeout;
-import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static utils.InstrumentationTestHelper.compileAndLoadClass;
+import static utils.TestHelper.assertWithTimeout;
+import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.agent.ClassesToRetransformFinder;
 import com.datadog.debugger.agent.Configuration;

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/origin/CodeOriginTest.java

@@ -1,6 +1,5 @@
 package com.datadog.debugger.origin;
 
-import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
 import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_FRAME;
 import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_PREFIX;
 import static datadog.trace.api.DDTags.DD_CODE_ORIGIN_TYPE;
@@ -14,6 +13,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static utils.InstrumentationTestHelper.compileAndLoadClass;
+import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.agent.CapturingTestBase;
 import com.datadog.debugger.codeorigin.DefaultCodeOriginRecorder;

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/trigger/TriggerProbeTest.java

@@ -1,13 +1,13 @@
 package com.datadog.debugger.trigger;
 
 import static com.datadog.debugger.el.DSL.*;
-import static com.datadog.debugger.util.TestHelper.setFieldInConfig;
 import static java.lang.String.format;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 import static utils.InstrumentationTestHelper.compileAndLoadClass;
+import static utils.TestHelper.setFieldInConfig;
 
 import com.datadog.debugger.agent.CapturingTestBase;
 import com.datadog.debugger.agent.Configuration;

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/util/ClassFileHelperTest.java

@@ -1,25 +1,37 @@
 package com.datadog.debugger.util;
 
-import static datadog.trace.util.Strings.getResourceName;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
-import java.net.URL;
+import java.net.URISyntaxException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import org.junit.jupiter.api.Test;
 
 public class ClassFileHelperTest {
-  public static byte[] getClassFileBytes(Class<?> clazz) {
-    URL resource = clazz.getResource("/" + getResourceName(clazz.getTypeName()));
-    byte[] buffer = new byte[4096];
-    ByteArrayOutputStream os = new ByteArrayOutputStream();
-    try (InputStream is = resource.openStream()) {
-      int readBytes;
-      while ((readBytes = is.read(buffer)) != -1) {
-        os.write(buffer, 0, readBytes);
-      }
-    } catch (IOException e) {
-      e.printStackTrace();
+
+  @Test
+  public void extractSourceFile() {
+    assertEquals(
+        "JDK8.java",
+        ClassFileHelper.extractSourceFile(
+            readClassFileBytes("/com/datadog/debugger/classfiles/JDK8.class")));
+    assertEquals(
+        "JDK23.java",
+        ClassFileHelper.extractSourceFile(
+            readClassFileBytes("/com/datadog/debugger/classfiles/JDK23.class")));
+    // big classfile (80KB)
+    assertEquals(
+        "CommandLine.java",
+        ClassFileHelper.extractSourceFile(
+            readClassFileBytes("/com/datadog/debugger/classfiles/CommandLine.class")));
+  }
+
+  private static byte[] readClassFileBytes(String fileName) {
+    try {
+      return Files.readAllBytes(Paths.get(ClassFileHelperTest.class.getResource(fileName).toURI()));
+    } catch (IOException | URISyntaxException e) {
+      throw new RuntimeException(e);
     }
-    return os.toByteArray();
   }
 }