IAST security controls: handling custom validation and sanitization methods (#7997)

What Does This Do
Add secure marks to custom input validators and sanitisers via env variable

RFC (Milestone 0)

Motivation
IAST frequently identifies vulnerabilities that are reported despite the fact that the application executes custom validation mechanisms before reaching those points, effectively preventing any exploitation. As a result, security teams receive unnecessary alerts about potential vulnerabilities that are not actually exploitable.

There is a need for a mechanism to recognise when these custom validation methods are in place, to avoid false positives.

Author: jandro996

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-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);
+  }
+}

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

@@ -0,0 +1,51 @@
+package com.datadog.iast.securitycontrol;
+
+import static org.objectweb.asm.ClassReader.SKIP_DEBUG;
+import static org.objectweb.asm.ClassReader.SKIP_FRAMES;
+
+import datadog.trace.api.iast.securitycontrol.SecurityControl;
+import java.lang.instrument.ClassFileTransformer;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class IastSecurityControlTransformer implements ClassFileTransformer {
+
+  private static final Logger LOGGER =
+      LoggerFactory.getLogger(IastSecurityControlTransformer.class);
+
+  private final Map<String, List<SecurityControl>> securityControls;
+
+  public IastSecurityControlTransformer(Map<String, List<SecurityControl>> securityControls) {
+    this.securityControls = securityControls;
+  }
+
+  @Override
+  @Nullable
+  public byte[] transform(
+      ClassLoader loader,
+      String className,
+      Class<?> classBeingRedefined,
+      java.security.ProtectionDomain protectionDomain,
+      byte[] classfileBuffer) {
+    List<SecurityControl> match = securityControls.get(className);
+    if (match == null || match.isEmpty()) {
+      return null; // Do not transform classes that do not have a security control
+    }
+    try {
+      ClassReader cr = new ClassReader(classfileBuffer);
+      ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES);
+      ClassVisitor cv = new SecurityControlMethodClassVisitor(cw, match);
+      cr.accept(cv, SKIP_DEBUG | SKIP_FRAMES);
+      return cw.toByteArray();
+    } catch (Throwable e) {
+      LOGGER.warn("Failed to transform class: {}", className, e);
+      return null;
+    }
+  }
+}

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

@@ -0,0 +1,63 @@
+package com.datadog.iast.securitycontrol;
+
+import static com.datadog.iast.securitycontrol.SecurityControlMethodClassVisitor.LOGGER;
+import static com.datadog.iast.securitycontrol.SecurityControlMethodClassVisitor.isPrimitive;
+
+import datadog.trace.api.iast.securitycontrol.SecurityControl;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+
+public class InputValidatorMethodAdapter extends AbstractMethodAdapter {
+
+  private final boolean isStatic;
+
+  public InputValidatorMethodAdapter(
+      final MethodVisitor mv,
+      final SecurityControl securityControl,
+      final int accessFlags,
+      final Type method) {
+    super(mv, securityControl, accessFlags, method);
+    isStatic = isStatic();
+  }
+
+  @Override
+  public void visitCode() {
+    super.visitCode();
+    processInputValidator();
+  }
+
+  private void processInputValidator() {
+    boolean allParameters = securityControl.getParametersToMark() == null;
+    Type[] types = method.getArgumentTypes();
+    int parametersCount = 0;
+    for (int i = 0; i < types.length; i++) {
+      Type type = types[i];
+      boolean isPrimitive = isPrimitive(type);
+      if (allParameters) {
+        if (!isPrimitive) {
+          callInputValidation(parametersCount);
+        }
+      } else if (securityControl.getParametersToMark().get(i)) {
+        if (isPrimitive) {
+          LOGGER.warn(
+              "Input validators should not be used on primitive types. Parameter {} with type {} .Security control: {}",
+              i,
+              type.getClassName(),
+              securityControl);
+        } else {
+          callInputValidation(parametersCount);
+        }
+      }
+      parametersCount += type.getSize();
+    }
+  }
+
+  private void callInputValidation(int i) {
+    // Load the parameter onto the stack
+    mv.visitVarInsn(
+        Opcodes.ALOAD,
+        isStatic ? i : i + 1); // instance methods have this as first element in the stack
+    loadMarksAndCallHelper();
+  }
+}

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

@@ -0,0 +1,31 @@
+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 SanitizerMethodAdapter extends AbstractMethodAdapter {
+
+  public SanitizerMethodAdapter(
+      final MethodVisitor mv,
+      final SecurityControl securityControl,
+      final int accessFlags,
+      final Type method) {
+    super(mv, securityControl, accessFlags, method);
+  }
+
+  @Override
+  public void visitInsn(int opcode) {
+    if (opcode == Opcodes.ARETURN) {
+      processSanitizer();
+    }
+    super.visitInsn(opcode);
+  }
+
+  private void processSanitizer() {
+    // Duplicate the return value on the stack
+    mv.visitInsn(Opcodes.DUP);
+    loadMarksAndCallHelper();
+  }
+}