IAST security controls: handling custom validation and sanitization methods

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