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..0f64da1b7b0 --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper.java @@ -0,0 +1,65 @@ +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; + } + // Skip lambdas declared by the agent itself to avoid self-instrumentation and recursion. + String targetName = targetClass.getName(); + if (targetName.startsWith("datadog.") || targetName.startsWith("net.bytebuddy.")) { + log.debug("Lambda {} skipped: declared by the agent", 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-bootstrap/src/main/resources/datadog/trace/bootstrap/instrumentation/java/lang/invoke/lambda_interface_name.trie b/dd-java-agent/agent-bootstrap/src/main/resources/datadog/trace/bootstrap/instrumentation/java/lang/invoke/lambda_interface_name.trie new file mode 100644 index 00000000000..0539a77402f --- /dev/null +++ b/dd-java-agent/agent-bootstrap/src/main/resources/datadog/trace/bootstrap/instrumentation/java/lang/invoke/lambda_interface_name.trie @@ -0,0 +1,11 @@ +# Reserved lambda interface manifest (inactive) + +# Production lambda transformation is intentionally not enabled for any functional interface. +# This file is not read at build time or runtime. Runtime selection is driven exclusively by +# enabled Instrumenter.ForLambda implementations and also requires trace.lambda.enabled. + +# If a validated manifest is introduced later, entries will use the ClassNameTrie format: +# +# 1 java.lang.Runnable + +# Tests register Runnable from TestRunnableLambdaInstrumentation; this file has no effect on them. 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 3a8c7065362..2f5fdeed4f5 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 @@ -5,6 +5,7 @@ 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.trace.agent.tooling.bytebuddy.SharedTypePools; import datadog.trace.agent.tooling.bytebuddy.iast.TaintableRedefinitionStrategyListener; @@ -19,6 +20,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; @@ -35,6 +39,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; @@ -163,6 +168,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 @@ -253,12 +267,86 @@ public void applied(Iterable instrumentationNames) { InstrumenterState.resetDefaultState(); try { - return transformerBuilder.installOn(inst); + ClassFileTransformer classFileTransformer = transformerBuilder.installOn(inst); + if (lambdaTransformationEnabled) { + registerLambdaTransformer(classFileTransformer, transformerBuilder.lambdaInterfaces()); + } + return classFileTransformer; } finally { SharedTypePools.endInstall(); } } + /** Registers the installed class-file transformer for generated lambdas. */ + private static void registerLambdaTransformer( + final ClassFileTransformer classFileTransformer, final String[] lambdaInterfaces) { + 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)) { + 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..c6cf17d0098 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()); @@ -75,6 +88,18 @@ public boolean matches( ids.clear(); long fromTick = InstrumenterMetrics.tick(); + String lambdaInterface = TypePoolFacade.lambdaInterface(); + if (null != lambdaInterface) { + LambdaMatchRecorder[] recorders = lambdaMatchers.get(lambdaInterface); + if (null != recorders) { + for (LambdaMatchRecorder recorder : recorders) { + recorder.record(target, classLoader, ids); + } + } + 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..ab8eb3708b0 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,6 +59,10 @@ 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; @@ -162,9 +166,41 @@ private void buildTypeInstrumentation(Instrumenter member) { } buildTypeMatcher(member, transformationId); + buildLambdaMatcher(member, transformationId); buildTypeAdvice(member, transformationId); } + private void buildLambdaMatcher(Instrumenter member, int transformationId) { + if (!(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 +327,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 +396,15 @@ 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 + // directly along with the instrumentation that requested this context store. + for (LambdaMatchRecorder recorder : lambdaRecorders) { + recorder.addTransformation(transformationId); + } + } } 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..629913181d0 --- /dev/null +++ b/dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/LambdaMatchRecorder.java @@ -0,0 +1,31 @@ +package datadog.trace.agent.tooling; + +import java.util.BitSet; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +/** Selects the transformations explicitly registered for a generated lambda interface. */ +final class LambdaMatchRecorder { + private final ElementMatcher typeMatcher; + private final ElementMatcher classLoaderMatcher; + private final BitSet transformationIds = new BitSet(); + + LambdaMatchRecorder( + int transformationId, + ElementMatcher typeMatcher, + ElementMatcher classLoaderMatcher) { + this.typeMatcher = typeMatcher; + this.classLoaderMatcher = classLoaderMatcher; + transformationIds.set(transformationId); + } + + void addTransformation(int transformationId) { + transformationIds.set(transformationId); + } + + void record(TypeDescription type, ClassLoader classLoader, BitSet matches) { + if (classLoaderMatcher.matches(classLoader) && typeMatcher.matches(type)) { + matches.or(transformationIds); + } + } +} 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/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..c36fe8c0b8e --- /dev/null +++ b/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/CombiningMatcherTest.java @@ -0,0 +1,48 @@ +package datadog.trace.agent.tooling; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +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); + + 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)); + } finally { + TypePoolFacade.endLambdaTransform(); + } + + TypePoolFacade.beginLambdaTransform(java.util.function.Supplier.class.getName()); + try { + assertFalse(matcher.matches(target, classLoader, null, null, null)); + } 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..9f7f3b3c3e3 --- /dev/null +++ b/dd-java-agent/agent-installer/src/test/java/datadog/trace/agent/tooling/LambdaTransformerTest.java @@ -0,0 +1,32 @@ +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.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformer; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.Test; + +class LambdaTransformerTest { + + @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", Object.class, new byte[0], Runnable.class.getName())); + transformed.set(false); + assertNull(transformer.transform("test/Lambda", Object.class, new byte[0], "other.Interface")); + 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/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..73c4d3eb40e 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; @@ -159,6 +161,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; @@ -231,6 +245,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 +275,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 +305,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 +389,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 +422,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..5e576849afa 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(); 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..089e18be343 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/bytebuddy/outline/TypeFactoryTest.java @@ -0,0 +1,128 @@ +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)); + } + + 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 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..ba468a4b7fd --- /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,232 @@ +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.debug( + "No injection 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/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..83453c5cec9 --- /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,70 @@ +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.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"); + } + + @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; + } + } +} 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..9653216290e --- /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,80 @@ +package testdog.trace.instrumentation.lambda; + +import static java.util.Collections.singletonMap; +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.Map; +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("test-runnable-lambda"); + } + + @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() { + return singletonMap(Runnable.class.getName(), State.class.getName()); + } + + 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/metadata/agent-jar-checks.properties b/metadata/agent-jar-checks.properties index a7bde3e2525..c02458675a1 100644 --- a/metadata/agent-jar-checks.properties +++ b/metadata/agent-jar-checks.properties @@ -110,6 +110,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 6c3fe354b68..6a97071521b 100644 --- a/metadata/supported-configurations.json +++ b/metadata/supported-configurations.json @@ -7972,6 +7972,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",