Limit the collections that the iast visitor can handle (#7768)

DataDog · Oct 16, 2024 · 68a5267 · 68a5267
1 parent 0bf75a3
commit 68a5267
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 7 deletions.
diff --git a/dd-java-agent/agent-iast/build.gradle b/dd-java-agent/agent-iast/build.gradle
@@ -56,6 +56,7 @@ dependencies {
   testImplementation libs.bytebuddy
   testImplementation('org.skyscreamer:jsonassert:1.5.1')
   testImplementation('org.codehaus.groovy:groovy-yaml:3.0.17')
+  testImplementation libs.guava
 
   testImplementation group: 'io.grpc', name: 'grpc-core', version: grpcVersion
   testImplementation group: 'io.grpc', name: 'grpc-protobuf', version: grpcVersion

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/ObjectVisitor.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/ObjectVisitor.java
@@ -9,8 +9,10 @@
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.IdentityHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Predicate;
@@ -23,6 +25,12 @@
 public class ObjectVisitor {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(ObjectVisitor.class);
+  private static final List<String> ALLOWED_COLLECTION_PKGS =
+      Arrays.asList(
+          "java.util",
+          "com.google.protobuf",
+          "org.apache.commons.collections",
+          "com.google.common.collect");
 
   private static final int MAX_VISITED_OBJECTS = 1000;
   private static final int MAX_DEPTH = 10;
@@ -122,6 +130,9 @@ private State visitArray(final int depth, final String path, final Object[] arra
   }
 
   private State visitMap(final int depth, final String path, final Map<?, ?> map) {
+    if (!isAllowedCollection(map)) {
+      return CONTINUE;
+    }
     final int mapDepth = depth + 1;
     for (final Map.Entry<?, ?> entry : map.entrySet()) {
       final Object key = entry.getKey();
@@ -145,6 +156,9 @@ private State visitMap(final int depth, final String path, final Map<?, ?> map)
   }
 
   private State visitIterable(final int depth, final String path, final Iterable<?> iterable) {
+    if (!isAllowedCollection(iterable)) {
+      return CONTINUE;
+    }
     final int iterableDepth = depth + 1;
     int index = 0;
     for (final Object item : iterable) {
@@ -188,6 +202,19 @@ private State visitObject(final int depth, final String path, final Object value
     return ObjectVisitor.State.CONTINUE;
   }
 
+  private static boolean isAllowedCollection(final Object value) {
+    if (value == null) {
+      return false;
+    }
+    final String packageName = value.getClass().getPackage().getName();
+    for (final String allowed : ALLOWED_COLLECTION_PKGS) {
+      if (packageName.startsWith(allowed)) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   public static boolean inspectClass(final Class<?> cls) {
     if (cls.isPrimitive()) {
       return false; // skip primitives

diff --git a/...va-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy b/...va-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy
@@ -66,7 +66,6 @@ class GrpcRequestMessageHandlerTest extends IastModuleImplTestBase {
     given:
     final visitor = Mock(ObjectVisitor.Visitor) {
       visit(_ as String, _ as Object) >> {
-        println 'feo'
         return CONTINUE
       }
     }

diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/ObjectVisitorTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/ObjectVisitorTest.groovy
@@ -1,5 +1,6 @@
 package com.datadog.iast.util
 
+import com.google.common.collect.Iterables
 import foo.bar.VisitableClass
 import spock.lang.Specification
 
@@ -39,12 +40,7 @@ class ObjectVisitorTest extends Specification {
     given:
     final visitor = Mock(ObjectVisitor.Visitor)
     final wrapped = ['1', '2', '3']
-    final target = new Iterable() {
-        @Override
-        Iterator iterator() {
-          return wrapped.iterator()
-        }
-      }
+    final target = Iterables.unmodifiableIterable(wrapped)
 
     when:
     ObjectVisitor.visit(target, visitor)
@@ -74,6 +70,29 @@ class ObjectVisitorTest extends Specification {
     0 * _
   }
 
+  void 'test visiting ignored collection'() {
+    given:
+    final visitor = Mock(ObjectVisitor.Visitor)
+    final target = new AbstractList<String>() {
+        @Override
+        String get(int index) {
+          assert index == 0
+          return 'value'
+        }
+
+        @Override
+        int size() {
+          return 1
+        }
+      }
+
+    when:
+    ObjectVisitor.visit(target, visitor)
+
+    then:
+    0 * _
+  }
+
   void 'test max depth'() {
     given:
     final visitor = Mock(ObjectVisitor.Visitor)