IAST security controls: handling custom validation and sanitization methods

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

@@ -58,6 +58,7 @@
 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;
@@ -129,7 +130,7 @@ private static void maybeApplySecurityControls(@Nullable Instrumentation instrum
     if (Config.get().getIastSecurityControlsConfiguration() == null || instrumentation == null) {
       return;
     }
-    List<SecurityControl> securityControls =
+    Map<String, List<SecurityControl>> securityControls =
         SecurityControlFormatter.format(Config.get().getIastSecurityControlsConfiguration());
     if (securityControls == null) {
       LOGGER.warn("No security controls to apply");

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

@@ -1,10 +1,12 @@
 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.Set;
-import java.util.stream.Collectors;
+import java.util.Map;
 import javax.annotation.Nullable;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
@@ -17,13 +19,10 @@ public class IastSecurityControlTransformer implements ClassFileTransformer {
   private static final Logger LOGGER =
       LoggerFactory.getLogger(IastSecurityControlTransformer.class);
 
-  private final List<SecurityControl> securityControls;
-  private final Set<String> classFilter;
+  private final Map<String, List<SecurityControl>> securityControls;
 
-  public IastSecurityControlTransformer(List<SecurityControl> securityControls) {
+  public IastSecurityControlTransformer(Map<String, List<SecurityControl>> securityControls) {
     this.securityControls = securityControls;
-    this.classFilter =
-        securityControls.stream().map(SecurityControl::getClassName).collect(Collectors.toSet());
   }
 
   @Override
@@ -34,28 +33,22 @@ public byte[] transform(
       Class<?> classBeingRedefined,
       java.security.ProtectionDomain protectionDomain,
       byte[] classfileBuffer) {
-    if (!classFilter.contains(className)) {
+    if (!securityControls.containsKey(className)) {
       return null; // Do not transform classes that are not in the classFilter
     }
-    List<SecurityControl> match = getSecurityControl(className);
+    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, 0);
+      cr.accept(cv, SKIP_DEBUG | SKIP_FRAMES);
       return cw.toByteArray();
     } catch (Throwable e) {
       LOGGER.warn("Failed to transform class: {}", className, e);
       return null;
     }
   }
-
-  private List<SecurityControl> getSecurityControl(final String className) {
-    return securityControls.stream()
-        .filter(sc -> sc.getClassName().equals(className))
-        .collect(Collectors.toList());
-  }
 }

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

@@ -19,39 +19,47 @@ public class SecurityControlMethodAdapter extends MethodVisitor {
   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) {
+      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) {
-    // Check if the opcode is a return instruction
-    if (opcode >= Opcodes.IRETURN && opcode <= Opcodes.RETURN) {
-      if (securityControl.getType() == SecurityControlType.INPUT_VALIDATOR) {
-        processInputValidator();
-      } else { // SecurityControlType.SANITIZER
+    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() {
-    Type returnType = Type.getReturnType(desc);
-    if (isPrimitive(returnType)) {
-      // no need to check primitives as we are not tainting them
-      LOGGER.warn(
-          "Sanitizers should not be used on non-primitive return types. Return type {}. Security control: {}",
-          returnType.getClassName(),
-          securityControl);
-      return;
-    }
     // Duplicate the return value on the stack
     mv.visitInsn(Opcodes.DUP);
     // Load the marks from securityControl onto the stack
@@ -71,7 +79,7 @@ private void processInputValidator() {
         if (!isPrimitive) {
           callInputValidation(parametersCount);
         }
-      } else if (securityControl.getParametersToMark().contains(i)) {
+      } else if (securityControl.getParametersToMark().get(i)) {
         if (isPrimitive) {
           LOGGER.warn(
               "Input validators should not be used on primitive types. Parameter {} with type {} .Security control: {}",
@@ -88,7 +96,9 @@ private void processInputValidator() {
 
   private void callInputValidation(int i) {
     // Duplicate the parameter value on the stack
-    mv.visitVarInsn(Opcodes.ALOAD, i);
+    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

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

@@ -1,5 +1,6 @@
 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;
@@ -36,7 +37,8 @@ public MethodVisitor visitMethod(
       }
     }
     if (match != null) {
-      return new SecurityControlMethodAdapter(mv, match, desc);
+      final boolean isStatic = (access & ACC_STATIC) > 0;
+      return new SecurityControlMethodAdapter(mv, match, desc, isStatic);
     }
     return mv;
   }

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

@@ -6,13 +6,29 @@ import datadog.trace.api.iast.propagation.PropagationModule
 import datadog.trace.api.iast.securitycontrol.SecurityControl
 import datadog.trace.api.iast.securitycontrol.SecurityControlFormatter
 import datadog.trace.test.util.DDSpecification
+import foo.bar.securitycontrol.SecurityControlStaticTestSuite
 import foo.bar.securitycontrol.SecurityControlTestSuite
 import net.bytebuddy.agent.ByteBuddyAgent
 
 import java.lang.instrument.Instrumentation
 
 class IastSecurityControlTransformerForkedTest extends DDSpecification{
 
+  //static methods
+  private static final String STATIC_SANITIZER = 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:sanitize'
+  private static final String STATIC_SANITIZE_VALIDATE_OBJECT = 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:sanitizeObject'
+  private static final String STATIC_SANITIZE_INPUTS= 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:sanitizeInputs'
+  private static final String STATIC_SANITIZE_MANY_INPUTS= 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:sanitizeManyInputs'
+  private static final String STATIC_SANITIZE_INT = 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:sanitizeInt'
+  private static final String STATIC_SANITIZE_LONG = 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:sanitizeLong'
+  private static final String STATIC_INPUT_VALIDATOR_VALIDATE_ALL = 'INPUT_VALIDATOR:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:validateAll'
+  private static final String STATIC_INPUT_VALIDATOR_VALIDATE_OVERLOADED = 'INPUT_VALIDATOR:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:validate:java.lang.Object,java.lang.String,java.lang.String:1,2'
+  private static final String STATIC_INPUT_VALIDATOR_VALIDATE_RETURNING_INT = 'INPUT_VALIDATOR:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:validateReturningInt'
+  private static final String STATIC_INPUT_VALIDATOR_VALIDATE_OBJECT = 'INPUT_VALIDATOR:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:validateObject'
+  private static final String STATIC_INPUT_VALIDATOR_VALIDATE_LONG = 'INPUT_VALIDATOR:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:validateLong'
+  private static final String STATIC_INPUT_VALIDATOR_VALIDATE_SELECTED_LONG = 'INPUT_VALIDATOR:XSS:foo.bar.securitycontrol.SecurityControlStaticTestSuite:validateSelectedLong:0'
+
+  //not static methods
   private static final String SANITIZER = 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlTestSuite:sanitize'
   private static final String SANITIZE_VALIDATE_OBJECT = 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlTestSuite:sanitizeObject'
   private static final String SANITIZE_INPUTS= 'SANITIZER:XSS:foo.bar.securitycontrol.SecurityControlTestSuite:sanitizeInputs'
@@ -29,10 +45,12 @@ class IastSecurityControlTransformerForkedTest extends DDSpecification{
 
 
   def setupSpec() {
-    final config = "${SANITIZER};${SANITIZE_VALIDATE_OBJECT};${SANITIZE_INPUTS};${SANITIZE_MANY_INPUTS};${SANITIZE_INT};${SANITIZE_LONG};${INPUT_VALIDATOR_VALIDATE_ALL};${INPUT_VALIDATOR_VALIDATE_OVERLOADED};${INPUT_VALIDATOR_VALIDATE_RETURNING_INT};${INPUT_VALIDATOR_VALIDATE_OBJECT};${INPUT_VALIDATOR_VALIDATE_LONG};${INPUT_VALIDATOR_VALIDATE_SELECTED_LONG}"
+    final staticConfig = "${STATIC_SANITIZER};${STATIC_SANITIZE_VALIDATE_OBJECT};${STATIC_SANITIZE_INPUTS};${STATIC_SANITIZE_MANY_INPUTS};${STATIC_SANITIZE_INT};${STATIC_SANITIZE_LONG};${STATIC_INPUT_VALIDATOR_VALIDATE_ALL};${STATIC_INPUT_VALIDATOR_VALIDATE_OVERLOADED};${STATIC_INPUT_VALIDATOR_VALIDATE_RETURNING_INT};${STATIC_INPUT_VALIDATOR_VALIDATE_OBJECT};${STATIC_INPUT_VALIDATOR_VALIDATE_LONG};${STATIC_INPUT_VALIDATOR_VALIDATE_SELECTED_LONG}"
+    final config =  "${SANITIZER};${SANITIZE_VALIDATE_OBJECT};${SANITIZE_INPUTS};${SANITIZE_MANY_INPUTS};${SANITIZE_INT};${SANITIZE_LONG};${INPUT_VALIDATOR_VALIDATE_ALL};${INPUT_VALIDATOR_VALIDATE_OVERLOADED};${INPUT_VALIDATOR_VALIDATE_RETURNING_INT};${INPUT_VALIDATOR_VALIDATE_OBJECT};${INPUT_VALIDATOR_VALIDATE_LONG};${INPUT_VALIDATOR_VALIDATE_SELECTED_LONG}"
+    final fullConfig = "${staticConfig};${config}"
     Instrumentation instrumentation =  ByteBuddyAgent.install()
-    List<SecurityControl> securityControls =
-      SecurityControlFormatter.format(config)
+    Map<String,List<SecurityControl>> securityControls =
+      SecurityControlFormatter.format(fullConfig)
     assert securityControls != null
     instrumentation.addTransformer(new IastSecurityControlTransformer(securityControls), true)
   }
@@ -45,7 +63,15 @@ class IastSecurityControlTransformerForkedTest extends DDSpecification{
     final marks = (VulnerabilityMarks.XSS_MARK | VulnerabilityMarks.CUSTOM_SECURITY_CONTROL_MARK)
 
     when:
-    SecurityControlTestSuite.&"$method".call(*args)
+    SecurityControlStaticTestSuite.&"$method".call(*args)
+
+    then:
+    expected * iastModule.markIfTainted( toSanitize, marks)
+    0 * _
+
+    when:
+    final suite = new SecurityControlTestSuite()
+    suite.&"$method".call(*args)
 
     then:
     expected * iastModule.markIfTainted( toSanitize, marks)
@@ -68,7 +94,17 @@ class IastSecurityControlTransformerForkedTest extends DDSpecification{
     final marks = (VulnerabilityMarks.XSS_MARK | VulnerabilityMarks.CUSTOM_SECURITY_CONTROL_MARK)
 
     when:
-    SecurityControlTestSuite.&"$method".call(*args)
+    SecurityControlStaticTestSuite.&"$method".call(*args)
+
+    then:
+    for (final validate : toValidate){
+      expected * iastModule.markIfTainted(validate, marks)
+    }
+    0 * _
+
+    when:
+    final suite = new SecurityControlTestSuite()
+    suite.&"$method".call(*args)
 
     then:
     for (final validate : toValidate){

dd-java-agent/agent-iast/src/test/java/foo/bar/securitycontrol/SecurityControlStaticTestSuite.java

@@ -0,0 +1,100 @@
+package foo.bar.securitycontrol;
+
+public class SecurityControlStaticTestSuite {
+
+  public static boolean validateAll(long input, String input2) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validateAll(String input) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validateAll(String input, String input2) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validateAll(
+      String input,
+      String input2,
+      String input3,
+      String input4,
+      String input5,
+      String input6,
+      String input7,
+      String input8,
+      String input9,
+      String input10) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validateLong(long input, String input2) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validateLong(String input, long input2) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validateLong(long intput1, String input2, long input3) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validateSelectedLong(long intput1) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validateSelectedLong(long input1, long intput2) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validate(String input) {
+    return true; // dummy implementation
+  }
+
+  public static boolean validate(Object o, String input, String input2) {
+    return true; // dummy implementation
+  }
+
+  public static int validateReturningInt(String input) {
+    return 1; // dummy implementation
+  }
+
+  public static int validateObject(Object input) {
+    return 1; // dummy implementation
+  }
+
+  public static String sanitize(String input) {
+    return "Sanitized";
+  }
+
+  public static Object sanitizeObject(String input) {
+    return "Sanitized";
+  }
+
+  public static String sanitizeInputs(String input, Object input2, int input3) {
+    return "Sanitized";
+  }
+
+  public static String sanitizeManyInputs(
+      String input,
+      String input2,
+      String input3,
+      String input4,
+      String input5,
+      String input6,
+      String input7,
+      String input8,
+      String input9,
+      String input10) {
+    return "Sanitized";
+  }
+
+  public static int sanitizeInt(int input) {
+    return input;
+  }
+
+  public static long sanitizeLong(long input) {
+    return input;
+  }
+}