From 5ea45caf8f284413438c5c5d9b78112f0e565d4f Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Sat, 17 Oct 2026 10:31:45 +0100 Subject: [PATCH 1/2] Add functional class injection test using startup instrumentation --- .../main/kotlin/instrument-testing.gradle.kts | 22 ++++ class-inject/build.gradle.kts | 1 + .../classinject/ClassInjectorTest.java | 105 ++++++++++++++++++ gradle.properties | 2 +- testing/build.gradle.kts | 9 ++ .../datadog/instrument/testing/TestAgent.java | 29 +++++ 6 files changed, 167 insertions(+), 1 deletion(-) create mode 100644 buildSrc/src/main/kotlin/instrument-testing.gradle.kts create mode 100644 class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java create mode 100644 testing/src/main/java/datadog/instrument/testing/TestAgent.java diff --git a/buildSrc/src/main/kotlin/instrument-testing.gradle.kts b/buildSrc/src/main/kotlin/instrument-testing.gradle.kts new file mode 100644 index 0000000..68baf4d --- /dev/null +++ b/buildSrc/src/main/kotlin/instrument-testing.gradle.kts @@ -0,0 +1,22 @@ +plugins { + id("java-common") +} + +val testAgent = configurations.create("testAgent") { + isCanBeConsumed = false + isTransitive = false +} + +dependencies { + testImplementation(project(":testing")) + testAgent(project(":testing")) +} + +class TestAgentArguments(@get:Classpath val agentClasspath: FileCollection) : CommandLineArgumentProvider { + override fun asArguments() = listOf("-javaagent:${agentClasspath.singleFile.absolutePath}") +} + +// Supply instrumentation to every test JVM, including the additional Java version tasks. +tasks.withType().configureEach { + jvmArgumentProviders.add(TestAgentArguments(testAgent)) +} diff --git a/class-inject/build.gradle.kts b/class-inject/build.gradle.kts index 6a41b36..8ee0107 100644 --- a/class-inject/build.gradle.kts +++ b/class-inject/build.gradle.kts @@ -1,6 +1,7 @@ plugins { id("java-common") id("instrument-glue") + id("instrument-testing") } dependencies { diff --git a/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java b/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java new file mode 100644 index 0000000..47eba61 --- /dev/null +++ b/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java @@ -0,0 +1,105 @@ +package datadog.instrument.classinject; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.objectweb.asm.Opcodes.*; + +import datadog.instrument.testing.TestAgent; +import java.security.ProtectionDomain; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.objectweb.asm.ClassWriter; +import org.objectweb.asm.MethodVisitor; + +class ClassInjectorTest { + + @BeforeAll + static void enableClassInjection() { + ClassInjector.enableClassInjection(TestAgent.getInstrumentation()); + } + + @Test + void injectsBootClass() throws Exception { + String name = "datadog.injected.BootSample"; + Class injectedClass = ClassInjector.injectBootClass(name, sampleClass(name)); + assertThat(injectedClass.getClassLoader()).isNull(); + assertThat(Class.forName(name, false, null)).isSameAs(injectedClass); + assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42); + + // injecting an already defined class returns the existing class + assertThat(ClassInjector.injectBootClass(name, sampleClass(name))).isSameAs(injectedClass); + } + + @Test + void injectsClassIntoGivenLoader() throws Exception { + ClassLoader loader = newCL(); + String name = "datadog.injected.LoaderSample"; + Class injectedClass = ClassInjector.injectClass(name, sampleClass(name), loader); + assertThat(injectedClass.getClassLoader()).isSameAs(loader); + assertThat(loader.loadClass(name)).isSameAs(injectedClass); + assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42); + + // injecting an already defined class returns the existing class + assertThat(ClassInjector.injectClass(name, sampleClass(name), loader)).isSameAs(injectedClass); + } + + @Test + void injectsClassWithGivenProtectionDomain() throws Exception { + ClassLoader loader = newCL(); + ProtectionDomain domain = new ProtectionDomain(null, null, loader, null); + String name = "datadog.injected.DomainSample"; + Class injectedClass = ClassInjector.injectClass(name, sampleClass(name), domain); + assertThat(injectedClass.getClassLoader()).isSameAs(loader); + assertThat(injectedClass.getProtectionDomain()).isSameAs(domain); + assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42); + } + + @Test + void injectsClassesThatReferenceEachOther() throws Exception { + ClassLoader loader = newCL(); + String caller = "datadog.injected.CallerSample"; + String callee = "datadog.injected.CalleeSample"; + Map bytecode = new LinkedHashMap<>(); + bytecode.put(caller, delegatingClass(caller, callee)); + bytecode.put(callee, sampleClass(callee)); + + List> injectedClasses = ClassInjector.injectClasses(bytecode, loader); + assertThat(injectedClasses).extracting(Class::getName).containsExactly(caller, callee); + assertThat(injectedClasses).allSatisfy(c -> assertThat(c.getClassLoader()).isSameAs(loader)); + assertThat(injectedClasses.get(0).getMethod("value").invoke(null)).isEqualTo(42); + } + + private static ClassLoader newCL() { + return new ClassLoader(null) {}; + } + + /** Generates a class with a static {@code value()} method that returns 42. */ + private static byte[] sampleClass(String name) { + ClassWriter writer = new ClassWriter(0); + writer.visit(V1_8, ACC_PUBLIC, name.replace('.', '/'), null, "java/lang/Object", null); + MethodVisitor method = writer.visitMethod(ACC_PUBLIC | ACC_STATIC, "value", "()I", null, null); + method.visitCode(); + method.visitIntInsn(BIPUSH, 42); + method.visitInsn(IRETURN); + method.visitMaxs(1, 0); + method.visitEnd(); + writer.visitEnd(); + return writer.toByteArray(); + } + + /** Generates a class with a static {@code value()} method that delegates to the target. */ + private static byte[] delegatingClass(String name, String target) { + ClassWriter writer = new ClassWriter(0); + writer.visit(V1_8, ACC_PUBLIC, name.replace('.', '/'), null, "java/lang/Object", null); + MethodVisitor method = writer.visitMethod(ACC_PUBLIC | ACC_STATIC, "value", "()I", null, null); + method.visitCode(); + method.visitMethodInsn(INVOKESTATIC, target.replace('.', '/'), "value", "()I", false); + method.visitInsn(IRETURN); + method.visitMaxs(1, 0); + method.visitEnd(); + writer.visitEnd(); + return writer.toByteArray(); + } +} diff --git a/gradle.properties b/gradle.properties index 651eaac..d6eec66 100644 --- a/gradle.properties +++ b/gradle.properties @@ -4,6 +4,6 @@ org.gradle.parallel=true org.gradle.workers.max=4 org.gradle.java.installations.auto-detect=false org.gradle.java.installations.auto-download=false -org.gradle.java.installations.fromEnv=JAVA_8_HOME,JAVA_11_HOME,JAVA_17_HOME,JAVA_21_HOME,JAVA_25_HOME +org.gradle.java.installations.fromEnv=JAVA_8_HOME,JAVA_11_HOME,JAVA_17_HOME,JAVA_21_HOME,JAVA_25_HOME,JAVA_27_HOME org.gradle.problems.report=false org.gradle.warning.mode=none diff --git a/testing/build.gradle.kts b/testing/build.gradle.kts index 2782a56..f160e21 100644 --- a/testing/build.gradle.kts +++ b/testing/build.gradle.kts @@ -1,3 +1,12 @@ plugins { id("java-common") } + +// Allow test JVMs to obtain instrumentation without dynamic agent attachment. +tasks.jar { + manifest { + attributes( + "Premain-Class" to "datadog.instrument.testing.TestAgent", + "Can-Retransform-Classes" to "true") + } +} diff --git a/testing/src/main/java/datadog/instrument/testing/TestAgent.java b/testing/src/main/java/datadog/instrument/testing/TestAgent.java new file mode 100644 index 0000000..26dd0c4 --- /dev/null +++ b/testing/src/main/java/datadog/instrument/testing/TestAgent.java @@ -0,0 +1,29 @@ +/* + * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. + * This product includes software developed at Datadog (https://www.datadoghq.com/). + * Copyright 2026-Present Datadog, Inc. + */ + +package datadog.instrument.testing; + +import java.lang.instrument.Instrumentation; + +/** Supplies instrumentation to tests when the testing jar is used as a startup agent. */ +public final class TestAgent { + private static volatile Instrumentation instrumentation; + + private TestAgent() {} + + /** Called by the JVM when the testing jar is supplied with {@code -javaagent}. */ + public static void premain(String arguments, Instrumentation inst) { + instrumentation = inst; + } + + /** Returns the instrumentation supplied at JVM startup. */ + public static Instrumentation getInstrumentation() { + if (instrumentation == null) { + throw new IllegalStateException("The testing jar must be supplied with -javaagent"); + } + return instrumentation; + } +} From 58f355d85ed29fecaa827858cf3345b449703612 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Tue, 20 Oct 2026 21:44:50 +0100 Subject: [PATCH 2/2] Expand functional class injection coverage --- .../classinject/ClassInjectorTest.java | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java b/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java index 47eba61..4ed6f02 100644 --- a/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java +++ b/class-inject/src/test/java/datadog/instrument/classinject/ClassInjectorTest.java @@ -54,6 +54,40 @@ void injectsClassWithGivenProtectionDomain() throws Exception { assertThat(injectedClass.getClassLoader()).isSameAs(loader); assertThat(injectedClass.getProtectionDomain()).isSameAs(domain); assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42); + + assertThat(ClassInjector.injectClass(name, sampleClass(name), domain)).isSameAs(injectedClass); + } + + @Test + void injectsBootClassWithGivenProtectionDomain() throws Exception { + ProtectionDomain domain = new ProtectionDomain(null, null, null, null); + String name = "datadog.injected.BootDomainSample"; + Class injectedClass = ClassInjector.injectClass(name, sampleClass(name), domain); + assertThat(injectedClass.getClassLoader()).isNull(); + assertThat(Class.forName(name, false, null)).isSameAs(injectedClass); + assertThat(injectedClass.getProtectionDomain()).isSameAs(domain); + assertThat(injectedClass.getMethod("value").invoke(null)).isEqualTo(42); + + assertThat(ClassInjector.injectClass(name, sampleClass(name), domain)).isSameAs(injectedClass); + } + + @Test + void injectsBootClassesThatReferenceEachOther() throws Exception { + String caller = "datadog.injected.BootCallerSample"; + String callee = "datadog.injected.BootCalleeSample"; + Map bytecode = new LinkedHashMap<>(); + bytecode.put(caller, delegatingClass(caller, callee)); + bytecode.put(callee, sampleClass(callee)); + + List> injectedClasses = ClassInjector.injectBootClasses(bytecode); + assertThat(injectedClasses).extracting(Class::getName).containsExactly(caller, callee); + assertThat(injectedClasses).allSatisfy(c -> assertThat(c.getClassLoader()).isNull()); + assertThat(Class.forName(caller, false, null)).isSameAs(injectedClasses.get(0)); + assertThat(Class.forName(callee, false, null)).isSameAs(injectedClasses.get(1)); + assertThat(injectedClasses.get(0).getMethod("value").invoke(null)).isEqualTo(42); + + assertThat(ClassInjector.injectBootClasses(bytecode)) + .containsExactlyElementsOf(injectedClasses); } @Test