Skip to content

Commit a25afaf

Browse files
amarzialiPerfectSlayermcculls
authored
Introduce @AppliesOn to override advices InstrumenterModule target system (#10404)
* Introduce @AppliesOn to override advices InstrumenterModule target system * Add some documentation * Fix tests and add logging * Move isApplicable logic to isEnabled and fix index generation * fix iast test * cleanup * cleanup * wip * refactor logback advice in the proper way * Simplify annotation * improve InstrumentModuleFilter * Change doc * remove leftover * Update dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java Co-authored-by: Bruce Bujon <PerfectSlayer@users.noreply.github.com> * fix javadoc * fix things * fix build * applyAdvices * add unfiltered modules() * Apply suggestions from code review Co-authored-by: Stuart McCulloch <stuart.mcculloch@datadoghq.com> Co-authored-by: Bruce Bujon <PerfectSlayer@users.noreply.github.com> * Fix compilation issues * move logger upper * first round of suggestions * second round of suggestions * Refactor LoggerConfigInstrumentation * refactor getSimpleName on Strings * Update dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterIndex.java Co-authored-by: Stuart McCulloch <stuart.mcculloch@datadoghq.com> * Update dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterIndex.java Co-authored-by: Stuart McCulloch <stuart.mcculloch@datadoghq.com> * spotless --------- Co-authored-by: Bruce Bujon <PerfectSlayer@users.noreply.github.com> Co-authored-by: Stuart McCulloch <stuart.mcculloch@datadoghq.com>
1 parent a34cb64 commit a25afaf

59 files changed

Lines changed: 773 additions & 219 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,8 @@ public static ClassFileTransformer installBytebuddyAgent(
182182
InstrumenterState.initialize(instrumenterIndex.instrumentationCount());
183183

184184
// combine known modules indexed at build-time with extensions contributed at run-time
185-
Iterable<InstrumenterModule> instrumenterModules = withExtensions(instrumenterIndex.modules());
185+
Iterable<InstrumenterModule> instrumenterModules =
186+
withExtensions(instrumenterIndex.modules(enabledSystems));
186187

187188
// This needs to be a separate loop through all instrumentations before we start adding
188189
// advice so that we can exclude field injection, since that will try to check exclusion
@@ -200,7 +201,7 @@ public static ClassFileTransformer installBytebuddyAgent(
200201
}
201202

202203
CombiningTransformerBuilder transformerBuilder =
203-
new CombiningTransformerBuilder(agentBuilder, instrumenterIndex);
204+
new CombiningTransformerBuilder(agentBuilder, instrumenterIndex, enabledSystems);
204205

205206
int installedCount = 0;
206207
for (InstrumenterModule module : instrumenterModules) {
@@ -295,7 +296,7 @@ public InstrumenterModule next() {
295296

296297
public static Set<InstrumenterModule.TargetSystem> getEnabledSystems() {
297298
EnumSet<InstrumenterModule.TargetSystem> enabledSystems =
298-
EnumSet.noneOf(InstrumenterModule.TargetSystem.class);
299+
EnumSet.of(InstrumenterModule.TargetSystem.CONTEXT_TRACKING);
299300
InstrumenterConfig cfg = InstrumenterConfig.get();
300301
if (cfg.isTraceEnabled()) {
301302
enabledSystems.add(InstrumenterModule.TargetSystem.TRACING);
@@ -309,6 +310,9 @@ public static Set<InstrumenterModule.TargetSystem> getEnabledSystems() {
309310
if (cfg.getIastActivation() != ProductActivation.FULLY_DISABLED) {
310311
enabledSystems.add(InstrumenterModule.TargetSystem.IAST);
311312
}
313+
if (cfg.isRaspEnabled()) {
314+
enabledSystems.add(InstrumenterModule.TargetSystem.RASP);
315+
}
312316
if (cfg.isCiVisibilityEnabled()) {
313317
enabledSystems.add(InstrumenterModule.TargetSystem.CIVISIBILITY);
314318
}

‎dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import java.util.HashMap;
2626
import java.util.List;
2727
import java.util.Map;
28+
import java.util.Set;
2829
import net.bytebuddy.agent.builder.AgentBuilder;
2930
import net.bytebuddy.asm.Advice;
3031
import net.bytebuddy.asm.AsmVisitorWrapper;
@@ -33,6 +34,8 @@
3334
import net.bytebuddy.dynamic.DynamicType;
3435
import net.bytebuddy.matcher.ElementMatcher;
3536
import net.bytebuddy.utility.JavaModule;
37+
import org.slf4j.Logger;
38+
import org.slf4j.LoggerFactory;
3639

3740
/**
3841
* Builds {@link InstrumenterModule}s into a single combining-matcher and splitting-transformer.
@@ -43,6 +46,8 @@
4346
public final class CombiningTransformerBuilder
4447
implements Instrumenter.TypeTransformer, Instrumenter.MethodTransformer {
4548

49+
private static final Logger log = LoggerFactory.getLogger(CombiningTransformerBuilder.class);
50+
4651
// Added here instead of byte-buddy's ignores because it's relatively
4752
// expensive. https://github.com/DataDog/dd-trace-java/pull/1045
4853
private static final ElementMatcher.Junction<TypeDescription> NOT_DECORATOR_MATCHER =
@@ -57,6 +62,7 @@ public final class CombiningTransformerBuilder
5762
private final AgentBuilder agentBuilder;
5863
private final InstrumenterIndex instrumenterIndex;
5964
private final int knownTransformationCount;
65+
private final Set<InstrumenterModule.TargetSystem> enabledSystems;
6066

6167
private final List<MatchRecorder> matchers = new ArrayList<>();
6268
private final BitSet knownTypesMask;
@@ -80,7 +86,9 @@ public final class CombiningTransformerBuilder
8086
private final List<AgentBuilder.Transformer> advice = new ArrayList<>();
8187

8288
public CombiningTransformerBuilder(
83-
AgentBuilder agentBuilder, InstrumenterIndex instrumenterIndex) {
89+
AgentBuilder agentBuilder,
90+
InstrumenterIndex instrumenterIndex,
91+
Set<InstrumenterModule.TargetSystem> enabledSystems) {
8492
this.agentBuilder = agentBuilder;
8593
this.instrumenterIndex = instrumenterIndex;
8694
int knownInstrumentationCount = instrumenterIndex.instrumentationCount();
@@ -89,6 +97,7 @@ public CombiningTransformerBuilder(
8997
this.transformers = new AdviceStack[knownTransformationCount];
9098
this.nextRuntimeInstrumentationId = knownInstrumentationCount;
9199
this.nextRuntimeTransformationId = knownTransformationCount;
100+
this.enabledSystems = enabledSystems;
92101
}
93102

94103
/** Builds matchers and transformers for an instrumentation module and its members. */
@@ -239,7 +248,26 @@ public void applyAdvice(Instrumenter.TransformingAdvice typeAdvice) {
239248
}
240249

241250
@Override
242-
public void applyAdvice(ElementMatcher<? super MethodDescription> matcher, String adviceClass) {
251+
public void applyAdvices(
252+
ElementMatcher<? super MethodDescription> matcher,
253+
String adviceClass,
254+
String... additionalAdviceClasses) {
255+
addAdviceIfEnabled(matcher, adviceClass);
256+
257+
if (additionalAdviceClasses != null) {
258+
for (String adviceClassName : additionalAdviceClasses) {
259+
addAdviceIfEnabled(matcher, adviceClassName);
260+
}
261+
}
262+
}
263+
264+
private void addAdviceIfEnabled(
265+
ElementMatcher<? super MethodDescription> matcher, String adviceClass) {
266+
if (!instrumenterIndex.isAdviceEnabled(adviceClass, enabledSystems)) {
267+
log.debug("Skipping advice class {} as it is not enabled", adviceClass);
268+
return;
269+
}
270+
log.debug("Installing advice class {}", adviceClass);
243271
Advice.WithCustomMapping customMapping = Advice.withCustomMapping();
244272
if (postProcessor != null) {
245273
customMapping = customMapping.with(postProcessor);

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AdviceAppliesOnScanner.java‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
package datadog.trace.agent.tooling;
2+
3+
import static java.util.Arrays.asList;
4+
import static java.util.Collections.emptyMap;
5+
6+
import datadog.trace.util.Strings;
7+
import java.io.IOException;
8+
import java.util.EnumSet;
9+
import java.util.HashMap;
10+
import java.util.HashSet;
11+
import java.util.Map;
12+
import java.util.Set;
13+
import net.bytebuddy.jar.asm.AnnotationVisitor;
14+
import net.bytebuddy.jar.asm.ClassReader;
15+
import net.bytebuddy.jar.asm.ClassVisitor;
16+
import net.bytebuddy.jar.asm.Opcodes;
17+
import org.slf4j.Logger;
18+
import org.slf4j.LoggerFactory;
19+
20+
/**
21+
* Scans an {@link InstrumenterModule} to extract, custom {@link
22+
* datadog.trace.agent.tooling.annotation.AppliesOn} annotation in any advices that are applied.
23+
*/
24+
public final class AdviceAppliesOnScanner {
25+
private static final Logger log = LoggerFactory.getLogger(AdviceAppliesOnScanner.class);
26+
27+
private static final String APPLIESON_ANNOTATION_DESC =
28+
"Ldatadog/trace/agent/tooling/annotation/AppliesOn;";
29+
30+
public static Map<String, Set<InstrumenterModule.TargetSystem>> extractTargetSystemOverrides(
31+
Instrumenter instrumenter) throws IOException {
32+
if (!(instrumenter instanceof Instrumenter.HasMethodAdvice)) {
33+
return emptyMap();
34+
}
35+
final String instrumenterClassName = instrumenter.getClass().getName();
36+
log.debug("Processing instrumenter class: {}", instrumenterClassName);
37+
final Map<String, Set<InstrumenterModule.TargetSystem>> map = new HashMap<>();
38+
final Set<InstrumenterModule.TargetSystem> overriddenTargetSystems =
39+
EnumSet.noneOf(InstrumenterModule.TargetSystem.class);
40+
// collect the advices
41+
final Set<String> adviceClassNames = new HashSet<>();
42+
((Instrumenter.HasMethodAdvice) instrumenter)
43+
.methodAdvice(
44+
(matcher, adviceClass, additionalClasses) -> {
45+
adviceClassNames.add(adviceClass);
46+
if (additionalClasses != null) {
47+
adviceClassNames.addAll(asList(additionalClasses));
48+
}
49+
});
50+
for (String adviceClassName : adviceClassNames) {
51+
// process each advice
52+
new ClassReader(adviceClassName)
53+
.accept(
54+
new ClassVisitor(Opcodes.ASM8) {
55+
private String className;
56+
57+
@Override
58+
public void visit(
59+
int version,
60+
int access,
61+
String name,
62+
String signature,
63+
String superName,
64+
String[] interfaces) {
65+
className = name.replace('/', '.');
66+
}
67+
68+
@Override
69+
public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
70+
if (APPLIESON_ANNOTATION_DESC.equals(descriptor)) {
71+
return new AnnotationVisitor(Opcodes.ASM8) {
72+
73+
@Override
74+
public AnnotationVisitor visitArray(String name) {
75+
if ("value".equals(name)) {
76+
return new AnnotationVisitor(Opcodes.ASM8) {
77+
@Override
78+
public void visitEnum(String name, String descriptor, String value) {
79+
try {
80+
overriddenTargetSystems.add(
81+
InstrumenterModule.TargetSystem.valueOf(value));
82+
} catch (IllegalArgumentException e) {
83+
log.warn("Unknown target system: {}", value);
84+
}
85+
}
86+
};
87+
}
88+
return null;
89+
}
90+
91+
@Override
92+
public void visitEnd() {
93+
if (!overriddenTargetSystems.isEmpty()) {
94+
log.debug(
95+
"Found @AppliesOn on {} → {}", className, overriddenTargetSystems);
96+
map.put(Strings.getSimpleName(adviceClassName), overriddenTargetSystems);
97+
}
98+
}
99+
};
100+
}
101+
return null;
102+
}
103+
},
104+
ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
105+
}
106+
return map;
107+
}
108+
}

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,15 @@ default void applyAdvice(AsmVisitorWrapper typeVisitor) {
110110

111111
/** Applies method advice from an instrumentation that {@link HasMethodAdvice}. */
112112
interface MethodTransformer {
113-
void applyAdvice(ElementMatcher<? super MethodDescription> matcher, String adviceClass);
113+
default void applyAdvice(
114+
ElementMatcher<? super MethodDescription> matcher, String adviceClass) {
115+
applyAdvices(matcher, adviceClass, (String[]) null);
116+
}
117+
118+
void applyAdvices(
119+
ElementMatcher<? super MethodDescription> matcher,
120+
String adviceClass,
121+
String... additionalAdviceClasses);
114122
}
115123

116124
/** Contributes a transformation step to the dynamic type builder. */

0 commit comments

Comments
 (0)