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/IastSecurityControlTransformer.java

@@ -0,0 +1,54 @@
+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) {
+    if (!securityControls.containsKey(className)) {
+      return null; // Do not transform classes that are not in the classFilter
+    }
+    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/SecurityControlMethodAdapter.java

@@ -0,0 +1,113 @@
+package com.datadog.iast.securitycontrol;
+
+import datadog.trace.api.iast.securitycontrol.SecurityControl;
+import datadog.trace.api.iast.securitycontrol.SecurityControlType;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class SecurityControlMethodAdapter extends MethodVisitor {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(SecurityControlMethodAdapter.class);
+
+  public static final String HELPER =
+      "datadog/trace/api/iast/securitycontrol/SecurityControlHelper";
+  public static final String METHOD = "setSecureMarks";
+  public static final String DESCRIPTOR = "(Ljava/lang/Object;I)V";
+  private final MethodVisitor mv;
+  private final SecurityControl securityControl;
+  private final String desc;
+  boolean isStatic;
+
+  public SecurityControlMethodAdapter(
+      final MethodVisitor mv,
+      final SecurityControl securityControl,
+      final String desc,
+      boolean isStatic) {
+    super(Opcodes.ASM9, mv);
+    this.mv = mv;
+    this.securityControl = securityControl;
+    this.desc = desc;
+    this.isStatic = isStatic;
+  }
+
+  @Override
+  public void visitCode() {
+    super.visitCode();
+    if (securityControl.getType() == SecurityControlType.INPUT_VALIDATOR) {
+      processInputValidator();
+    }
+  }
+
+  @Override
+  public void visitInsn(int opcode) {
+    if (securityControl.getType() == SecurityControlType.SANITIZER) {
+      // Only process the return value if is an Object as we are not tainting primitives
+      if (opcode == Opcodes.ARETURN) {
+        processSanitizer();
+      } else {
+        Type returnType = Type.getReturnType(desc);
+        // no need to check primitives as we are not tainting them
+        LOGGER.warn(
+            "Sanitizers should not be used on primitive return types. Return type {}. Security control: {}",
+            returnType.getClassName(),
+            securityControl);
+      }
+    }
+    super.visitInsn(opcode);
+  }
+
+  private void processSanitizer() {
+    // Duplicate the return value on the stack
+    mv.visitInsn(Opcodes.DUP);
+    // 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);
+  }
+
+  private void processInputValidator() {
+    boolean allParameters = securityControl.getParametersToMark() == null;
+    Type[] types = Type.getArgumentTypes(desc);
+    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 += types[i].getSize();
+    }
+  }
+
+  private void callInputValidation(int i) {
+    // Duplicate the parameter value on the stack
+    mv.visitVarInsn(
+        Opcodes.ALOAD,
+        isStatic ? i : i + 1); // instance methods have this as first element in the stack
+    // Load the marks from securityControl onto the stack
+    mv.visitLdcInsn(securityControl.getMarks());
+    // Insert the call to setSecureMarks with the parameter value and marks as parameters
+    mv.visitMethodInsn(Opcodes.INVOKESTATIC, HELPER, METHOD, DESCRIPTOR, false);
+  }
+
+  private static boolean isPrimitive(Type type) {
+    // Check if is a primitive type
+    int sort = type.getSort();
+    return sort >= Type.BOOLEAN && sort <= Type.DOUBLE;
+  }
+}

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

@@ -0,0 +1,73 @@
+package com.datadog.iast.securitycontrol;
+
+import static org.objectweb.asm.Opcodes.ACC_STATIC;
+import static org.objectweb.asm.Opcodes.ASM8;
+
+import datadog.trace.api.iast.securitycontrol.SecurityControl;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.ClassWriter;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Type;
+
+public class SecurityControlMethodClassVisitor extends ClassVisitor {
+
+  private final List<SecurityControl> securityControls;
+
+  public SecurityControlMethodClassVisitor(
+      final ClassWriter cw, final List<SecurityControl> securityControls) {
+    super(ASM8, cw);
+    this.securityControls = securityControls;
+  }
+
+  @Override
+  @Nullable
+  public MethodVisitor visitMethod(
+      int access, String name, String desc, String signature, String[] exceptions) {
+    MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
+    if (mv == null) {
+      return null;
+    }
+    SecurityControl match = null;
+    for (SecurityControl securityControl : securityControls) {
+      if (shouldBeAdapted(securityControl, name, desc)) {
+        match = securityControl;
+        break;
+      }
+    }
+    if (match != null) {
+      final boolean isStatic = (access & ACC_STATIC) > 0;
+      return new SecurityControlMethodAdapter(mv, match, desc, isStatic);
+    }
+    return mv;
+  }
+
+  public void visitEnd() {
+    cv.visitEnd();
+  }
+
+  private boolean shouldBeAdapted(SecurityControl securityControl, String name, String desc) {
+
+    if (!securityControl.getMethod().equals(name)) {
+      return false;
+    }
+
+    if (securityControl.getParameterTypes() == null) {
+      return true;
+    }
+
+    Type[] types = Type.getArgumentTypes(desc);
+    if (types.length != securityControl.getParameterTypes().size()) {
+      return false;
+    }
+
+    for (int i = 0; i < types.length; i++) {
+      if (!types[i].getClassName().equals(securityControl.getParameterTypes().get(i))) {
+        return false;
+      }
+    }
+
+    return true;
+  }
+}

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/IastSystemTest.groovy

@@ -136,7 +136,7 @@ class IastSystemTest extends DDSpecification {
     InstrumentationBridge.clearIastModules()
 
     when:
-    Agent.maybeStartIast(null, null)
+    Agent.maybeStartIast(null, null, null)
 
     then:
     InstrumentationBridge.iastModules.each {