diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformer.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformer.java new file mode 100644 index 00000000000..0c9bb065e88 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformer.java @@ -0,0 +1,14 @@ +package datadog.trace.bootstrap.instrumentation.java.lang.invoke; + +/** Transforms a generated lambda class before it is defined. */ +public interface LambdaTransformer { + /** + * @param slashClassName internal (slash-separated) name of the generated lambda class + * @param targetClass the class declaring the lambda + * @param classBytes the freshly generated lambda class bytes + * @param interfaceClassName the functional interface implemented by the lambda + * @return the transformed bytes, or {@code null}/the original bytes if unchanged + */ + byte[] transform( + String slashClassName, Class targetClass, byte[] classBytes, String interfaceClassName); +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper.java new file mode 100644 index 00000000000..82572cad015 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper.java @@ -0,0 +1,59 @@ +package datadog.trace.bootstrap.instrumentation.java.lang.invoke; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Transforms eligible lambda bytes before definition, falling back to the original on failure. */ +public final class LambdaTransformerHelper { + private static final Logger log = LoggerFactory.getLogger(LambdaTransformerHelper.class); + + // Agent transformation may itself create lambdas. + private static final ThreadLocal TRANSFORMING = new ThreadLocal<>(); + + private LambdaTransformerHelper() {} + + /** + * @param classBytes the generated lambda class bytes + * @param lambdaClassName internal (slash-separated) name of the generated lambda class + * @param targetClass the class declaring the lambda + * @param interfaceClass the functional interface implemented by the lambda + * @return possibly transformed bytes; the original bytes on any failure + */ + public static byte[] transform( + byte[] classBytes, String lambdaClassName, Class targetClass, Class interfaceClass) { + try { + if (interfaceClass == null) { + return classBytes; + } + String interfaceName = interfaceClass.getName(); + LambdaTransformer transformer = LambdaTransformerHolder.get(); + if (transformer == null) { + log.debug("Lambda {} skipped: no transformer registered", lambdaClassName); + return classBytes; + } + if (targetClass == null) { + log.debug("Lambda {} skipped: no target class", lambdaClassName); + return classBytes; + } + if (Boolean.TRUE.equals(TRANSFORMING.get())) { + log.debug("Lambda {} skipped: re-entrant transform", lambdaClassName); + return classBytes; + } + TRANSFORMING.set(Boolean.TRUE); + try { + byte[] result = + transformer.transform(lambdaClassName, targetClass, classBytes, interfaceName); + if (result == null) { + log.debug("Lambda {} not transformed", lambdaClassName); + return classBytes; + } + return result; + } finally { + TRANSFORMING.remove(); + } + } catch (Throwable e) { + log.debug("Lambda {} skipped: {}", lambdaClassName, e.toString()); + return classBytes; + } + } +} diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHolder.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHolder.java new file mode 100644 index 00000000000..9c596a659e4 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHolder.java @@ -0,0 +1,19 @@ +package datadog.trace.bootstrap.instrumentation.java.lang.invoke; + +/** + * Holds the {@link LambdaTransformer} registered by the agent installer. Lives on the bootstrap + * class path so it is reachable from instrumented {@code java.lang.invoke} code. + */ +public final class LambdaTransformerHolder { + private static volatile LambdaTransformer transformer; + + private LambdaTransformerHolder() {} + + public static void set(LambdaTransformer transformer) { + LambdaTransformerHolder.transformer = transformer; + } + + public static LambdaTransformer get() { + return transformer; + } +} diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java index db69109651a..03ca32f284f 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java @@ -2,9 +2,11 @@ import static datadog.trace.agent.tooling.ExtensionFinder.findExtensions; import static datadog.trace.agent.tooling.ExtensionLoader.loadExtensions; +import static datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnores.isIgnored; import static datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnoresMatcher.globalIgnoresMatcher; import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer; +import datadog.environment.JavaVirtualMachine; import datadog.environment.SystemProperties; import datadog.instrument.fieldinject.GlobalObjectStore; import datadog.trace.agent.tooling.bytebuddy.SharedTypePools; @@ -20,6 +22,9 @@ import datadog.trace.api.telemetry.IntegrationsCollector; import datadog.trace.bootstrap.FieldBackedContextAccessor; import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformer; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHelper; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHolder; import datadog.trace.bootstrap.instrumentation.java.module.JpmsHelper; import datadog.trace.util.AgentTaskScheduler; import de.thetaphi.forbiddenapis.SuppressForbidden; @@ -36,6 +41,7 @@ import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.TimeUnit; import java.util.function.BooleanSupplier; +import java.util.function.Function; import net.bytebuddy.ByteBuddy; import net.bytebuddy.agent.builder.AgentBuilder; import net.bytebuddy.description.type.TypeDescription; @@ -166,6 +172,15 @@ public static ClassFileTransformer installBytebuddyAgent( // .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED) .ignore(globalIgnoresMatcher(skipAdditionalLibraryMatcher)); + boolean lambdaTransformationEnabled = + !Platform.isNativeImageBuilder() + && InstrumenterConfig.get() + .isIntegrationEnabled(Collections.singleton("lambda"), false); + if (lambdaTransformationEnabled) { + // The injected metafactory call needs java.base to read the bootstrap helper's module. + agentBuilder = agentBuilder.assureReadEdgeTo(inst, LambdaTransformerHelper.class); + } + if (DEBUG) { agentBuilder = agentBuilder @@ -216,7 +231,8 @@ public static ClassFileTransformer installBytebuddyAgent( } CombiningTransformerBuilder transformerBuilder = - new CombiningTransformerBuilder(agentBuilder, instrumenterIndex, enabledSystems); + new CombiningTransformerBuilder( + agentBuilder, instrumenterIndex, enabledSystems, lambdaTransformationEnabled); int installedCount = 0; for (InstrumenterModule module : instrumenterModules) { @@ -265,12 +281,96 @@ public void applied(Iterable instrumentationNames) { InstrumenterState.resetDefaultState(); try { - return transformerBuilder.installOn(inst); + ClassFileTransformer classFileTransformer = transformerBuilder.installOn(inst); + registerLambdaTransformer( + lambdaTransformationEnabled, classFileTransformer, transformerBuilder.lambdaInterfaces()); + return classFileTransformer; } finally { SharedTypePools.endInstall(); } } + /** Registers the installed class-file transformer for generated lambdas. */ + static void registerLambdaTransformer( + final boolean enabled, + final ClassFileTransformer classFileTransformer, + final String[] lambdaInterfaces) { + if (!enabled) { + // Agent installation can be repeated in tests and embedded environments. + LambdaTransformerHolder.set(null); + return; + } + LambdaTransformer transformer = + lambdaInterfaces.length == 0 ? null : newLambdaTransformer(classFileTransformer); + LambdaTransformerHolder.set(filterLambdaTransformer(transformer, lambdaInterfaces)); + } + + static LambdaTransformer filterLambdaTransformer( + final LambdaTransformer transformer, final String[] lambdaInterfaces) { + if (transformer == null) { + return null; + } + return (className, targetClass, classBytes, interfaceName) -> { + for (String enabledInterface : lambdaInterfaces) { + if (enabledInterface.equals(interfaceName)) { + // Apply the system-level name filter before entering the full transformer pipeline. + if (isIgnored(targetClass.getName(), true)) { + return null; + } + return transformer.transform(className, targetClass, classBytes, interfaceName); + } + } + return null; + }; + } + + /** + * Java 9+ requires the module-aware transformer for injected read edges. Failure must disable + * lambda transformation rather than fall back to the module-less overload. + */ + @SuppressWarnings("unchecked") + private static LambdaTransformer newLambdaTransformer( + final ClassFileTransformer classFileTransformer) { + if (JavaVirtualMachine.isJavaVersionAtLeast(9)) { + try { + Function factory = + (Function) + Instrumenter.class + .getClassLoader() + .loadClass("datadog.trace.agent.tooling.bytebuddy.DDJava9LambdaTransformer") + .getField("FACTORY") + .get(null); + return factory.apply(classFileTransformer); + } catch (Throwable e) { + log.debug("Problem loading Java 9 lambda transformer, disabling lambda transformation", e); + return null; + } + } + // Avoid invoking the instrumented metafactory while installing its transformer. + return new LambdaTransformer() { + @Override + public byte[] transform( + String slashClassName, + Class targetClass, + byte[] classBytes, + String interfaceClassName) { + TypePoolFacade.beginLambdaTransform(interfaceClassName); + try { + return classFileTransformer.transform( + targetClass.getClassLoader(), + slashClassName, + null, + targetClass.getProtectionDomain(), + classBytes); + } catch (Throwable ignored) { + return null; + } finally { + TypePoolFacade.endLambdaTransform(); + } + } + }; + } + /** Returns an iterable that combines the original sequence with any discovered extensions. */ private static Iterable withExtensions(Iterable initial) { String extensionsPath = InstrumenterConfig.get().getTraceExtensionsPath(); diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningMatcher.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningMatcher.java index 5ac7ea2e67c..bb66f312c71 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningMatcher.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningMatcher.java @@ -2,9 +2,11 @@ import static datadog.trace.api.config.TraceInstrumentationConfig.EXPERIMENTAL_DEFER_INTEGRATIONS_UNTIL; import static datadog.trace.util.AgentThreadFactory.AgentThread.RETRANSFORMER; +import static java.util.Collections.unmodifiableMap; import datadog.trace.agent.tooling.bytebuddy.matcher.CustomExcludes; import datadog.trace.agent.tooling.bytebuddy.matcher.ProxyClassIgnores; +import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade; import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.time.TimeUtils; import datadog.trace.util.AgentTaskScheduler; @@ -12,8 +14,10 @@ import java.security.ProtectionDomain; import java.util.ArrayList; import java.util.BitSet; +import java.util.HashMap; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.concurrent.TimeUnit; import net.bytebuddy.agent.builder.AgentBuilder; @@ -45,13 +49,22 @@ final class CombiningMatcher implements AgentBuilder.RawMatcher { private final BitSet knownTypesMask; private final MatchRecorder[] matchers; + private final Map lambdaMatchers; private volatile boolean deferring; CombiningMatcher( - Instrumentation instrumentation, BitSet knownTypesMask, List matchers) { + Instrumentation instrumentation, + BitSet knownTypesMask, + List matchers, + Map> lambdaMatchers) { this.knownTypesMask = knownTypesMask; this.matchers = matchers.toArray(new MatchRecorder[0]); + Map lambdaMatchersByInterface = new HashMap<>(); + lambdaMatchers.forEach( + (name, recorders) -> + lambdaMatchersByInterface.put(name, recorders.toArray(new LambdaMatchRecorder[0]))); + this.lambdaMatchers = unmodifiableMap(lambdaMatchersByInterface); if (DEFER_MATCHING) { scheduleResumeMatching(instrumentation, InstrumenterConfig.get().deferIntegrationsUntil()); @@ -66,8 +79,15 @@ public boolean matches( Class classBeingRedefined, ProtectionDomain pd) { + String lambdaInterface = lambdaMatchers.isEmpty() ? null : TypePoolFacade.lambdaInterface(); + // check initial requests to see if we should defer matching until retransformation - if (DEFER_MATCHING && null == classBeingRedefined && deferring && isDeferred(classLoader)) { + // Generated lambda classes cannot be retransformed because they are not yet defined here. + if (null == lambdaInterface + && DEFER_MATCHING + && null == classBeingRedefined + && deferring + && isDeferred(classLoader)) { return false; } @@ -75,6 +95,27 @@ public boolean matches( ids.clear(); long fromTick = InstrumenterMetrics.tick(); + if (null != lambdaInterface) { + LambdaMatchRecorder[] recorders = lambdaMatchers.get(lambdaInterface); + if (null != recorders) { + for (LambdaMatchRecorder recorder : recorders) { + try { + recorder.record(target, classLoader, ids); + } catch (Throwable e) { + if (log.isDebugEnabled()) { + log.debug( + "Lambda instrumentation matcher unexpected exception - interface={} matcher={}", + lambdaInterface, + recorder.describe(), + e); + } + } + } + } + InstrumenterMetrics.matchType(fromTick); + return !ids.isEmpty(); + } + knownTypesIndex.apply(target.getName(), knownTypesMask, ids); if (ids.isEmpty()) { InstrumenterMetrics.knownTypeMiss(fromTick); diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java index 5a7e4b9c3df..4540e1f34f8 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java @@ -59,10 +59,15 @@ public final class CombiningTransformerBuilder private final Map, ElementMatcher> contextStoreInjection = new HashMap<>(); + private final Map> lambdaMatchers = new HashMap<>(); + private final Map, List> + lambdaContextStoreInjection = new HashMap<>(); + private final AgentBuilder agentBuilder; private final InstrumenterIndex instrumenterIndex; private final int knownTransformationCount; private final Set enabledSystems; + private final boolean lambdaTransformationEnabled; private final List matchers = new ArrayList<>(); private final BitSet knownTypesMask; @@ -88,7 +93,8 @@ public final class CombiningTransformerBuilder public CombiningTransformerBuilder( AgentBuilder agentBuilder, InstrumenterIndex instrumenterIndex, - Set enabledSystems) { + Set enabledSystems, + boolean lambdaTransformationEnabled) { this.agentBuilder = agentBuilder; this.instrumenterIndex = instrumenterIndex; int knownInstrumentationCount = instrumenterIndex.instrumentationCount(); @@ -98,6 +104,7 @@ public CombiningTransformerBuilder( this.nextRuntimeInstrumentationId = knownInstrumentationCount; this.nextRuntimeTransformationId = knownTransformationCount; this.enabledSystems = enabledSystems; + this.lambdaTransformationEnabled = lambdaTransformationEnabled; } /** Builds matchers and transformers for an instrumentation module and its members. */ @@ -162,9 +169,41 @@ private void buildTypeInstrumentation(Instrumenter member) { } buildTypeMatcher(member, transformationId); + buildLambdaMatcher(member, transformationId); buildTypeAdvice(member, transformationId); } + private void buildLambdaMatcher(Instrumenter member, int transformationId) { + if (!lambdaTransformationEnabled || !(member instanceof Instrumenter.ForLambda)) { + return; + } + + Instrumenter.ForLambda lambdaInstrumenter = (Instrumenter.ForLambda) member; + ElementMatcher typeMatcher = lambdaInstrumenter.lambdaMatcher(); + if (member instanceof Instrumenter.WithTypeStructure) { + typeMatcher = + new ElementMatcher.Junction.Conjunction<>( + typeMatcher, ((Instrumenter.WithTypeStructure) member).structureMatcher()); + } + + LambdaMatchRecorder recorder = + new LambdaMatchRecorder( + transformationId, typeMatcher, requireBoth(classLoaderMatcher, muzzle)); + lambdaMatchers + .computeIfAbsent(lambdaInstrumenter.lambdaInterface(), ignored -> new ArrayList<>()) + .add(recorder); + + for (Map.Entry store : contextStore.entrySet()) { + lambdaContextStoreInjection + .computeIfAbsent(store, ignored -> new ArrayList<>()) + .add(recorder); + } + } + + String[] lambdaInterfaces() { + return lambdaMatchers.keySet().toArray(new String[0]); + } + private void buildTypeMatcher(Instrumenter member, int transformationId) { if (member instanceof Instrumenter.ForSingleType) { @@ -291,7 +330,7 @@ public ClassFileTransformer installOn(Instrumentation instrumentation) { } return agentBuilder - .type(new CombiningMatcher(instrumentation, knownTypesMask, matchers)) + .type(new CombiningMatcher(instrumentation, knownTypesMask, matchers, lambdaMatchers)) .and(NOT_DECORATOR_MATCHER) .transform(defaultTransformers()) .transform(new SplittingTransformer(transformers)) @@ -360,6 +399,19 @@ private void applyContextStoreInjection( matchers.add(new MatchRecorder.ForContextStore(transformationId, activation, contextMatcher)); transformers[transformationId] = new AdviceStack(new VisitingTransformer(contextAdvice)); + + List lambdaRecorders = lambdaContextStoreInjection.get(contextStore); + if (null != lambdaRecorders) { + // Lambda transformation happens before definition, so its field injector can be selected + // along with the instrumentation that requested this context store. Keep the normal + // assignability check because a module may declare stores for unrelated context keys. + for (LambdaMatchRecorder recorder : lambdaRecorders) { + recorder.addTransformation( + transformationId, + target -> contextMatcher.matches(target, null), + contextMatcher.describe()); + } + } } static final class VisitingTransformer implements AgentBuilder.Transformer { diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/LambdaMatchRecorder.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/LambdaMatchRecorder.java new file mode 100644 index 00000000000..bf879180750 --- /dev/null +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/LambdaMatchRecorder.java @@ -0,0 +1,71 @@ +package datadog.trace.agent.tooling; + +import java.util.ArrayList; +import java.util.BitSet; +import java.util.List; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Selects the transformations explicitly registered for a generated lambda interface. */ +final class LambdaMatchRecorder { + private static final Logger log = LoggerFactory.getLogger(LambdaMatchRecorder.class); + + private final ElementMatcher typeMatcher; + private final ElementMatcher classLoaderMatcher; + private final BitSet transformationIds = new BitSet(); + private final List conditionalTransformations = new ArrayList<>(); + + LambdaMatchRecorder( + int transformationId, + ElementMatcher typeMatcher, + ElementMatcher classLoaderMatcher) { + this.typeMatcher = typeMatcher; + this.classLoaderMatcher = classLoaderMatcher; + transformationIds.set(transformationId); + } + + void addTransformation( + int transformationId, ElementMatcher matcher, String matcherDescription) { + conditionalTransformations.add( + new ConditionalTransformation(transformationId, matcher, matcherDescription)); + } + + void record(TypeDescription type, ClassLoader classLoader, BitSet matches) { + if (classLoaderMatcher.matches(classLoader) && typeMatcher.matches(type)) { + matches.or(transformationIds); + for (ConditionalTransformation transformation : conditionalTransformations) { + try { + if (transformation.matcher.matches(type)) { + matches.set(transformation.id); + } + } catch (Throwable e) { + if (log.isDebugEnabled()) { + log.debug( + "Lambda transformation matcher unexpected exception - {}", + transformation.description, + e); + } + } + } + } + } + + String describe() { + return typeMatcher.toString(); + } + + private static final class ConditionalTransformation { + private final int id; + private final ElementMatcher matcher; + private final String description; + + private ConditionalTransformation( + int id, ElementMatcher matcher, String description) { + this.id = id; + this.matcher = matcher; + this.description = description; + } + } +} diff --git a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/bytebuddy/DDClassFileTransformer.java b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/bytebuddy/DDClassFileTransformer.java index c42cac42336..520c45df8b2 100644 --- a/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/bytebuddy/DDClassFileTransformer.java +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/bytebuddy/DDClassFileTransformer.java @@ -37,7 +37,7 @@ public byte[] transform( return classFileTransformer.transform( classLoader, internalClassName, classBeingRedefined, protectionDomain, classFileBuffer); } finally { - SharedTypePools.endTransform(); + SharedTypePools.endTransform(classFileBuffer); } } } diff --git a/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9ClassFileTransformer.java b/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9ClassFileTransformer.java index 7b7f14afc10..b478b4759b9 100644 --- a/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9ClassFileTransformer.java +++ b/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9ClassFileTransformer.java @@ -38,7 +38,7 @@ public byte[] transform( return classFileTransformer.transform( classLoader, internalClassName, classBeingRedefined, protectionDomain, classFileBuffer); } finally { - SharedTypePools.endTransform(); + SharedTypePools.endTransform(classFileBuffer); } } @@ -65,7 +65,7 @@ public byte[] transform( protectionDomain, classFileBuffer); } finally { - SharedTypePools.endTransform(); + SharedTypePools.endTransform(classFileBuffer); } } } diff --git a/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9LambdaTransformer.java b/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9LambdaTransformer.java new file mode 100644 index 00000000000..d385d385589 --- /dev/null +++ b/dd-java-agent/agent-installer/src/main/java11/datadog/trace/agent/tooling/bytebuddy/DDJava9LambdaTransformer.java @@ -0,0 +1,44 @@ +package datadog.trace.agent.tooling.bytebuddy; + +import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformer; +import java.lang.instrument.ClassFileTransformer; +import java.util.function.Function; + +/** Routes generated lambdas through the module-aware Java 9+ transformer overload. */ +public final class DDJava9LambdaTransformer implements LambdaTransformer { + + /** Read reflectively by the agent installer, which cannot name {@link Module} itself. */ + public static final Function FACTORY = + new Function() { + @Override + public LambdaTransformer apply(ClassFileTransformer classFileTransformer) { + return new DDJava9LambdaTransformer(classFileTransformer); + } + }; + + private final ClassFileTransformer classFileTransformer; + + public DDJava9LambdaTransformer(ClassFileTransformer classFileTransformer) { + this.classFileTransformer = classFileTransformer; + } + + @Override + public byte[] transform( + String slashClassName, Class targetClass, byte[] classBytes, String interfaceClassName) { + TypePoolFacade.beginLambdaTransform(interfaceClassName); + try { + return classFileTransformer.transform( + targetClass.getModule(), + targetClass.getClassLoader(), + slashClassName, + null, + targetClass.getProtectionDomain(), + classBytes); + } catch (Throwable ignored) { + return null; + } finally { + TypePoolFacade.endLambdaTransform(); + } + } +} diff --git a/dd-java-agent/agent-installer/src/test/java/datadog/application/LambdaTarget.java b/dd-java-agent/agent-installer/src/test/java/datadog/application/LambdaTarget.java new file mode 100644 index 00000000000..862f71bfd8c --- /dev/null +++ b/dd-java-agent/agent-installer/src/test/java/datadog/application/LambdaTarget.java @@ -0,0 +1,6 @@ +package datadog.application; + +/** Application-owned class under the top-level Datadog namespace. */ +public final class LambdaTarget { + private LambdaTarget() {} +} diff --git a/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/CombiningMatcherTest.java b/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/CombiningMatcherTest.java new file mode 100644 index 00000000000..2bb40d950d3 --- /dev/null +++ b/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/CombiningMatcherTest.java @@ -0,0 +1,91 @@ +package datadog.trace.agent.tooling; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static java.util.Arrays.asList; +import static java.util.Collections.emptyList; +import static java.util.Collections.singletonList; +import static java.util.Collections.singletonMap; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade; +import java.util.BitSet; +import java.util.List; +import java.util.Map; +import net.bytebuddy.description.type.TypeDescription; +import org.junit.jupiter.api.Test; + +class CombiningMatcherTest { + + @Test + void selectsOnlyTransformationsRegisteredForLambdaInterface() { + TypeDescription target = new TypeDescription.ForLoadedType(getClass()); + ClassLoader classLoader = getClass().getClassLoader(); + LambdaMatchRecorder recorder = + new LambdaMatchRecorder(3, named(getClass().getName()), loader -> loader == classLoader); + recorder.addTransformation(7, named(getClass().getName()), "matching context store"); + recorder.addTransformation(11, named(String.class.getName()), "unrelated context store"); + recorder.addTransformation( + 13, + ignored -> { + throw new IllegalStateException("broken conditional matcher"); + }, + "broken context store"); + recorder.addTransformation(17, named(getClass().getName()), "later context store"); + + Map> lambdaMatchers = + singletonMap(Runnable.class.getName(), singletonList(recorder)); + CombiningMatcher matcher = + new CombiningMatcher(null, new BitSet(), emptyList(), lambdaMatchers); + + TypePoolFacade.beginLambdaTransform(Runnable.class.getName()); + try { + assertTrue(matcher.matches(target, classLoader, null, null, null)); + assertTrue(CombiningMatcher.recordedMatches.get().get(3)); + assertTrue(CombiningMatcher.recordedMatches.get().get(7)); + assertFalse(CombiningMatcher.recordedMatches.get().get(11)); + assertFalse(CombiningMatcher.recordedMatches.get().get(13)); + assertTrue(CombiningMatcher.recordedMatches.get().get(17)); + } finally { + TypePoolFacade.endLambdaTransform(); + } + + TypePoolFacade.beginLambdaTransform(java.util.function.Supplier.class.getName()); + try { + assertFalse(matcher.matches(target, classLoader, null, null, null)); + } finally { + TypePoolFacade.endLambdaTransform(); + } + } + + @Test + void isolatesFailuresBetweenLambdaMatchers() { + TypeDescription target = new TypeDescription.ForLoadedType(getClass()); + ClassLoader classLoader = getClass().getClassLoader(); + LambdaMatchRecorder broken = + new LambdaMatchRecorder( + 3, + ignored -> { + throw new IllegalStateException("broken matcher"); + }, + ignored -> true); + LambdaMatchRecorder working = + new LambdaMatchRecorder(7, named(getClass().getName()), ignored -> true); + + CombiningMatcher matcher = + new CombiningMatcher( + null, + new BitSet(), + emptyList(), + singletonMap(Runnable.class.getName(), asList(broken, working))); + + TypePoolFacade.beginLambdaTransform(Runnable.class.getName()); + try { + assertTrue(matcher.matches(target, classLoader, null, null, null)); + assertFalse(CombiningMatcher.recordedMatches.get().get(3)); + assertTrue(CombiningMatcher.recordedMatches.get().get(7)); + } finally { + TypePoolFacade.endLambdaTransform(); + } + } +} diff --git a/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/LambdaTransformerTest.java b/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/LambdaTransformerTest.java new file mode 100644 index 00000000000..66c2ac7cf04 --- /dev/null +++ b/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/LambdaTransformerTest.java @@ -0,0 +1,69 @@ +package datadog.trace.agent.tooling; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; + +import datadog.application.LambdaTarget; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformer; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHolder; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.Test; + +class LambdaTransformerTest { + + @Test + void disabledInstallationClearsPreviousTransformer() { + LambdaTransformer previous = (className, targetClass, classBytes, interfaceName) -> classBytes; + LambdaTransformerHolder.set(previous); + try { + AgentInstaller.registerLambdaTransformer(false, null, new String[0]); + + assertNull(LambdaTransformerHolder.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + @Test + void onlyTransformsEnabledInterfaces() { + byte[] transformedBytes = new byte[0]; + AtomicBoolean transformed = new AtomicBoolean(); + LambdaTransformer delegate = + (className, targetClass, classBytes, interfaceName) -> { + transformed.set(true); + return transformedBytes; + }; + LambdaTransformer transformer = + AgentInstaller.filterLambdaTransformer(delegate, new String[] {Runnable.class.getName()}); + + assertSame( + transformedBytes, + transformer.transform( + "test/Lambda", LambdaTarget.class, new byte[0], Runnable.class.getName())); + transformed.set(false); + assertNull( + transformer.transform("test/Lambda", LambdaTarget.class, new byte[0], "other.Interface")); + assertFalse(transformed.get()); + } + + @Test + void skipsGloballyIgnoredTargets() { + AtomicBoolean transformed = new AtomicBoolean(); + LambdaTransformer delegate = + (className, targetClass, classBytes, interfaceName) -> { + transformed.set(true); + return classBytes; + }; + LambdaTransformer transformer = + AgentInstaller.filterLambdaTransformer(delegate, new String[] {Runnable.class.getName()}); + + assertNull( + transformer.transform( + "datadog/trace/Lambda", + LambdaTransformerTest.class, + new byte[0], + Runnable.class.getName())); + assertFalse(transformed.get()); + } +} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java index b146121d273..1c8174d6fa5 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java @@ -36,6 +36,21 @@ interface ForTypeHierarchy { /** Instrumentation that transforms types on the bootstrap class-path. */ interface ForBootstrap {} + /** + * Instrumentation selected directly for an exact lambda interface. + * + *

Lambda transformation also requires the opt-in {@code trace.lambda.enabled} setting. An + * enabled {@code ForLambda} instrumentation registers its interface but does not enable the + * lambda metafactory instrumentation by itself. + */ + interface ForLambda { + /** Exact functional interface name reported by the lambda metafactory. */ + String lambdaInterface(); + + /** Additional checks not already guaranteed by the lambda metafactory. */ + ElementMatcher lambdaMatcher(); + } + /** * Instrumentation that matches a series of types configured at runtime. This is used for last * minute additions in the field such as testing a new JDBC driver that is not yet in the allowed diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/SharedTypePools.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/SharedTypePools.java index f6029fc8f97..b4b9b511005 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/SharedTypePools.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/SharedTypePools.java @@ -32,9 +32,9 @@ public static void endInstall() { SUPPLIER.endInstall(); } - /** Hints that the javaagent has finished calling {@link ClassFileTransformer#transform}. */ - public static void endTransform() { - SUPPLIER.endTransform(); + /** Ends the transform associated with this callback's input bytes, if it reached resolution. */ + public static void endTransform(byte[] classFileBuffer) { + SUPPLIER.endTransform(classFileBuffer); } public static void clear() { @@ -61,8 +61,8 @@ public interface Supplier { /** Hints that the javaagent has finished installing as a {@link ClassFileTransformer}. */ void endInstall(); - /** Hints that the javaagent has finished calling {@link ClassFileTransformer#transform}. */ - void endTransform(); + /** Ends the transform associated with this callback's input bytes, if it reached resolution. */ + void endTransform(byte[] classFileBuffer); void clear(); } @@ -82,7 +82,7 @@ public void annotationOfInterest(String name) {} public void endInstall() {} @Override - public void endTransform() {} + public void endTransform(byte[] classFileBuffer) {} @Override public void clear() {} diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java index 6c8d01152ca..e6454169f2d 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/memoize/Memoizer.java @@ -155,7 +155,7 @@ static final class MemoizingMatcher @Override protected boolean doMatch(TypeDescription target) { String targetName = target.getName(); - if (noMatchFilter.contains(targetName) + if ((isCacheable(target) && noMatchFilter.contains(targetName)) || "java.lang.Object".equals(targetName) || target.isPrimitive()) { return false; @@ -166,7 +166,7 @@ protected boolean doMatch(TypeDescription target) { } static BitSet memoizeHierarchy(TypeDescription type, Map localMemos) { - if (noMatchFilter.contains(type.getName())) { + if (isCacheable(type) && noMatchFilter.contains(type.getName())) { return NO_MATCH; } else { return doMemoize(type, localMemos); @@ -176,13 +176,14 @@ static BitSet memoizeHierarchy(TypeDescription type, Map localMe static BitSet doMemoize(TypeDescription type, Map localMemos) { String name = type.getName(); + boolean cacheable = isCacheable(type); BitSet memo = localMemos.get(name); if (null != memo) { return memo; // short-circuit circular references } long fromTick = InstrumenterMetrics.tick(); - SharedTypeInfo sharedMemo = memos.find(name); + SharedTypeInfo sharedMemo = cacheable ? memos.find(name) : null; if (null != sharedMemo) { if (namesAreUnique || name.startsWith("java.") || sameOrigin(type, sharedMemo)) { InstrumenterMetrics.reuseTypeMemo(fromTick); @@ -233,16 +234,20 @@ static BitSet doMemoize(TypeDescription type, Map localMemos) { // update no-match filter if there's no interesting matches and result is complete if (memo.nextSetBit(INTERNAL_MATCHERS) < 0 && !memo.get(isPartial.matcherId)) { - noMatchFilter.add(name); + if (cacheable) { + noMatchFilter.add(name); + } return NO_MATCH; } - // otherwise share result for this location (other locations may have different results) - if (namesAreUnique || name.startsWith("java.") || !(type instanceof WithLocation)) { - memos.share(name, 0, null, memo); - } else { - WithLocation origin = (WithLocation) type; - memos.share(name, origin.getClassLoaderId(), origin.getClassFile(), memo); + if (cacheable) { + // otherwise share result for this location (other locations may have different results) + if (namesAreUnique || name.startsWith("java.") || !(type instanceof WithLocation)) { + memos.share(name, 0, null, memo); + } else { + WithLocation origin = (WithLocation) type; + memos.share(name, origin.getClassLoaderId(), origin.getClassFile(), memo); + } } return memo; @@ -253,6 +258,10 @@ static boolean potentialMatch(String name) { return !noMatchFilter.contains(name); } + private static boolean isCacheable(TypeDescription type) { + return !(type instanceof WithLocation) || ((WithLocation) type).isCacheable(); + } + private static boolean sameOrigin(TypeDescription type, SharedTypeInfo sharedMemo) { return !(type instanceof WithLocation) || sharedMemo.sameClassLoader(((WithLocation) type).getClassLoaderId()) diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java index 51650b3d2b2..029f43cacba 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactory.java @@ -100,6 +100,8 @@ final class TypeFactory { boolean createOutlines = OUTLINING_ENABLED; + String lambdaInterface; + ClassLoader originalClassLoader; ClassLoader currentClassLoader; @@ -145,10 +147,11 @@ static void clear() { /** * New transform request; begins with type matching that only requires outline descriptions. * - *

Byte-buddy's circularity lock makes sure we won't have any nested transform calls, but we - * may be asked to transform support types while deciding which loaded types need re-transforming - * when first installing the agent. If that happens then we need to remember the original context - * used for matching and restore it afterwards. + *

Byte-buddy's circularity lock prevents nested callbacks from reaching this step, but those + * callbacks still pass through the outer transformer wrapper. We may also be asked to transform + * support types while deciding which loaded types need re-transforming when first installing the + * agent. If that happens then we need to remember the original context used for matching and + * restore it afterwards. */ void beginTransform(String name, byte[] bytecode) { targetName = name; @@ -159,6 +162,18 @@ void beginTransform(String name, byte[] bytecode) { } } + void beginLambdaTransform(String interfaceClassName) { + lambdaInterface = interfaceClassName; + } + + void endLambdaTransform() { + lambdaInterface = null; + } + + String lambdaInterface() { + return lambdaInterface; + } + /** Once matching is complete we need full descriptions for the actual transformation. */ void enableFullDescriptions() { createOutlines = false; @@ -171,7 +186,14 @@ boolean disableFullDescriptions() { return wasEnabled; } - /** Cleans-up local caches to minimise memory use once we're done with the type-factory. */ + /** Cleans up local caches if this callback owns the active transformation. */ + void endTransform(byte[] classFileBuffer) { + if (targetBytecode == classFileBuffer) { + endTransform(); + } + } + + /** Cleans up local caches to minimise memory use once we're done with the type-factory. */ void endTransform() { if (null == targetName) { return; // transformation didn't reach resolve step @@ -231,6 +253,10 @@ private TypeDescription deferTypeResolution(String name) { return deferredTypes.computeIfAbsent(name, deferType); } + private boolean isLambdaTarget(String name) { + return null != lambdaInterface && name.equals(targetName); + } + /** Attempts to resolve the named type using the current context. */ TypeDescription resolveType(LazyType request) { if (null != classFileLocator) { @@ -257,9 +283,10 @@ private TypeDescription lookupType( int classLoaderId = request.getClassLoaderId(); boolean isOutline = typeParser == outlineTypeParser; long fromTick = InstrumenterMetrics.tick(); + // Hidden lambda names may later be reused by an ordinary class definition. + boolean cacheable = !isLambdaTarget(name); - // existing type description from same classloader? - SharedTypeInfo sharedType = types.find(name); + SharedTypeInfo sharedType = cacheable ? types.find(name) : null; if (null != sharedType && (name.startsWith("java.") || sharedType.sameClassLoader(classLoaderId))) { InstrumenterMetrics.reuseTypeDescription(fromTick, isOutline); @@ -286,14 +313,16 @@ private TypeDescription lookupType( InstrumenterMetrics.buildTypeDescription(fromTick, isOutline); - if (MEMOIZING_ENABLED && null != type) { + if (cacheable && MEMOIZING_ENABLED && null != type) { if (type.isPublic()) { isPublicFilter.add(name); } } - // share result, whether we found it or not - types.share(name, classLoaderId, classFile, type); + if (cacheable) { + // share result, whether we found it or not + types.share(name, classLoaderId, classFile, type); + } return type; } @@ -368,6 +397,11 @@ public byte[] getBytecode() { return null; } + @Override + public boolean isCacheable() { + return !isLambdaTarget(name); + } + private ClassFileLocator.Resolution locateClassFile() { if (name.equals(targetName)) { return new ClassFileLocator.Resolution.Explicit(targetBytecode); @@ -396,7 +430,7 @@ public TypeList.Generic getInterfaces() { @Override public boolean isPublic() { - return isPublicFilter.contains(name) || super.isPublic(); + return (isCacheable() && isPublicFilter.contains(name)) || super.isPublic(); } private TypeDescription outline() { diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java index b5b2e3b39a8..228f0895e25 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/TypePoolFacade.java @@ -49,6 +49,18 @@ public static void beginTransform(String name, byte[] bytecode) { typeFactory.get().beginTransform(name, bytecode); } + public static void beginLambdaTransform(String interfaceClassName) { + typeFactory.get().beginLambdaTransform(interfaceClassName); + } + + public static void endLambdaTransform() { + typeFactory.get().endLambdaTransform(); + } + + public static String lambdaInterface() { + return typeFactory.get().lambdaInterface(); + } + /** Switch to full descriptions, needed for the actual class transformation. */ public static void enableFullDescriptions() { typeFactory.get().enableFullDescriptions(); @@ -60,8 +72,8 @@ public static boolean disableFullDescriptions() { } @Override - public void endTransform() { - typeFactory.get().endTransform(); + public void endTransform(byte[] classFileBuffer) { + typeFactory.get().endTransform(classFileBuffer); } @Override diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/WithLocation.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/WithLocation.java index 55bde94f0ba..3e1aa6f329b 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/WithLocation.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/outline/WithLocation.java @@ -9,4 +9,9 @@ public interface WithLocation { URL getClassFile(); byte[] getBytecode(); + + /** Whether details for this type may be shared by symbolic name. */ + default boolean isCacheable() { + return true; + } } diff --git a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie index 36cedaf6081..939b8181417 100644 --- a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie +++ b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie @@ -57,6 +57,8 @@ 0 java.lang.Runtime # allow context tracking for VirtualThread 0 java.lang.VirtualThread +# allow opt-in instrumentation of the lambda metafactory +0 java.lang.invoke.InnerClassLambdaMetafactory 0 java.net.http.* 0 java.net.HttpURLConnection 0 java.net.InetAddress diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTest.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTest.java new file mode 100644 index 00000000000..dbb3bdec592 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTest.java @@ -0,0 +1,150 @@ +package datadog.trace.agent.tooling.bytebuddy.outline; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +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 datadog.trace.agent.tooling.bytebuddy.memoize.MemoizedMatchers; +import java.util.concurrent.Callable; +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.description.modifier.Visibility; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; +import org.junit.jupiter.api.Test; + +class TypeFactoryTest { + @Test + void reusesCachedDescriptionForRegularTransformationTarget() { + String name = getClass().getName() + "$RegularTarget"; + + assertEquals( + Runnable.class.getName(), resolveInterface(name, bytes(name, Runnable.class), false)); + assertEquals( + Runnable.class.getName(), resolveInterface(name, bytes(name, Callable.class), false)); + } + + @Test + void rebuildsLambdaTransformationTargetFromSuppliedBytes() { + String name = getClass().getName() + "$LambdaTarget"; + + assertEquals( + Runnable.class.getName(), resolveInterface(name, bytes(name, Runnable.class), false)); + assertEquals( + Callable.class.getName(), resolveInterface(name, bytes(name, Callable.class), true)); + } + + @Test + void doesNotCacheLambdaTransformationTarget() { + String name = getClass().getName() + "$UncachedLambdaTarget"; + + assertEquals( + Runnable.class.getName(), resolveInterface(name, bytes(name, Runnable.class), true)); + assertEquals( + Callable.class.getName(), resolveInterface(name, bytes(name, Callable.class), false)); + } + + @Test + void doesNotMemoizeLambdaMatcherResults() { + ElementMatcher implementsRunnable = + new MemoizedMatchers().hasInterface(named(Runnable.class.getName())); + + String lambdaFirstName = getClass().getName() + "$LambdaMatcherFirst"; + assertTrue( + matches(lambdaFirstName, bytes(lambdaFirstName, Runnable.class), true, implementsRunnable)); + assertFalse( + matches( + lambdaFirstName, bytes(lambdaFirstName, Callable.class), false, implementsRunnable)); + + String lambdaSecondName = getClass().getName() + "$LambdaMatcherSecond"; + assertFalse( + matches( + lambdaSecondName, bytes(lambdaSecondName, Callable.class), false, implementsRunnable)); + assertTrue( + matches( + lambdaSecondName, bytes(lambdaSecondName, Runnable.class), true, implementsRunnable)); + } + + @Test + void doesNotReuseCachedVisibilityForLambda() { + ElementMatcher isPublic = isPublic(); + String name = getClass().getName() + "$LambdaVisibility"; + + assertTrue(matches(name, bytes(name, Runnable.class, Visibility.PUBLIC), false, isPublic)); + assertFalse( + matches(name, bytes(name, Runnable.class, Visibility.PACKAGE_PRIVATE), true, isPublic)); + } + + @Test + void nestedCallbackDoesNotEndOuterTransform() { + String name = getClass().getName() + "$NestedCallbackTarget"; + TypeFactory typeFactory = TypeFactory.typeFactory.get(); + typeFactory.switchContext(TypeFactoryTest.class.getClassLoader()); + + byte[] outerBytecode = bytes(name, Runnable.class); + typeFactory.beginTransform(name, outerBytecode); + try { + typeFactory.endTransform(bytes(name + "$Nested", Runnable.class)); + + assertEquals(Runnable.class.getName(), resolveCurrentTargetInterface(name)); + } finally { + typeFactory.endTransform(outerBytecode); + } + } + + private static String resolveInterface(String name, byte[] bytecode, boolean lambda) { + TypeFactory typeFactory = TypeFactory.typeFactory.get(); + typeFactory.switchContext(TypeFactoryTest.class.getClassLoader()); + if (lambda) { + typeFactory.beginLambdaTransform(Runnable.class.getName()); + } + typeFactory.beginTransform(name, bytecode); + try { + TypeDescription type = TypeFactory.findType(name); + return type.getInterfaces().getOnly().asErasure().getName(); + } finally { + typeFactory.endTransform(); + if (lambda) { + typeFactory.endLambdaTransform(); + } + } + } + + private static String resolveCurrentTargetInterface(String name) { + TypeDescription type = TypeFactory.findType(name); + return type.getInterfaces().getOnly().asErasure().getName(); + } + + private static boolean matches( + String name, byte[] bytecode, boolean lambda, ElementMatcher matcher) { + TypeFactory typeFactory = TypeFactory.typeFactory.get(); + typeFactory.switchContext(TypeFactoryTest.class.getClassLoader()); + if (lambda) { + typeFactory.beginLambdaTransform(Runnable.class.getName()); + } + typeFactory.beginTransform(name, bytecode); + try { + return matcher.matches(TypeFactory.findType(name)); + } finally { + typeFactory.endTransform(); + if (lambda) { + typeFactory.endLambdaTransform(); + } + } + } + + private static byte[] bytes(String name, Class implementedInterface) { + return bytes(name, implementedInterface, Visibility.PUBLIC); + } + + private static byte[] bytes(String name, Class implementedInterface, Visibility visibility) { + return new ByteBuddy() + .subclass(Object.class) + .name(name) + .implement(implementedInterface) + .modifiers(visibility) + .make() + .getBytes(); + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/main/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentation.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/main/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentation.java new file mode 100644 index 00000000000..93cbea124a6 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/main/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentation.java @@ -0,0 +1,233 @@ +package datadog.trace.instrumentation.java.lang.invoke; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.Platform; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHelper; +import net.bytebuddy.asm.AsmVisitorWrapper; +import net.bytebuddy.description.field.FieldDescription; +import net.bytebuddy.description.field.FieldList; +import net.bytebuddy.description.method.MethodList; +import net.bytebuddy.description.type.TypeDefinition; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.implementation.Implementation; +import net.bytebuddy.jar.asm.ClassVisitor; +import net.bytebuddy.jar.asm.ClassWriter; +import net.bytebuddy.jar.asm.MethodVisitor; +import net.bytebuddy.jar.asm.Opcodes; +import net.bytebuddy.jar.asm.Type; +import net.bytebuddy.matcher.ElementMatcher; +import net.bytebuddy.pool.TypePool; +import net.bytebuddy.utility.OpenedClassReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Routes generated lambda bytes through the agent transformer before definition, allowing + * instrumentations registered for an exact functional interface to transform them. + * + *

An ASM visitor is required because the transform call must be inserted immediately after the + * lambda bytes are generated, in the middle of the metafactory method. + * + *

The injected call executes inside {@code java.lang.invoke}, so it only requires a module read + * edge to the bootstrap helper; the package does not need to be opened reflectively. + */ +@AutoService(InstrumenterModule.class) +public final class LambdaMetafactoryInstrumentation extends InstrumenterModule.ContextTracking + implements Instrumenter.ForBootstrap, + Instrumenter.ForSingleType, + Instrumenter.HasTypeAdvice, + Instrumenter.WithTypeStructure { + + private static final Logger log = LoggerFactory.getLogger(LambdaMetafactoryInstrumentation.class); + + private static final String METAFACTORY = "java.lang.invoke.InnerClassLambdaMetafactory"; + + private static final String LAMBDA_CLASS_NAME_FIELD = "lambdaClassName"; + private static final String TARGET_CLASS_FIELD = "targetClass"; + private static final String INTERFACE_CLASS_FIELD = "interfaceClass"; + private static final String LEGACY_INTERFACE_CLASS_FIELD = "samBase"; + + public LambdaMetafactoryInstrumentation() { + super("lambda"); + } + + @Override + protected boolean defaultEnabled() { + return false; + } + + @Override + public boolean isEnabled() { + return super.isEnabled() && !Platform.isNativeImageBuilder(); + } + + @Override + public String instrumentedType() { + return METAFACTORY; + } + + /** Require every field read by the injected bytecode. */ + @Override + public ElementMatcher structureMatcher() { + return HasMetafactoryFields.INSTANCE; + } + + /** Public because this matcher is loaded across agent class-loader boundaries. */ + public static final class HasMetafactoryFields implements ElementMatcher { + public static final HasMetafactoryFields INSTANCE = new HasMetafactoryFields(); + + @Override + public boolean matches(TypeDescription target) { + return declaresField(target, LAMBDA_CLASS_NAME_FIELD, String.class.getName()) + && declaresField(target, TARGET_CLASS_FIELD, Class.class.getName()) + && interfaceClassField(target) != null; + } + + static String interfaceClassField(TypeDescription type) { + // JDK 8 and 11 use samBase; newer JDKs use interfaceClass. + if (declaresField(type, INTERFACE_CLASS_FIELD, Class.class.getName())) { + return INTERFACE_CLASS_FIELD; + } + if (declaresField(type, LEGACY_INTERFACE_CLASS_FIELD, Class.class.getName())) { + return LEGACY_INTERFACE_CLASS_FIELD; + } + return null; + } + + private static boolean declaresField(TypeDescription type, String name, String fieldType) { + for (TypeDefinition current = type; current != null; current = current.getSuperClass()) { + for (FieldDescription field : current.asErasure().getDeclaredFields()) { + if (name.equals(field.getName()) + && fieldType.equals(field.getType().asErasure().getName())) { + return true; + } + } + } + return false; + } + } + + @Override + public void typeAdvice(TypeTransformer transformer) { + transformer.applyAdvice(new MetafactoryVisitorWrapper()); + } + + public static final class MetafactoryVisitorWrapper implements AsmVisitorWrapper { + @Override + public int mergeWriter(int flags) { + return flags | ClassWriter.COMPUTE_MAXS; + } + + @Override + public int mergeReader(int flags) { + return flags; + } + + @Override + public ClassVisitor wrap( + TypeDescription instrumentedType, + ClassVisitor classVisitor, + Implementation.Context implementationContext, + TypePool typePool, + FieldList fields, + MethodList methods, + int writerFlags, + int readerFlags) { + return new MetafactoryClassVisitor( + classVisitor, + instrumentedType.getInternalName(), + HasMetafactoryFields.interfaceClassField(instrumentedType)); + } + } + + private static final class MetafactoryClassVisitor extends ClassVisitor { + private final String slashClassName; + private final String interfaceClassField; + private boolean injected; + + MetafactoryClassVisitor(ClassVisitor cv, String slashClassName, String interfaceClassField) { + super(OpenedClassReader.ASM_API, cv); + this.slashClassName = slashClassName; + this.interfaceClassField = interfaceClassField; + } + + @Override + public MethodVisitor visitMethod( + int access, String name, String descriptor, String signature, String[] exceptions) { + MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions); + // The byte-generation method changed in JDK 25. + if (("spinInnerClass".equals(name) || "generateInnerClass".equals(name)) + && "()Ljava/lang/Class;".equals(descriptor)) { + return new MetafactoryMethodVisitor(api, mv, slashClassName, interfaceClassField, this); + } + return mv; + } + + @Override + public void visitEnd() { + super.visitEnd(); + if (!injected) { + log.warn( + "No supported byte-generation site found in {}; lambda transformation is inactive", + slashClassName); + } + } + } + + private static final class MetafactoryMethodVisitor extends MethodVisitor { + private final String slashClassName; + private final String interfaceClassField; + private final MetafactoryClassVisitor declaringVisitor; + + MetafactoryMethodVisitor( + int api, + MethodVisitor mv, + String slashClassName, + String interfaceClassField, + MetafactoryClassVisitor declaringVisitor) { + super(api, mv); + this.slashClassName = slashClassName; + this.interfaceClassField = interfaceClassField; + this.declaringVisitor = declaringVisitor; + } + + @Override + public void visitMethodInsn( + int opcode, String owner, String name, String descriptor, boolean isInterface) { + super.visitMethodInsn(opcode, owner, name, descriptor, isInterface); + // Match repackaged JDK APIs while excluding unrelated byte-array producers. The generated + // byte[] remains on the operand stack after the original call. + if ((opcode == Opcodes.INVOKEVIRTUAL + && "toByteArray".equals(name) + && "()[B".equals(descriptor) + && owner.endsWith("/ClassWriter")) + || (opcode == Opcodes.INVOKEINTERFACE + && "build".equals(name) + && descriptor.endsWith(")[B") + && owner.endsWith("/ClassFile"))) { + // stack: ..., byte[] + super.visitVarInsn(Opcodes.ALOAD, 0); + super.visitFieldInsn( + Opcodes.GETFIELD, slashClassName, LAMBDA_CLASS_NAME_FIELD, "Ljava/lang/String;"); + super.visitVarInsn(Opcodes.ALOAD, 0); + // Resolves the defining class loader and module. + super.visitFieldInsn( + Opcodes.GETFIELD, slashClassName, TARGET_CLASS_FIELD, "Ljava/lang/Class;"); + super.visitVarInsn(Opcodes.ALOAD, 0); + // Allows the helper to reject interfaces without a registered ForLambda instrumentation. + super.visitFieldInsn( + Opcodes.GETFIELD, slashClassName, interfaceClassField, "Ljava/lang/Class;"); + super.visitMethodInsn( + Opcodes.INVOKESTATIC, + Type.getInternalName(LambdaTransformerHelper.class), + "transform", + "([BLjava/lang/String;Ljava/lang/Class;Ljava/lang/Class;)[B", + false); + // stack: ..., transformed byte[] + declaringVisitor.injected = true; + } + } + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentationTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentationTest.java new file mode 100644 index 00000000000..99a42059e30 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/datadog/trace/instrumentation/java/lang/invoke/LambdaMetafactoryInstrumentationTest.java @@ -0,0 +1,309 @@ +package datadog.trace.instrumentation.java.lang.invoke; + +import static net.bytebuddy.jar.asm.Opcodes.ACC_PRIVATE; +import static net.bytebuddy.jar.asm.Opcodes.ACC_PUBLIC; +import static net.bytebuddy.jar.asm.Opcodes.ACONST_NULL; +import static net.bytebuddy.jar.asm.Opcodes.ARETURN; +import static net.bytebuddy.jar.asm.Opcodes.ASM7; +import static net.bytebuddy.jar.asm.Opcodes.INVOKEINTERFACE; +import static net.bytebuddy.jar.asm.Opcodes.INVOKESTATIC; +import static net.bytebuddy.jar.asm.Opcodes.INVOKEVIRTUAL; +import static net.bytebuddy.jar.asm.Opcodes.POP; +import static net.bytebuddy.jar.asm.Opcodes.V1_8; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHelper; +import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHolder; +import datadog.trace.instrumentation.java.lang.invoke.LambdaMetafactoryInstrumentation.MetafactoryVisitorWrapper; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.jar.asm.ClassReader; +import net.bytebuddy.jar.asm.ClassVisitor; +import net.bytebuddy.jar.asm.ClassWriter; +import net.bytebuddy.jar.asm.MethodVisitor; +import org.junit.jupiter.api.Test; + +class LambdaMetafactoryInstrumentationTest { + + private static final String HELPER = + "datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper"; + + private static boolean injectsTransformCall( + String methodName, String methodDescriptor, ClassBody body) { + ClassWriter in = new ClassWriter(0); + in.visit(V1_8, ACC_PUBLIC, "Dummy", null, "java/lang/Object", null); + MethodVisitor mv = in.visitMethod(ACC_PRIVATE, methodName, methodDescriptor, null, null); + mv.visitCode(); + body.write(mv); + mv.visitMaxs(0, 0); + mv.visitEnd(); + in.visitEnd(); + + ClassWriter out = new ClassWriter(0); + ClassVisitor visitor = + new MetafactoryVisitorWrapper() + .wrap(realMetafactoryDescription(), out, null, null, null, null, 0, 0); + new ClassReader(in.toByteArray()).accept(visitor, 0); + + AtomicBoolean found = new AtomicBoolean(false); + new ClassReader(out.toByteArray()) + .accept( + new ClassVisitor(ASM7) { + @Override + public MethodVisitor visitMethod( + int access, String name, String desc, String sig, String[] ex) { + return new MethodVisitor(ASM7) { + @Override + public void visitMethodInsn( + int opcode, String owner, String name, String desc, boolean itf) { + if (opcode == INVOKESTATIC + && HELPER.equals(owner) + && "transform".equals(name) + && "([BLjava/lang/String;Ljava/lang/Class;Ljava/lang/Class;)[B" + .equals(desc)) { + found.set(true); + } + } + }; + } + }, + 0); + return found.get(); + } + + @Test + void injectsAfterToByteArrayInSpinInnerClass() { + assertTrue( + injectsTransformCall( + "spinInnerClass", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEVIRTUAL, + "jdk/internal/org/objectweb/asm/ClassWriter", + "toByteArray", + "()[B", + false); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + @Test + void injectsAfterToByteArrayInGenerateInnerClass() { + assertTrue( + injectsTransformCall( + "generateInnerClass", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEVIRTUAL, + "jdk/internal/org/objectweb/asm/ClassWriter", + "toByteArray", + "()[B", + false); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + @Test + void injectsAfterBuildOnClassFileApi() { + assertTrue( + injectsTransformCall( + "spinInnerClass", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEINTERFACE, + "java/lang/classfile/ClassFile", + "build", + "(Ljava/lang/classfile/constantpool/ClassEntry;" + + "Ljava/lang/classfile/constantpool/ConstantPoolBuilder;" + + "Ljava/util/function/Consumer;)[B", + true); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + @Test + void doesNotInjectInUnrelatedMethod() { + assertFalse( + injectsTransformCall( + "someOtherMethod", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEVIRTUAL, + "jdk/internal/org/objectweb/asm/ClassWriter", + "toByteArray", + "()[B", + false); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + @Test + void doesNotInjectOnUnrelatedToByteArrayOwner() { + assertFalse( + injectsTransformCall( + "spinInnerClass", + "()Ljava/lang/Class;", + mv -> { + mv.visitInsn(ACONST_NULL); + mv.visitMethodInsn( + INVOKEVIRTUAL, "java/io/ByteArrayOutputStream", "toByteArray", "()[B", false); + mv.visitInsn(POP); + mv.visitInsn(ACONST_NULL); + mv.visitInsn(ARETURN); + })); + } + + /** Verifies every field read by the injected bytecode, including inherited fields. */ + @Test + void structureMatcherAcceptsTheRealMetafactory() { + assertTrue( + new LambdaMetafactoryInstrumentation() + .structureMatcher() + .matches(realMetafactoryDescription())); + } + + @Test + void structureMatcherAcceptsCurrentInterfaceField() { + assertTrue( + new LambdaMetafactoryInstrumentation() + .structureMatcher() + .matches(TypeDescription.ForLoadedType.of(CurrentMetafactoryFields.class))); + } + + @Test + void structureMatcherAcceptsLegacyInterfaceField() { + assertTrue( + new LambdaMetafactoryInstrumentation() + .structureMatcher() + .matches(TypeDescription.ForLoadedType.of(LegacyMetafactoryFields.class))); + } + + @Test + void structureMatcherRejectsTypeWithoutTheFields() { + assertFalse( + new LambdaMetafactoryInstrumentation() + .structureMatcher() + .matches(TypeDescription.ForLoadedType.of(Object.class))); + } + + private static TypeDescription realMetafactoryDescription() { + try { + return TypeDescription.ForLoadedType.of( + Class.forName("java.lang.invoke.InnerClassLambdaMetafactory")); + } catch (ClassNotFoundException e) { + throw new AssertionError(e); + } + } + + private static final class CurrentMetafactoryFields { + private String lambdaClassName; + private Class targetClass; + private Class interfaceClass; + } + + private static final class LegacyMetafactoryFields { + private String lambdaClassName; + private Class targetClass; + private Class samBase; + } + + @Test + void registeredTransformerReceivesFunctionalInterface() { + byte[] originalBytes = new byte[0]; + AtomicBoolean transformed = new AtomicBoolean(); + AtomicReference transformedInterface = new AtomicReference<>(); + LambdaTransformerHolder.set( + (className, targetClass, classBytes, interfaceClassName) -> { + transformed.set(true); + transformedInterface.set(interfaceClassName); + return classBytes; + }); + try { + byte[] result = + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class); + + assertSame(originalBytes, result); + assertTrue(transformed.get()); + assertEquals(Runnable.class.getName(), transformedInterface.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + @Test + void transformerFailureFallsBackAndDoesNotPoisonNextLambda() { + byte[] originalBytes = new byte[0]; + byte[] transformedBytes = new byte[1]; + AtomicInteger calls = new AtomicInteger(); + LambdaTransformerHolder.set( + (className, targetClass, classBytes, interfaceClassName) -> { + if (calls.getAndIncrement() == 0) { + throw new IllegalStateException("expected test failure"); + } + return transformedBytes; + }); + try { + assertSame( + originalBytes, + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class)); + assertSame( + transformedBytes, + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class)); + assertEquals(2, calls.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + @Test + void nullTransformFallsBackAndDoesNotPoisonNextLambda() { + byte[] originalBytes = new byte[0]; + byte[] transformedBytes = new byte[1]; + AtomicInteger calls = new AtomicInteger(); + LambdaTransformerHolder.set( + (className, targetClass, classBytes, interfaceClassName) -> + calls.getAndIncrement() == 0 ? null : transformedBytes); + try { + assertSame( + originalBytes, + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class)); + assertSame( + transformedBytes, + LambdaTransformerHelper.transform( + originalBytes, "test/Lambda", Object.class, Runnable.class)); + assertEquals(2, calls.get()); + } finally { + LambdaTransformerHolder.set(null); + } + } + + @FunctionalInterface + private interface ClassBody { + void write(MethodVisitor mv); + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDeferredForkedTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDeferredForkedTest.java new file mode 100644 index 00000000000..2cff1a0cb37 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDeferredForkedTest.java @@ -0,0 +1,22 @@ +package testdog.trace.instrumentation.lambda; + +import static org.junit.jupiter.api.Assertions.assertTrue; +import static testdog.trace.instrumentation.lambda.TestRunnableLambdaInstrumentation.ADVICE_MARKER_FIELD; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.FieldBackedContextAccessor; +import datadog.trace.test.junit.utils.config.WithConfig; +import org.junit.jupiter.api.Test; + +@WithConfig(key = "trace.lambda.enabled", value = "true") +@WithConfig(key = "experimental.defer.integrations.until", value = "30s") +public class LambdaMetafactoryDeferredForkedTest extends AbstractInstrumentationTest { + + @Test + void generatedLambdaIsTransformedWhileRegularMatchingIsDeferred() throws Exception { + Runnable lambda = () -> {}; + + assertTrue(lambda instanceof FieldBackedContextAccessor); + assertTrue(lambda.getClass().getDeclaredField(ADVICE_MARKER_FIELD).isSynthetic()); + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDisabledForkedTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDisabledForkedTest.java new file mode 100644 index 00000000000..543e16a3f84 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryDisabledForkedTest.java @@ -0,0 +1,32 @@ +package testdog.trace.instrumentation.lambda; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static testdog.trace.instrumentation.lambda.TestRunnableLambdaInstrumentation.ADVICE_MARKER_FIELD; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.FieldBackedContextAccessor; +import datadog.trace.bootstrap.instrumentation.java.module.JpmsHelper; +import org.junit.jupiter.api.Test; + +public class LambdaMetafactoryDisabledForkedTest extends AbstractInstrumentationTest { + + @Test + void lambdaTransformationIsDisabledByDefault() { + Runnable lambda = () -> {}; + + assertFalse(lambda instanceof FieldBackedContextAccessor); + assertFalse(hasAdviceMarker(lambda)); + assertFalse( + JpmsHelper.getAllTriggers().contains("java.lang.invoke.InnerClassLambdaMetafactory"), + "disabled lambda instrumentation should not register a JPMS clearance trigger"); + } + + private static boolean hasAdviceMarker(Object lambda) { + try { + lambda.getClass().getDeclaredField(ADVICE_MARKER_FIELD); + return true; + } catch (NoSuchFieldException ignored) { + return false; + } + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryIntegrationTest.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryIntegrationTest.java new file mode 100644 index 00000000000..f23ca339e15 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/LambdaMetafactoryIntegrationTest.java @@ -0,0 +1,85 @@ +package testdog.trace.instrumentation.lambda; + +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 testdog.trace.instrumentation.lambda.TestRunnableLambdaInstrumentation.ADVICE_MARKER_FIELD; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.bootstrap.FieldBackedContextAccessor; +import datadog.trace.test.junit.utils.config.WithConfig; +import java.lang.reflect.Field; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; +import org.junit.jupiter.api.Test; + +/** Lambda integration tests outside the ignored {@code datadog.*} prefix. */ +@WithConfig(key = "trace.lambda.enabled", value = "true") +public class LambdaMetafactoryIntegrationTest extends AbstractInstrumentationTest { + + @Test + void registeredLambdaReceivesItsInstrumentationTransformations() { + // Link after the agent is installed. + Runnable lambda = () -> {}; + + assertTrue( + lambda instanceof FieldBackedContextAccessor, + "test instrumentation should field-inject Runnable lambdas"); + assertTrue( + hasAdviceMarker(lambda), + "test instrumentation should apply its own type advice to Runnable lambdas"); + assertEquals( + 1, + contextFieldCount(lambda), + "lambda should only receive context fields for keys it implements"); + } + + @Test + void testInstrumentationHandlesRunnableCaptureShapes() { + AtomicInteger counter = new AtomicInteger(); + int delta = 7; + Runnable[] lambdas = {() -> {}, counter::incrementAndGet, () -> counter.addAndGet(delta)}; + + for (Runnable lambda : lambdas) { + assertTrue( + lambda instanceof FieldBackedContextAccessor, + "every Runnable lambda shape should be field-injected"); + assertTrue( + hasAdviceMarker(lambda), + "every Runnable lambda shape should receive the test instrumentation advice"); + lambda.run(); + } + assertEquals(8, counter.get()); + } + + @Test + void unregisteredLambdaInterfaceIsNotTransformed() { + Supplier lambda = Object::new; + + assertFalse( + lambda instanceof FieldBackedContextAccessor, + "only interfaces registered by a lambda instrumentation should be transformed"); + assertFalse( + hasAdviceMarker(lambda), + "an unregistered lambda interface should not receive the test instrumentation advice"); + } + + private static boolean hasAdviceMarker(Object lambda) { + try { + lambda.getClass().getDeclaredField(ADVICE_MARKER_FIELD); + return true; + } catch (NoSuchFieldException ignored) { + return false; + } + } + + private static long contextFieldCount(Object lambda) { + long count = 0; + for (Field field : lambda.getClass().getDeclaredFields()) { + if (field.getName().startsWith("__datadogContext$")) { + count++; + } + } + return count; + } +} diff --git a/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/TestRunnableLambdaInstrumentation.java b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/TestRunnableLambdaInstrumentation.java new file mode 100644 index 00000000000..e1768f23600 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-lang/java-lang-1.8/src/test/java/testdog/trace/instrumentation/lambda/TestRunnableLambdaInstrumentation.java @@ -0,0 +1,84 @@ +package testdog.trace.instrumentation.lambda; + +import static net.bytebuddy.matcher.ElementMatchers.any; + +import com.google.auto.service.AutoService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.bootstrap.instrumentation.java.concurrent.State; +import java.util.HashMap; +import java.util.Map; +import java.util.function.Supplier; +import net.bytebuddy.asm.AsmVisitorWrapper; +import net.bytebuddy.description.field.FieldDescription; +import net.bytebuddy.description.field.FieldList; +import net.bytebuddy.description.method.MethodList; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.implementation.Implementation; +import net.bytebuddy.jar.asm.ClassVisitor; +import net.bytebuddy.jar.asm.FieldVisitor; +import net.bytebuddy.jar.asm.Opcodes; +import net.bytebuddy.matcher.ElementMatcher; +import net.bytebuddy.pool.TypePool; +import net.bytebuddy.utility.OpenedClassReader; + +/** Test-only consumer of generated Runnable lambda transformation. */ +@AutoService(InstrumenterModule.class) +public final class TestRunnableLambdaInstrumentation extends InstrumenterModule.ContextTracking + implements Instrumenter.ForLambda, Instrumenter.HasTypeAdvice { + + static final String ADVICE_MARKER_FIELD = "__datadog_test_for_lambda"; + + public TestRunnableLambdaInstrumentation() { + super("java-lang"); + } + + @Override + public String lambdaInterface() { + return Runnable.class.getName(); + } + + @Override + public ElementMatcher lambdaMatcher() { + return any(); + } + + @Override + public void typeAdvice(TypeTransformer transformer) { + transformer.applyAdvice(new AdviceMarkerVisitor()); + } + + @Override + public Map contextStore() { + Map stores = new HashMap<>(); + stores.put(Runnable.class.getName(), State.class.getName()); + stores.put(Supplier.class.getName(), State.class.getName()); + return stores; + } + + private static final class AdviceMarkerVisitor extends AsmVisitorWrapper.AbstractBase { + @Override + public ClassVisitor wrap( + TypeDescription instrumentedType, + ClassVisitor classVisitor, + Implementation.Context implementationContext, + TypePool typePool, + FieldList fields, + MethodList methods, + int writerFlags, + int readerFlags) { + return new ClassVisitor(OpenedClassReader.ASM_API, classVisitor) { + @Override + public void visitEnd() { + FieldVisitor marker = + cv.visitField( + Opcodes.ACC_PUBLIC | Opcodes.ACC_SYNTHETIC, ADVICE_MARKER_FIELD, "Z", null, null); + if (marker != null) { + marker.visitEnd(); + } + super.visitEnd(); + } + }; + } + } +} diff --git a/docs/how_instrumentations_work.md b/docs/how_instrumentations_work.md index bfe4e6232bb..ea4b6ff9b3b 100644 --- a/docs/how_instrumentations_work.md +++ b/docs/how_instrumentations_work.md @@ -231,6 +231,7 @@ that specifies how target types will be selected for instrumentation. | [`ForSingleType`](https://github.com/DataDog/dd-trace-java/blob/297b575f0f265c1dc78f9958e7b4b9365c80d1f9/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java#L70) | `String instrumentedType()` | Instruments only a single class name known at compile time.(see [Json2FactoryInstrumentation](https://github.com/DataDog/dd-trace-java/blob/9a28dc3f0333e781b2defc378c9020bf0a44ee9a/dd-java-agent/instrumentation/jackson-core/src/main/java/datadog/trace/instrumentation/jackson/core/Json2FactoryInstrumentation.java#L19)) | | [`ForKnownTypes`](https://github.com/DataDog/dd-trace-java/blob/297b575f0f265c1dc78f9958e7b4b9365c80d1f9/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java#L75) | `String[] knownMatchingTypes()` | Instruments multiple class names known at compile time. | | [`ForTypeHierarchy`](https://github.com/DataDog/dd-trace-java/blob/297b575f0f265c1dc78f9958e7b4b9365c80d1f9/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java#L80) | `String hierarchyMarkerType()``ElementMatcher hierarchyMatcher()` | Composes more complex matchers using chained [HierarchyMatchers](https://github.com/DataDog/dd-trace-java/blob/9a28dc3f0333e781b2defc378c9020bf0a44ee9a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/matcher/HierarchyMatchers.java#L18) methods. The `hierarchyMarkerType()` method should return a type name. Classloaders without this type can skip the more expensive `hierarchyMatcher()` method. (see [HttpClientInstrumentation](https://github.com/DataDog/dd-trace-java/blob/9a28dc3f0333e781b2defc378c9020bf0a44ee9a/dd-java-agent/instrumentation/java-http-client/src/main/java/datadog/trace/instrumentation/httpclient/HttpClientInstrumentation.java#L43)) | +| [`ForLambda`](../dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java) | `String lambdaInterface()` `ElementMatcher lambdaMatcher()` | Selects generated implementations of one exact functional interface. This mode is opt-in; see [Instrumenting Generated Lambda Classes](#instrumenting-generated-lambda-classes). | | [`ForConfiguredType`](https://github.com/DataDog/dd-trace-java/blob/297b575f0f265c1dc78f9958e7b4b9365c80d1f9/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java#L93) | `Collection configuredMatchingTypes()` | **_Do not implement this interface_**_._Use `ForKnownType` instead. `ForConfiguredType` is only used for last minute additions in the field - such as when a customer has a new JDBC driver that's not in the allowed list and we need to test it and provide a workaround until the next release. | | [`ForConfiguredTypes`](https://github.com/DataDog/dd-trace-java/blob/297b575f0f265c1dc78f9958e7b4b9365c80d1f9/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java#L88) | `String configuredMatchingType();` | **_Do not implement this interface._** __Like `ForConfiguredType,` for multiple classes | @@ -262,6 +263,45 @@ or [`UrlInstrumentation`](https://github.com/DataDog/dd-trace-java/blob/3e81c006 > Without classloader available, helper classes for bootstrap instrumentation must be place into the > `:dd-java-agent:agent-bootstrap` module rather than loaded using [the default mechanism](#helper-classes). +### Instrumenting Generated Lambda Classes + +Lambda classes need special handling. The JDK creates their bytecode inside `LambdaMetafactory` and defines the class +immediately, so the usual agent transformation callback cannot reliably add fields or advice afterward. The lambda mode +routes those bytes through the existing transformer just before the class is defined. + +This mode is disabled by default. Both of the following are required: + +- Set `dd.trace.lambda.enabled=true` before the agent is installed. +- Have an enabled instrumentation implement `Instrumenter.ForLambda` for the functional interface it supports. + +The configuration flag alone does not select any lambdas. Each consumer names one exact functional interface and may +apply additional checks to the generated class: + +```java +@Override +public String lambdaInterface() { + return "com.example.Callback"; +} + +@Override +public ElementMatcher lambdaMatcher() { + return any(); +} +``` + +There is no separate allowlist to maintain. Enabled `ForLambda` instrumentations provide the interface names used by the +runtime fast path. Once selected, their normal advice, context-store field injection, helper injection, classloader +checks, and muzzle checks still apply. Context stores declared for unrelated key types are not injected. + +Lambda matching happens during generation because these classes cannot be retransformed later. For the same reason, +tests should enable the setting before agent installation and create the lambda from a call site that has not already +been linked. + +The mechanism is deliberately best-effort. Reentrant generation is skipped to avoid transformer recursion, and +CDS/AppCDS classes that bypass runtime generation are not covered. GraalVM Native Image builder processes are excluded. +If a future JDK changes the internal byte-generation shape, the agent leaves the metafactory unchanged and logs a +warning. + ### Method Matching After the type is selected, the type’s target members(e.g., methods) must next be selected using the Instrumentation diff --git a/metadata/agent-jar-checks.properties b/metadata/agent-jar-checks.properties index 383b701c52e..8c2d2da1651 100644 --- a/metadata/agent-jar-checks.properties +++ b/metadata/agent-jar-checks.properties @@ -111,6 +111,7 @@ expected.integrations = IastInstrumentation,\ jwt,\ kafka,\ kotlin_coroutine,\ + lambda,\ lettuce,\ liberty,\ log4j,\ diff --git a/metadata/supported-configurations.json b/metadata/supported-configurations.json index 1892f63130c..4599d642615 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -7988,6 +7988,14 @@ "aliases": ["DD_LEGACY_E2E_DURATION_ENABLED"] } ], + "DD_TRACE_LAMBDA_ENABLED": [ + { + "version": "A", + "type": "boolean", + "default": "false", + "aliases": ["DD_TRACE_INTEGRATION_LAMBDA_ENABLED", "DD_INTEGRATION_LAMBDA_ENABLED"] + } + ], "DD_TRACE_LETTUCE_4_ASYNC_ENABLED": [ { "version": "A",