Fix exact return-type matching and expand class-match tests

class-match/build.gradle.kts

@@ -4,6 +4,7 @@ plugins {
 
 dependencies {
   implementation(project(":utils"))
+  testImplementation(libs.asm)
 
   sampleBytecode(libs.asm.test)
   sampleBytecode("org.springframework:spring-web:7.0.9")

class-match/src/main/java/datadog/instrument/classmatch/MethodMatcher.java

@@ -186,8 +186,8 @@ default MethodMatcher parameter(int paramIndex, TypeMatcher typeMatcher) {
    * @return matcher of methods returning the same type
    */
   default MethodMatcher returning(String returnType) {
-    String returnDescriptor = descriptor(returnType);
-    return and(m -> m.descriptor.endsWith(returnDescriptor));
+    String returnSuffix = ')' + descriptor(returnType);
+    return and(m -> m.descriptor.endsWith(returnSuffix));
   }
 
   /**
@@ -197,8 +197,8 @@ default MethodMatcher returning(String returnType) {
    * @return matcher of methods returning the same type
    */
   default MethodMatcher returning(Class<?> returnType) {
-    String returnDescriptor = descriptor(returnType);
-    return and(m -> m.descriptor.endsWith(returnDescriptor));
+    String returnSuffix = ')' + descriptor(returnType);
+    return and(m -> m.descriptor.endsWith(returnSuffix));
   }
 
   /**

class-match/src/test/java/datadog/instrument/classmatch/AccessMatcherTest.java

@@ -0,0 +1,77 @@
+package datadog.instrument.classmatch;
+
+import static datadog.instrument.classmatch.AccessMatcher.ABSTRACT;
+import static datadog.instrument.classmatch.AccessMatcher.CLASS;
+import static datadog.instrument.classmatch.AccessMatcher.CONCRETE;
+import static datadog.instrument.classmatch.AccessMatcher.FINAL;
+import static datadog.instrument.classmatch.AccessMatcher.INSTANCE;
+import static datadog.instrument.classmatch.AccessMatcher.INTERFACE;
+import static datadog.instrument.classmatch.AccessMatcher.NON_FINAL;
+import static datadog.instrument.classmatch.AccessMatcher.PACKAGE_PRIVATE;
+import static datadog.instrument.classmatch.AccessMatcher.PRIVATE;
+import static datadog.instrument.classmatch.AccessMatcher.PROTECTED;
+import static datadog.instrument.classmatch.AccessMatcher.PUBLIC;
+import static datadog.instrument.classmatch.AccessMatcher.STATIC;
+import static datadog.instrument.classmatch.AccessMatcher.SYNCHRONIZED;
+import static datadog.instrument.classmatch.AccessMatcher.TRANSIENT;
+import static datadog.instrument.classmatch.AccessMatcher.VARARGS;
+import static datadog.instrument.classmatch.AccessMatcher.VOLATILE;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+class AccessMatcherTest {
+  static Stream<Arguments> flags() {
+    return Stream.of(
+        Arguments.of("public", PUBLIC, 0x0001, false),
+        Arguments.of("private", PRIVATE, 0x0002, false),
+        Arguments.of("protected", PROTECTED, 0x0004, false),
+        Arguments.of("package-private", PACKAGE_PRIVATE, 0x0007, true),
+        Arguments.of("static", STATIC, 0x0008, false),
+        Arguments.of("instance", INSTANCE, 0x0008, true),
+        Arguments.of("final", FINAL, 0x0010, false),
+        Arguments.of("non-final", NON_FINAL, 0x0010, true),
+        Arguments.of("synchronized", SYNCHRONIZED, 0x0020, false),
+        Arguments.of("volatile", VOLATILE, 0x0040, false),
+        Arguments.of("varargs", VARARGS, 0x0080, false),
+        Arguments.of("transient", TRANSIENT, 0x0080, false),
+        Arguments.of("interface", INTERFACE, 0x0200, false),
+        Arguments.of("class", CLASS, 0x0200, true),
+        Arguments.of("abstract", ABSTRACT, 0x0400, false),
+        Arguments.of("concrete", CONCRETE, 0x0400, true));
+  }
+
+  @ParameterizedTest(name = "{0}")
+  @MethodSource("flags")
+  void flagsIgnoreUnrelatedBits(String name, AccessMatcher matcher, int flag, boolean inverted) {
+    // Exercise every combination of the JVM's 16 access bits, including synthetic/bridge bits.
+    for (int access = 0; access <= 0xffff; access++) {
+      boolean expected = ((access & flag) != 0) != inverted;
+      if (matcher.test(access) != expected) {
+        fail(name + ": expected " + expected + " for access " + access);
+      }
+    }
+  }
+
+  @Test
+  void compositionAndNegation() {
+    for (int access = 0; access < 32; access++) {
+      assertEquals(PUBLIC.test(access) && STATIC.test(access), PUBLIC.and(STATIC).test(access));
+      assertEquals(PUBLIC.test(access) || PRIVATE.test(access), PUBLIC.or(PRIVATE).test(access));
+      assertEquals(!PUBLIC.test(access), StandardMatchers.not(PUBLIC).test(access));
+    }
+    AccessMatcher unexpected =
+        access -> {
+          throw new AssertionError("should short circuit");
+        };
+    assertFalse(PUBLIC.and(unexpected).test(0));
+    assertTrue(PUBLIC.or(unexpected).test(1));
+  }
+}

class-match/src/test/java/datadog/instrument/classmatch/AnnotationRegistrationProbe.java

@@ -0,0 +1,41 @@
+package datadog.instrument.classmatch;
+
+import static datadog.instrument.classmatch.BytecodeTestSupport.annotatedClass;
+
+import java.util.Arrays;
+import java.util.Collections;
+
+/** Entry point for testing initial static state in a separate JVM, independent of JUnit order. */
+public final class AnnotationRegistrationProbe {
+  private AnnotationRegistrationProbe() {}
+
+  public static void main(String[] args) {
+    byte[] bytes = annotatedClass("tests/First", "tests/Second");
+    expect(bytes); // no registration has occurred in this JVM
+    if ("single".equals(args[0])) {
+      ClassFile.annotationOfInterest("tests/First");
+    } else {
+      ClassFile.annotationsOfInterest(Collections.singletonList("tests/First"));
+    }
+    expect(bytes, "tests/First");
+    // Both registration APIs must preserve earlier names and tolerate repeated registration.
+    ClassFile.annotationsOfInterest(Arrays.asList("tests/First", "tests/Second"));
+    ClassFile.annotationOfInterest("tests/First");
+    ClassFile.annotationsOfInterest(Collections.emptyList());
+    expect(bytes, "tests/First", "tests/Second");
+  }
+
+  private static void expect(byte[] bytes, String... expected) {
+    ClassOutline outline = ClassFile.outline(bytes);
+    if (!Arrays.equals(expected, outline.annotations)
+        || !Arrays.equals(expected, outline.methods[0].annotations)) {
+      throw new AssertionError(
+          "Expected "
+              + Arrays.toString(expected)
+              + ", class="
+              + Arrays.toString(outline.annotations)
+              + ", method="
+              + Arrays.toString(outline.methods[0].annotations));
+    }
+  }
+}

class-match/src/test/java/datadog/instrument/classmatch/AnnotationRegistrationTest.java

@@ -0,0 +1,109 @@
+package datadog.instrument.classmatch;
+
+import static datadog.instrument.classmatch.BytecodeTestSupport.annotatedClass;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+class AnnotationRegistrationTest {
+  @TempDir Path temporary;
+
+  @ParameterizedTest
+  @ValueSource(strings = {"single", "collection"})
+  void firstRegistrationWorksInAFreshJvm(String mode) throws Exception {
+    Path output = temporary.resolve("probe-" + mode + ".log");
+    Process process =
+        new ProcessBuilder(
+                Paths.get(System.getProperty("java.home"), "bin", "java").toString(),
+                "-cp",
+                System.getProperty("java.class.path"),
+                AnnotationRegistrationProbe.class.getName(),
+                mode)
+            .redirectErrorStream(true)
+            .redirectOutput(output.toFile())
+            .start();
+    try {
+      assertTrue(process.waitFor(20, TimeUnit.SECONDS), "annotation probe timed out");
+      assertEquals(
+          0, process.exitValue(), new String(Files.readAllBytes(output), StandardCharsets.UTF_8));
+    } finally {
+      process.destroyForcibly();
+      process.waitFor(5, TimeUnit.SECONDS);
+    }
+  }
+
+  @Test
+  void concurrentRegistrationPublishesCompleteSnapshotsWithoutLosingNames() throws Exception {
+    int writers = 4, readers = 2, registrations = 12;
+    List<String> names = new ArrayList<>();
+    for (int i = 0; i < writers * registrations; i++) {
+      names.add("tests/Concurrent/Marker" + i);
+    }
+    byte[] bytes = annotatedClass(names.toArray(new String[0]));
+    Set<String> allowed = new HashSet<>(names);
+    CountDownLatch start = new CountDownLatch(1);
+    ExecutorService pool = Executors.newFixedThreadPool(writers + readers);
+    List<Future<?>> tasks = new ArrayList<>();
+    try {
+      for (int worker = 0; worker < writers + readers; worker++) {
+        final int id = worker;
+        tasks.add(
+            pool.submit(
+                () -> {
+                  start.await();
+                  Set<String> registered = new HashSet<>();
+                  for (int i = 0; i < registrations; i++) {
+                    if (id < writers) {
+                      String name = names.get(id * registrations + i);
+                      if ((id & 1) == 0) {
+                        ClassFile.annotationOfInterest(name);
+                      } else {
+                        ClassFile.annotationsOfInterest(Arrays.asList(name, name));
+                      }
+                      registered.add(name);
+                    }
+                    assertSnapshot(bytes, registered, allowed);
+                  }
+                  return null;
+                }));
+      }
+      start.countDown();
+      for (Future<?> task : tasks) {
+        task.get(30, TimeUnit.SECONDS);
+      }
+      assertSnapshot(bytes, allowed, allowed);
+    } finally {
+      pool.shutdownNow();
+      assertTrue(pool.awaitTermination(15, TimeUnit.SECONDS), "annotation workers did not stop");
+    }
+  }
+
+  private static void assertSnapshot(byte[] bytes, Set<String> required, Set<String> allowed) {
+    ClassOutline outline = ClassFile.outline(bytes);
+    for (String[] annotations :
+        new String[][] {outline.annotations, outline.methods[0].annotations}) {
+      Set<String> actual = new HashSet<>(Arrays.asList(annotations));
+      assertEquals(annotations.length, actual.size(), "snapshot contains duplicate names");
+      assertTrue(actual.containsAll(required), "snapshot lost registered names");
+      assertTrue(allowed.containsAll(actual), "snapshot contains unexpected names");
+    }
+  }
+}

class-match/src/test/java/datadog/instrument/classmatch/BytecodeTestSupport.java

@@ -0,0 +1,32 @@
+package datadog.instrument.classmatch;
+
+import static org.objectweb.asm.Opcodes.ACC_ABSTRACT;
+import static org.objectweb.asm.Opcodes.ACC_PUBLIC;
+import static org.objectweb.asm.Opcodes.V1_8;
+
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.MethodVisitor;
+
+/** ASM fixtures are independent of the parser and descriptor conversion under test. */
+final class BytecodeTestSupport {
+  private BytecodeTestSupport() {}
+
+  static ClassWriter classWriter(String name) {
+    ClassWriter writer = new ClassWriter(0);
+    writer.visit(V1_8, ACC_PUBLIC | ACC_ABSTRACT, name, null, "java/lang/Object", null);
+    return writer;
+  }
+
+  static byte[] annotatedClass(String... names) {
+    ClassWriter writer = classWriter("tests/Annotated");
+    MethodVisitor method =
+        writer.visitMethod(ACC_PUBLIC | ACC_ABSTRACT, "marked", "()V", null, null);
+    for (String name : names) {
+      writer.visitAnnotation("L" + name + ";", true).visitEnd();
+      method.visitAnnotation("L" + name + ";", true).visitEnd();
+    }
+    method.visitEnd();
+    writer.visitEnd();
+    return writer.toByteArray();
+  }
+}

class-match/src/test/java/datadog/instrument/classmatch/ClassFileByteKeyTest.java

@@ -0,0 +1,45 @@
+package datadog.instrument.classmatch;
+
+import static java.nio.charset.StandardCharsets.US_ASCII;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+class ClassFileByteKeyTest {
+  @Test
+  void byteRangesUseTheirContentAndIgnoreSurroundingBytes() {
+    byte[] value = {0, -1, 127, -128};
+    byte[] surrounded = {42, 0, -1, 127, -128, 99};
+    ClassFile.UtfKey whole = new ClassFile.UtfKey(value);
+    ClassFile.UtfKey window = new ClassFile.UtfKey(surrounded, 1, value.length);
+    assertEquals(Arrays.hashCode(value), whole.hashCode());
+    assertEquals(whole.hashCode(), window.hashCode());
+    assertEquals(whole, window);
+    assertEquals(window, whole);
+    assertTrue(ClassFile.sameBytes(surrounded, 1, 5, value, 0));
+    assertEquals(Arrays.hashCode(value), ClassFile.hashBytes(surrounded, 1, 5));
+    assertEquals(1, ClassFile.hashBytes(value, 2, 2));
+    assertTrue(ClassFile.sameBytes(value, 2, 2, new byte[0], 0));
+    assertFalse(whole.equals(null));
+    assertFalse(whole.equals(new byte[] {0, -1, 127, -128}));
+    assertFalse(whole.equals(new ClassFile.UtfKey(new byte[0])));
+  }
+
+  @Test
+  void distinctBytesWithTheSameHashRemainDifferentKeys() {
+    ClassFile.UtfKey first = new ClassFile.UtfKey("Aa".getBytes(US_ASCII));
+    ClassFile.UtfKey second = new ClassFile.UtfKey("BB".getBytes(US_ASCII));
+    assertEquals(first.hashCode(), second.hashCode());
+    assertFalse(first.equals(second));
+    Map<ClassFile.UtfKey, String> lookup = new HashMap<>();
+    lookup.put(first, "first");
+    lookup.put(second, "second");
+    assertEquals("first", lookup.get(new ClassFile.UtfKey("!Aa!".getBytes(US_ASCII), 1, 2)));
+    assertEquals("second", lookup.get(new ClassFile.UtfKey("!BB!".getBytes(US_ASCII), 1, 2)));
+  }
+}

class-match/src/test/java/datadog/instrument/classmatch/ClassFileUnicodeTest.java

@@ -0,0 +1,51 @@
+package datadog.instrument.classmatch;
+
+import static datadog.instrument.classmatch.BytecodeTestSupport.classWriter;
+import static datadog.instrument.classmatch.FieldMatcher.field;
+import static datadog.instrument.classmatch.MatcherTestSupport.assertMatches;
+import static datadog.instrument.classmatch.MethodMatcher.method;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.objectweb.asm.Opcodes.ACC_ABSTRACT;
+import static org.objectweb.asm.Opcodes.ACC_PUBLIC;
+
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.objectweb.asm.ClassWriter;
+
+class ClassFileUnicodeTest {
+  @ParameterizedTest
+  @ValueSource(strings = {"Café", "Grüße", "例クラス", "\uD801\uDC00", "Asciié例\uD801\uDC00End"})
+  void modifiedUtf8NamesRoundTripAtEveryPositionAndAtAnOffset(String name) {
+    String className = "tests/" + name;
+    String descriptor = "L" + className + ";";
+    ClassWriter writer = classWriter(className);
+    writer.visitField(ACC_PUBLIC, "field" + name, descriptor, null, null).visitEnd();
+    writer
+        .visitMethod(
+            ACC_PUBLIC | ACC_ABSTRACT,
+            "method" + name,
+            "(" + descriptor + "I)" + descriptor,
+            null,
+            null)
+        .visitEnd();
+    writer.visitEnd();
+    byte[] bytes = writer.toByteArray();
+    for (int offset : new int[] {0, 29}) {
+      byte[] input = new byte[bytes.length + offset];
+      System.arraycopy(bytes, 0, input, offset, bytes.length);
+      ClassHeader header = ClassFile.header(input, offset);
+      ClassOutline outline = ClassFile.outline(input, offset);
+      assertEquals(className, header.className);
+      assertEquals(className, outline.className);
+      assertEquals("field" + name, outline.fields[0].fieldName);
+      assertEquals(descriptor, outline.fields[0].descriptor);
+      assertEquals("method" + name, outline.methods[0].methodName);
+      assertEquals("(" + descriptor + "I)" + descriptor, outline.methods[0].descriptor);
+      assertMatches(field("field" + name).type(className), outline.fields[0], true);
+      assertMatches(
+          method("method" + name).parameters(className, "int").returning(className),
+          outline.methods[0],
+          true);
+    }
+  }
+}