Introduce @AppliesOn to override advices InstrumenterModule target system

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

@@ -182,7 +182,8 @@ public static ClassFileTransformer installBytebuddyAgent(
     InstrumenterState.initialize(instrumenterIndex.instrumentationCount());
 
     // combine known modules indexed at build-time with extensions contributed at run-time
-    Iterable<InstrumenterModule> instrumenterModules = withExtensions(instrumenterIndex.modules());
+    Iterable<InstrumenterModule> instrumenterModules =
+        withExtensions(instrumenterIndex.modules(enabledSystems));
 
     // This needs to be a separate loop through all instrumentations before we start adding
     // advice so that we can exclude field injection, since that will try to check exclusion
@@ -200,7 +201,7 @@ public static ClassFileTransformer installBytebuddyAgent(
     }
 
     CombiningTransformerBuilder transformerBuilder =
-        new CombiningTransformerBuilder(agentBuilder, instrumenterIndex);
+        new CombiningTransformerBuilder(agentBuilder, instrumenterIndex, enabledSystems);
 
     int installedCount = 0;
     for (InstrumenterModule module : instrumenterModules) {
@@ -295,7 +296,7 @@ public InstrumenterModule next() {
 
   public static Set<InstrumenterModule.TargetSystem> getEnabledSystems() {
     EnumSet<InstrumenterModule.TargetSystem> enabledSystems =
-        EnumSet.noneOf(InstrumenterModule.TargetSystem.class);
+        EnumSet.of(InstrumenterModule.TargetSystem.CONTEXT_TRACKING);
     InstrumenterConfig cfg = InstrumenterConfig.get();
     if (cfg.isTraceEnabled()) {
       enabledSystems.add(InstrumenterModule.TargetSystem.TRACING);
@@ -309,6 +310,9 @@ public static Set<InstrumenterModule.TargetSystem> getEnabledSystems() {
     if (cfg.getIastActivation() != ProductActivation.FULLY_DISABLED) {
       enabledSystems.add(InstrumenterModule.TargetSystem.IAST);
     }
+    if (cfg.isRaspEnabled()) {
+      enabledSystems.add(InstrumenterModule.TargetSystem.RASP);
+    }
     if (cfg.isCiVisibilityEnabled()) {
       enabledSystems.add(InstrumenterModule.TargetSystem.CIVISIBILITY);
     }

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

@@ -25,6 +25,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.asm.AsmVisitorWrapper;
@@ -33,6 +34,8 @@
 import net.bytebuddy.dynamic.DynamicType;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.utility.JavaModule;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Builds {@link InstrumenterModule}s into a single combining-matcher and splitting-transformer.
@@ -43,6 +46,8 @@
 public final class CombiningTransformerBuilder
     implements Instrumenter.TypeTransformer, Instrumenter.MethodTransformer {
 
+  private static final Logger log = LoggerFactory.getLogger(CombiningTransformerBuilder.class);
+
   // Added here instead of byte-buddy's ignores because it's relatively
   // expensive. https://github.com/DataDog/dd-trace-java/pull/1045
   private static final ElementMatcher.Junction<TypeDescription> NOT_DECORATOR_MATCHER =
@@ -57,6 +62,7 @@ public final class CombiningTransformerBuilder
   private final AgentBuilder agentBuilder;
   private final InstrumenterIndex instrumenterIndex;
   private final int knownTransformationCount;
+  private final Set<InstrumenterModule.TargetSystem> enabledSystems;
 
   private final List<MatchRecorder> matchers = new ArrayList<>();
   private final BitSet knownTypesMask;
@@ -80,7 +86,9 @@ public final class CombiningTransformerBuilder
   private final List<AgentBuilder.Transformer> advice = new ArrayList<>();
 
   public CombiningTransformerBuilder(
-      AgentBuilder agentBuilder, InstrumenterIndex instrumenterIndex) {
+      AgentBuilder agentBuilder,
+      InstrumenterIndex instrumenterIndex,
+      Set<InstrumenterModule.TargetSystem> enabledSystems) {
     this.agentBuilder = agentBuilder;
     this.instrumenterIndex = instrumenterIndex;
     int knownInstrumentationCount = instrumenterIndex.instrumentationCount();
@@ -89,6 +97,7 @@ public CombiningTransformerBuilder(
     this.transformers = new AdviceStack[knownTransformationCount];
     this.nextRuntimeInstrumentationId = knownInstrumentationCount;
     this.nextRuntimeTransformationId = knownTransformationCount;
+    this.enabledSystems = enabledSystems;
   }
 
   /** Builds matchers and transformers for an instrumentation module and its members. */
@@ -239,7 +248,26 @@ public void applyAdvice(Instrumenter.TransformingAdvice typeAdvice) {
   }
 
   @Override
-  public void applyAdvice(ElementMatcher<? super MethodDescription> matcher, String adviceClass) {
+  public void applyAdvices(
+      ElementMatcher<? super MethodDescription> matcher,
+      String adviceClass,
+      String... additionalAdviceClasses) {
+    addAdviceIfEnabled(matcher, adviceClass);
+
+    if (additionalAdviceClasses != null) {
+      for (String adviceClassName : additionalAdviceClasses) {
+        addAdviceIfEnabled(matcher, adviceClassName);
+      }
+    }
+  }
+
+  private void addAdviceIfEnabled(
+      ElementMatcher<? super MethodDescription> matcher, String adviceClass) {
+    if (!instrumenterIndex.isAdviceEnabled(adviceClass, enabledSystems)) {
+      log.debug("Skipping advice class {} as it is not enabled", adviceClass);
+      return;
+    }
+    log.debug("Installing advice class {}", adviceClass);
     Advice.WithCustomMapping customMapping = Advice.withCustomMapping();
     if (postProcessor != null) {
       customMapping = customMapping.with(postProcessor);

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

@@ -0,0 +1,108 @@
+package datadog.trace.agent.tooling;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptyMap;
+
+import datadog.trace.util.Strings;
+import java.io.IOException;
+import java.util.EnumSet;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import net.bytebuddy.jar.asm.AnnotationVisitor;
+import net.bytebuddy.jar.asm.ClassReader;
+import net.bytebuddy.jar.asm.ClassVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Scans an {@link InstrumenterModule} to extract, custom {@link
+ * datadog.trace.agent.tooling.annotation.AppliesOn} annotation in any advices that are applied.
+ */
+public final class AdviceAppliesOnScanner {
+  private static final Logger log = LoggerFactory.getLogger(AdviceAppliesOnScanner.class);
+
+  private static final String APPLIESON_ANNOTATION_DESC =
+      "Ldatadog/trace/agent/tooling/annotation/AppliesOn;";
+
+  public static Map<String, Set<InstrumenterModule.TargetSystem>> extractTargetSystemOverrides(
+      Instrumenter instrumenter) throws IOException {
+    if (!(instrumenter instanceof Instrumenter.HasMethodAdvice)) {
+      return emptyMap();
+    }
+    final String instrumenterClassName = instrumenter.getClass().getName();
+    log.debug("Processing instrumenter class: {}", instrumenterClassName);
+    final Map<String, Set<InstrumenterModule.TargetSystem>> map = new HashMap<>();
+    final Set<InstrumenterModule.TargetSystem> overriddenTargetSystems =
+        EnumSet.noneOf(InstrumenterModule.TargetSystem.class);
+    // collect the advices
+    final Set<String> adviceClassNames = new HashSet<>();
+    ((Instrumenter.HasMethodAdvice) instrumenter)
+        .methodAdvice(
+            (matcher, adviceClass, additionalClasses) -> {
+              adviceClassNames.add(adviceClass);
+              if (additionalClasses != null) {
+                adviceClassNames.addAll(asList(additionalClasses));
+              }
+            });
+    for (String adviceClassName : adviceClassNames) {
+      // process each advice
+      new ClassReader(adviceClassName)
+          .accept(
+              new ClassVisitor(Opcodes.ASM8) {
+                private String className;
+
+                @Override
+                public void visit(
+                    int version,
+                    int access,
+                    String name,
+                    String signature,
+                    String superName,
+                    String[] interfaces) {
+                  className = name.replace('/', '.');
+                }
+
+                @Override
+                public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
+                  if (APPLIESON_ANNOTATION_DESC.equals(descriptor)) {
+                    return new AnnotationVisitor(Opcodes.ASM8) {
+
+                      @Override
+                      public AnnotationVisitor visitArray(String name) {
+                        if ("value".equals(name)) {
+                          return new AnnotationVisitor(Opcodes.ASM8) {
+                            @Override
+                            public void visitEnum(String name, String descriptor, String value) {
+                              try {
+                                overriddenTargetSystems.add(
+                                    InstrumenterModule.TargetSystem.valueOf(value));
+                              } catch (IllegalArgumentException e) {
+                                log.warn("Unknown target system: {}", value);
+                              }
+                            }
+                          };
+                        }
+                        return null;
+                      }
+
+                      @Override
+                      public void visitEnd() {
+                        if (!overriddenTargetSystems.isEmpty()) {
+                          log.debug(
+                              "Found @AppliesOn on {} → {}", className, overriddenTargetSystems);
+                          map.put(Strings.getSimpleName(adviceClassName), overriddenTargetSystems);
+                        }
+                      }
+                    };
+                  }
+                  return null;
+                }
+              },
+              ClassReader.SKIP_CODE | ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);
+    }
+    return map;
+  }
+}

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

@@ -110,7 +110,15 @@ default void applyAdvice(AsmVisitorWrapper typeVisitor) {
 
   /** Applies method advice from an instrumentation that {@link HasMethodAdvice}. */
   interface MethodTransformer {
-    void applyAdvice(ElementMatcher<? super MethodDescription> matcher, String adviceClass);
+    default void applyAdvice(
+        ElementMatcher<? super MethodDescription> matcher, String adviceClass) {
+      applyAdvices(matcher, adviceClass, (String[]) null);
+    }
+
+    void applyAdvices(
+        ElementMatcher<? super MethodDescription> matcher,
+        String adviceClass,
+        String... additionalAdviceClasses);
   }
 
   /** Contributes a transformation step to the dynamic type builder. */