Merge branch 'master' into mario.vidal/taint_tracking_string_builder_…

…set_length

.gitlab-ci.yml

@@ -220,8 +220,9 @@ onboarding_tests_installer:
         SCENARIO: [ SIMPLE_INSTALLER_AUTO_INJECTION, SIMPLE_AUTO_INJECTION_PROFILING ]
 
 onboarding_tests_k8s_injection:
-  variables:
-    WEBLOG_VARIANT: dd-lib-java-init-test-app
+  parallel:
+    matrix:
+      - WEBLOG_VARIANT: [dd-lib-java-init-test-app]
 
 create_key:
   stage: generate-signing-key

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/CallSiteSpecification.java

@@ -570,11 +570,20 @@ public AfterSpecification(
 
     @Override
     protected void validateAdvice(@Nonnull final ValidationContext context) {
-      if (advice.isVoidReturn()) {
-        context.addError(ErrorCode.ADVICE_AFTER_SHOULD_NOT_RETURN_VOID);
-      }
-      if (findReturn() == null) {
-        context.addError(ErrorCode.ADVICE_AFTER_SHOULD_HAVE_RETURN);
+      if (shouldNotUseReturn(pointcut)) {
+        if (!advice.isVoidReturn()) {
+          context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID);
+        }
+        if (findReturn() != null) {
+          context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN);
+        }
+      } else {
+        if (advice.isVoidReturn()) {
+          context.addError(ErrorCode.ADVICE_AFTER_SHOULD_NOT_RETURN_VOID);
+        }
+        if (findReturn() == null) {
+          context.addError(ErrorCode.ADVICE_AFTER_SHOULD_HAVE_RETURN);
+        }
       }
       if (!pointcut.isConstructor()) {
         if (!isStaticPointcut() && !includeThis()) {
@@ -588,6 +597,10 @@ protected void validateAdvice(@Nonnull final ValidationContext context) {
       super.validateAdvice(context);
     }
 
+    private boolean shouldNotUseReturn(final MethodType type) {
+      return !type.isConstructor() && type.isVoidReturn();
+    }
+
     @Override
     public String toString() {
       return "@CallSite.After(" + signature + ")";

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/ErrorCode.java

@@ -359,6 +359,20 @@ public String apply(final Object[] objects) {
     }
   },
 
+  ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID {
+    @Override
+    public String apply(final Object[] objects) {
+      return "After advice for void method should return void";
+    }
+  },
+
+  ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN {
+    @Override
+    public String apply(final Object[] objects) {
+      return "After advice for void method should not contain @Return annotated parameters";
+    }
+  },
+
   ADVICE_AFTER_SHOULD_HAVE_RETURN {
     @Override
     public String apply(final Object[] objects) {

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy

@@ -473,6 +473,36 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
     }
   }
 
+
+  @CallSite(spi = CallSites)
+  class AfterAdviceWithVoidReturn {
+    @CallSite.After("void java.lang.StringBuilder.setLength(int)")
+    static void after(@CallSite.This StringBuilder self, @CallSite.Argument(0) int length) {
+    }
+  }
+
+  void 'test after advice with void return'() {
+    setup:
+    final spec = buildClassSpecification(AfterAdviceWithVoidReturn)
+    final generator = buildAdviceGenerator(buildDir)
+
+    when:
+    final result = generator.generate(spec)
+
+    then:
+    assertNoErrors result
+    assertCallSites(result.file) {
+      advices(0) {
+        pointcut('java/lang/StringBuilder', 'setLength', '(I)V')
+        statements(
+          'handler.dupInvoke(owner, descriptor, StackDupMode.COPY);',
+          'handler.method(opcode, owner, name, descriptor, isInterface);',
+          'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceWithVoidReturn", "after", "(Ljava/lang/StringBuilder;I)V");',
+        )
+      }
+    }
+  }
+
   private static AdviceGenerator buildAdviceGenerator(final File targetFolder) {
     return new AdviceGeneratorImpl(targetFolder, pointcutParser())
   }

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceSpecificationTest.groovy

@@ -542,4 +542,26 @@ class AdviceSpecificationTest extends BaseCsiPluginTest {
     then:
     0 * context.addError(_, _)
   }
+
+
+  @CallSite(spi = CallSites)
+  class AfterWithVoidWrongAdvice {
+    @CallSite.After("void java.lang.String.getChars(int, int, char[], int)")
+    static String after(@CallSite.AllArguments final Object[] args, @CallSite.Return final String result) {
+      return result;
+    }
+  }
+
+  void 'test after advice with void should not use @Return'() {
+    setup:
+    final context = mockValidationContext()
+    final spec = buildClassSpecification(AfterWithVoidWrongAdvice)
+
+    when:
+    spec.advices.each { it.validate(context) }
+
+    then:
+    1 * context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID, _)
+    1 * context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN, _)
+  }
 }

communication/src/main/java/datadog/communication/ddagent/SharedCommunicationObjects.java

@@ -1,6 +1,7 @@
 package datadog.communication.ddagent;
 
 import static datadog.communication.ddagent.TracerVersion.TRACER_VERSION;
+import static datadog.trace.util.AgentThreadFactory.AGENT_THREAD_GROUP;
 
 import datadog.common.container.ContainerInfo;
 import datadog.common.socket.SocketUtils;
@@ -9,6 +10,7 @@
 import datadog.remoteconfig.ConfigurationPoller;
 import datadog.remoteconfig.DefaultConfigurationPoller;
 import datadog.trace.api.Config;
+import datadog.trace.util.AgentTaskScheduler;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 import okhttp3.HttpUrl;
@@ -98,8 +100,11 @@ public DDAgentFeaturesDiscovery featuresDiscovery(Config config) {
               agentUrl,
               config.isTraceAgentV05Enabled(),
               config.isTracerMetricsEnabled());
-      if (!"true".equalsIgnoreCase(System.getProperty("dd.test.no.early.discovery"))) {
-        featuresDiscovery.discover();
+      if (AGENT_THREAD_GROUP.equals(Thread.currentThread().getThreadGroup())) {
+        featuresDiscovery.discover(); // safe to run on same thread
+      } else {
+        // avoid performing blocking I/O operation on application thread
+        AgentTaskScheduler.INSTANCE.execute(featuresDiscovery::discover);
       }
     }
     return featuresDiscovery;

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -529,7 +529,7 @@ public void execute() {
 
       installDatadogTracer(initTelemetry, scoClass, sco);
       maybeStartAppSec(scoClass, sco);
-      maybeStartIast(scoClass, sco);
+      maybeStartIast(instrumentation, scoClass, sco);
       maybeStartCiVisibility(instrumentation, scoClass, sco);
       maybeStartLogsIntake(scoClass, sco);
       // start debugger before remote config to subscribe to it before starting to poll
@@ -820,14 +820,14 @@ private static boolean isSupportedAppSecArch() {
     return true;
   }
 
-  private static void maybeStartIast(Class<?> scoClass, Object o) {
+  private static void maybeStartIast(Instrumentation instrumentation, Class<?> scoClass, Object o) {
     if (iastEnabled || !iastFullyDisabled) {
 
       StaticEventLogger.begin("IAST");
 
       try {
         SubscriptionService ss = AgentTracer.get().getSubscriptionService(RequestContextSlot.IAST);
-        startIast(ss, scoClass, o);
+        startIast(instrumentation, ss, scoClass, o);
       } catch (Exception e) {
         log.error("Error starting IAST subsystem", e);
       }
@@ -836,12 +836,13 @@ private static void maybeStartIast(Class<?> scoClass, Object o) {
     }
   }
 
-  private static void startIast(SubscriptionService ss, Class<?> scoClass, Object sco) {
+  private static void startIast(
+      Instrumentation instrumentation, SubscriptionService ss, Class<?> scoClass, Object sco) {
     try {
       final Class<?> appSecSysClass = AGENT_CLASSLOADER.loadClass("com.datadog.iast.IastSystem");
       final Method iastInstallerMethod =
-          appSecSysClass.getMethod("start", SubscriptionService.class);
-      iastInstallerMethod.invoke(null, ss);
+          appSecSysClass.getMethod("start", Instrumentation.class, SubscriptionService.class);
+      iastInstallerMethod.invoke(null, instrumentation, ss);
     } catch (final Throwable e) {
       log.warn("Not starting IAST subsystem", e);
     }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/jdbc/DBInfo.java

@@ -91,9 +91,6 @@ public static class Builder {
 
     public Builder type(String type) {
       this.type = type;
-      // Those DBs use the full text of the query including the comments as a cache key,
-      // so we disable full propagation support for them to avoid destroying the cache.
-      if (type.equals("oracle")) this.fullPropagationSupport = false;
       return this;
     }
 

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerAgentTest.java

@@ -122,11 +122,11 @@ public void runEnabledWithDatadogAgent() throws InterruptedException, IOExceptio
     ConfigurationPoller configurationPoller =
         (ConfigurationPoller) sharedCommunicationObjects.configurationPoller(Config.get());
     configurationPoller.start();
-    RecordedRequest request = datadogAgentServer.takeRequest(5, TimeUnit.SECONDS);
-    assertNotNull(request);
-    assertEquals("/info", request.getPath());
-    request = datadogAgentServer.takeRequest(5, TimeUnit.SECONDS);
-    assertNotNull(request);
+    RecordedRequest request;
+    do {
+      request = datadogAgentServer.takeRequest(5, TimeUnit.SECONDS);
+      assertNotNull(request);
+    } while ("/info".equals(request.getPath()));
     assertEquals("/v0.7/config", request.getPath());
     DebuggerAgent.stop();
     datadogAgentServer.shutdown();

dd-java-agent/agent-iast/build.gradle

@@ -50,6 +50,7 @@ dependencies {
   implementation project(':internal-api')
   implementation project(':internal-api:internal-api-9')
   implementation libs.moshi
+  implementation libs.bundles.asm
 
   testFixturesApi project(':dd-java-agent:testing')
   testFixturesApi project(':utils:test-utils')

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastSystem.java

@@ -8,6 +8,7 @@
 import com.datadog.iast.propagation.CodecModuleImpl;
 import com.datadog.iast.propagation.PropagationModuleImpl;
 import com.datadog.iast.propagation.StringModuleImpl;
+import com.datadog.iast.securitycontrol.IastSecurityControlTransformer;
 import com.datadog.iast.sink.ApplicationModuleImpl;
 import com.datadog.iast.sink.CommandInjectionModuleImpl;
 import com.datadog.iast.sink.HardcodedSecretModuleImpl;
@@ -47,12 +48,17 @@
 import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.IastModule;
 import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.securitycontrol.SecurityControl;
+import datadog.trace.api.iast.securitycontrol.SecurityControlFormatter;
 import datadog.trace.api.iast.telemetry.IastMetricCollector;
 import datadog.trace.api.iast.telemetry.Verbosity;
 import datadog.trace.util.AgentTaskScheduler;
 import datadog.trace.util.stacktrace.StackWalkerFactory;
+import java.lang.instrument.Instrumentation;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.UndeclaredThrowableException;
+import java.util.List;
+import java.util.Map;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
@@ -67,11 +73,21 @@ public class IastSystem {
   public static Verbosity VERBOSITY = Verbosity.OFF;
 
   public static void start(final SubscriptionService ss) {
-    start(ss, null);
+    start(null, ss, null);
+  }
+
+  public static void start(final SubscriptionService ss, OverheadController overheadController) {
+    start(null, ss, overheadController);
+  }
+
+  public static void start(final Instrumentation instrumentation, final SubscriptionService ss) {
+    start(instrumentation, ss, null);
   }
 
   public static void start(
-      final SubscriptionService ss, @Nullable OverheadController overheadController) {
+      @Nullable final Instrumentation instrumentation,
+      final SubscriptionService ss,
+      @Nullable OverheadController overheadController) {
     final Config config = Config.get();
     final ProductActivation iast = config.getIastActivation();
     final ProductActivation appSec = config.getAppSecActivation();
@@ -106,9 +122,23 @@ public static void start(
     registerRequestEndedCallback(ss, addTelemetry, dependencies);
     registerHeadersCallback(ss);
     registerGrpcServerRequestMessageCallback(ss);
+    maybeApplySecurityControls(instrumentation);
     LOGGER.debug("IAST started");
   }
 
+  private static void maybeApplySecurityControls(@Nullable Instrumentation instrumentation) {
+    if (Config.get().getIastSecurityControlsConfiguration() == null || instrumentation == null) {
+      return;
+    }
+    Map<String, List<SecurityControl>> securityControls =
+        SecurityControlFormatter.format(Config.get().getIastSecurityControlsConfiguration());
+    if (securityControls == null) {
+      LOGGER.warn("No security controls to apply");
+      return;
+    }
+    instrumentation.addTransformer(new IastSecurityControlTransformer(securityControls), true);
+  }
+
   private static IastContext.Provider contextProvider(
       final ProductActivation iast, final boolean global) {
     if (iast != FULLY_ENABLED) {

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java

@@ -652,6 +652,22 @@ public Taintable.Source findSource(
     return highestPrioritySource(to, target);
   }
 
+  @Override
+  public void markIfTainted(@Nullable Object target, int mark) {
+    if (target == null) {
+      return;
+    }
+    final IastContext ctx = IastContext.Provider.get();
+    if (ctx == null) {
+      return;
+    }
+    TaintedObjects taintedObjects = ctx.getTaintedObjects();
+    TaintedObject taintedObject = taintedObjects.get(target);
+    if (taintedObject != null) {
+      taintedObject.setRanges(markRanges(taintedObject.getRanges(), mark));
+    }
+  }
+
   @Override
   public boolean isTainted(@Nullable final Object target) {
     if (target == null) {

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/securitycontrol/AbstractMethodAdapter.java

@@ -0,0 +1,45 @@
+package com.datadog.iast.securitycontrol;
+
+import datadog.trace.api.iast.securitycontrol.SecurityControl;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+
+public class AbstractMethodAdapter extends MethodVisitor {
+
+  private static final String HELPER =
+      "datadog/trace/api/iast/securitycontrol/SecurityControlHelper";
+  private static final String METHOD = "setSecureMarks";
+  private static final String DESCRIPTOR = "(Ljava/lang/Object;I)V";
+  protected final MethodVisitor mv;
+  protected final SecurityControl securityControl;
+  protected final int accessFlags;
+  protected final Type method;
+
+  public AbstractMethodAdapter(
+      final MethodVisitor mv,
+      final SecurityControl securityControl,
+      final int accessFlags,
+      final Type method) {
+    super(Opcodes.ASM9, mv);
+    this.mv = mv;
+    this.securityControl = securityControl;
+    this.accessFlags = accessFlags;
+    this.method = method;
+  }
+
+  protected boolean isStatic() {
+    return (accessFlags & Opcodes.ACC_STATIC) > 0;
+  }
+
+  /**
+   * This method loads the current secure marks defined in the control and then calls the helper
+   * method
+   */
+  protected void loadMarksAndCallHelper() {
+    // Load the marks from securityControl onto the stack
+    mv.visitLdcInsn(securityControl.getMarks());
+    // Insert the call to setSecureMarks with the return value and marks as parameters
+    mv.visitMethodInsn(Opcodes.INVOKESTATIC, HELPER, METHOD, DESCRIPTOR, false);
+  }
+}