From 68a5267476d65ddacd323d78f584fae1abed4a6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Wed, 16 Oct 2024 11:44:37 +0200 Subject: [PATCH] Limit the collections that the iast visitor can handle (#7768) --- dd-java-agent/agent-iast/build.gradle | 1 + .../com/datadog/iast/util/ObjectVisitor.java | 27 ++++++++++++++++ .../iast/GrpcRequestMessageHandlerTest.groovy | 1 - .../iast/util/ObjectVisitorTest.groovy | 31 +++++++++++++++---- 4 files changed, 53 insertions(+), 7 deletions(-) diff --git a/dd-java-agent/agent-iast/build.gradle b/dd-java-agent/agent-iast/build.gradle index a9aa0c79571..b9ce5daae90 100644 --- 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 index 01f5e568061..73501b85ad9 100644 --- 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 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/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy index 58dff1966f9..7c8ee884385 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/GrpcRequestMessageHandlerTest.groovy +++ b/dd-java-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 index 1f1c1eea745..a21c7f21f88 100644 --- 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() { + @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)