diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TraceApiTransformer.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TraceApiTransformer.java new file mode 100644 index 00000000000..bc5c24cb94d --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TraceApiTransformer.java @@ -0,0 +1,105 @@ +package datadog.trace.bootstrap.aot; + +import static datadog.instrument.asm.Opcodes.ACC_ABSTRACT; +import static datadog.instrument.asm.Opcodes.ACC_NATIVE; +import static datadog.instrument.asm.Opcodes.ASM9; +import static datadog.instrument.asm.Opcodes.INVOKEINTERFACE; +import static datadog.instrument.asm.Opcodes.POP2; + +import datadog.instrument.asm.ClassReader; +import datadog.instrument.asm.ClassVisitor; +import datadog.instrument.asm.ClassWriter; +import datadog.instrument.asm.MethodVisitor; +import java.lang.instrument.ClassFileTransformer; +import java.security.ProtectionDomain; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Workaround a potential AOT bug where {@link datadog.trace.api.interceptor.TraceInterceptor} is + * mistakenly restored from the system class-loader in production, even though it was visible from + * the boot class-loader during training, resulting in {@link LinkageError}s. + * + *

Any call to {@link datadog.trace.api.Tracer#addTraceInterceptor} from application code in the + * system class-loader appears to trigger this bug. The workaround is to replace these calls during + * training with opcodes that pop the tracer and argument, and push the expected return value. + * + *

Note this transformation is not persisted, so in production the original method is invoked. + */ +final class TraceApiTransformer implements ClassFileTransformer { + private static final ClassLoader SYSTEM_CLASS_LOADER = ClassLoader.getSystemClassLoader(); + + @Override + public byte[] transform( + ClassLoader loader, + String className, + Class classBeingRedefined, + ProtectionDomain pd, + byte[] bytecode) { + + // workaround only needed in the system class-loader + if (loader == SYSTEM_CLASS_LOADER) { + try { + ClassReader cr = new ClassReader(bytecode); + ClassWriter cw = new ClassWriter(cr, 0); + AtomicBoolean modified = new AtomicBoolean(); + cr.accept(new CallerPatch(cw, modified), 0); + // only return something when we've modified the bytecode + if (modified.get()) { + return cw.toByteArray(); + } + } catch (Throwable ignore) { + // skip this class + } + } + return null; // tells the JVM to keep the original bytecode + } + + /** Patches callers of {@link datadog.trace.api.Tracer#addTraceInterceptor}. */ + static final class CallerPatch extends ClassVisitor { + private final AtomicBoolean modified; + + CallerPatch(ClassVisitor cv, AtomicBoolean modified) { + super(ASM9, cv); + this.modified = modified; + } + + @Override + public MethodVisitor visitMethod( + int access, String name, String descriptor, String signature, String[] exceptions) { + MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions); + if ((access & (ACC_ABSTRACT | ACC_NATIVE)) == 0) { + return new InvokePatch(mv, modified); + } else { + return mv; // no need to patch abstract/native methods + } + } + } + + /** Removes calls to {@link datadog.trace.api.Tracer#addTraceInterceptor}. */ + static final class InvokePatch extends MethodVisitor { + private final AtomicBoolean modified; + + InvokePatch(MethodVisitor mv, AtomicBoolean modified) { + super(ASM9, mv); + this.modified = modified; + } + + @Override + public void visitMethodInsn( + int opcode, String owner, String name, String descriptor, boolean isInterface) { + if (INVOKEINTERFACE == opcode + && "datadog/trace/api/Tracer".equals(owner) + && "addTraceInterceptor".equals(name) + && "(Ldatadog/trace/api/interceptor/TraceInterceptor;)Z".equals(descriptor)) { + // pop tracer and trace interceptor argument from call stack + mv.visitInsn(POP2); + // push true return value + mv.visitLdcInsn(true); + // flag that we've modified the bytecode + modified.set(true); + } else { + super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + } + } + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TrainingAgent.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TrainingAgent.java new file mode 100644 index 00000000000..07573bf6ab6 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/aot/TrainingAgent.java @@ -0,0 +1,19 @@ +package datadog.trace.bootstrap.aot; + +import java.lang.instrument.Instrumentation; +import java.net.URL; + +/** Prepares the agent for Ahead-of-Time training. */ +public final class TrainingAgent { + public static void start( + final Object bootstrapInitTelemetry, + final Instrumentation inst, + final URL agentJarURL, + final String agentArgs) { + + // apply TraceInterceptor LinkageError workaround + inst.addTransformer(new TraceApiTransformer()); + + // don't start services, they won't be cached as they use a custom classloader + } +} diff --git a/dd-java-agent/build.gradle b/dd-java-agent/build.gradle index 7f9cfa24a0f..1b6ccfba789 100644 --- a/dd-java-agent/build.gradle +++ b/dd-java-agent/build.gradle @@ -23,12 +23,16 @@ tasks.named("processResources") { dependsOn(includedJarFileTree) } -// The special pre-check should be compiled with Java 6 to detect unsupported Java versions -// and prevent issues for users that still using them. sourceSets { + // The special pre-check must be compiled with Java 6 to detect unsupported + // Java versions and prevent issues for users that still using them. "main_java6" { java.srcDirs "${project.projectDir}/src/main/java6" } + // Additional checks that use the Java 11 API. + "main_java11" { + java.srcDirs "${project.projectDir}/src/main/java11" + } main.resources.srcDir(includedAgentDir) } @@ -36,11 +40,18 @@ def java6CompileTask = tasks.named("compileMain_java6Java") { configureCompiler(it, 8, JavaVersion.VERSION_1_6) } +def java11CompileTask = tasks.named("compileMain_java11Java") { + configureCompiler(it, 11) +} + tasks.named("compileJava") { dependsOn(java6CompileTask) + dependsOn(java11CompileTask) } dependencies { + implementation sourceSets.main_java11.output + main_java11CompileOnly libs.forbiddenapis main_java6CompileOnly libs.forbiddenapis testImplementation sourceSets.main_java6.output } @@ -249,6 +260,8 @@ includeShadowJar(traceShadowJar, 'trace', includedJarFileTree) tasks.named("shadowJar", ShadowJar) { // Include AgentPreCheck compiled with Java 6. from sourceSets.main_java6.output + // Include additional checks compiled with Java 11. + from sourceSets.main_java11.output generalShadowJarConfig(it) diff --git a/dd-java-agent/src/main/java/datadog/trace/bootstrap/AgentBootstrap.java b/dd-java-agent/src/main/java/datadog/trace/bootstrap/AgentBootstrap.java index 5e8b8567849..df1344a752d 100644 --- a/dd-java-agent/src/main/java/datadog/trace/bootstrap/AgentBootstrap.java +++ b/dd-java-agent/src/main/java/datadog/trace/bootstrap/AgentBootstrap.java @@ -142,10 +142,17 @@ private static void agentmainImpl( recordInstrumentationSource("cmd_line"); } + String agentClassName; + if (isAotTraining(agentArgs, inst)) { + agentClassName = "datadog.trace.bootstrap.aot.TrainingAgent"; + } else { + agentClassName = "datadog.trace.bootstrap.Agent"; + } + final URL agentJarURL = installAgentJar(inst); final Class agentClass; try { - agentClass = Class.forName("datadog.trace.bootstrap.Agent", true, null); + agentClass = Class.forName(agentClassName, true, null); } catch (ClassNotFoundException | LinkageError e) { throw new IllegalStateException("Unable to load DD Java Agent.", e); } @@ -422,4 +429,17 @@ private static void checkJarManifestMainClassIsThis(final URL jarUrl) throws IOE + jarUrl + "'. Make sure you don't have this .class-file anywhere, besides dd-java-agent.jar"); } + + /** Returns {@code true} if the JVM is training, i.e. writing to a CDS/AOT archive. */ + private static boolean isAotTraining(String agentArgs, Instrumentation inst) { + if (!JavaVirtualMachine.isJavaVersionAtLeast(25)) { + return false; // agent doesn't support training mode before Java 25 + } else if ("aot_training".equalsIgnoreCase(agentArgs)) { + return true; // training mode explicitly enabled via -javaagent + } else if ("false".equalsIgnoreCase(EnvironmentVariables.get("DD_DETECT_AOT_TRAINING_MODE"))) { + return false; // detection of training mode disabled via DD_DETECT_AOT_TRAINING_MODE=false + } else { + return AdvancedAgentChecks.isAotTraining(inst); // check JVM status + } + } } diff --git a/dd-java-agent/src/main/java11/datadog/trace/bootstrap/AdvancedAgentChecks.java b/dd-java-agent/src/main/java11/datadog/trace/bootstrap/AdvancedAgentChecks.java new file mode 100644 index 00000000000..a92fb94b491 --- /dev/null +++ b/dd-java-agent/src/main/java11/datadog/trace/bootstrap/AdvancedAgentChecks.java @@ -0,0 +1,39 @@ +package datadog.trace.bootstrap; + +import static java.util.Collections.emptyMap; +import static java.util.Collections.emptySet; +import static java.util.Collections.singleton; +import static java.util.Collections.singletonMap; + +import de.thetaphi.forbiddenapis.SuppressForbidden; +import java.lang.instrument.Instrumentation; +import java.lang.reflect.Method; + +/** Additional agent checks that require Java 11+. */ +public final class AdvancedAgentChecks { + + /** Returns {@code true} if the JVM is writing to a CDS/AOT archive, i.e. is in training mode. */ + @SuppressForbidden + public static boolean isAotTraining(Instrumentation inst) { + try { + Class cds = Class.forName("jdk.internal.misc.CDS"); + + // ensure the module containing CDS exports it to our unnamed module + Module cdsModule = cds.getModule(); + Module unnamedModule = AdvancedAgentChecks.class.getClassLoader().getUnnamedModule(); + inst.redefineModule( + cdsModule, + emptySet(), + singletonMap("jdk.internal.misc", singleton(unnamedModule)), + emptyMap(), + emptySet(), + emptyMap()); + + // if the JVM is writing to a CDS/AOT archive then it's in training mode + Method isDumpingArchive = cds.getMethod("isDumpingArchive"); + return (boolean) isDumpingArchive.invoke(null); + } catch (Throwable ignore) { + return false; // if we don't have access then assume we're not training + } + } +} diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 6c4408bc11d..23eeee9ba80 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -105,6 +105,14 @@ "aliases": [] } ], + "DD_DETECT_AOT_TRAINING_MODE": [ + { + "version": "A", + "type": "boolean", + "default": null, + "aliases": [] + } + ], "DD_API_KEY": [ { "version": "A",