From 988b767ea7abca1d9e1da6cf941bd835ea8079d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Tue, 10 Dec 2024 16:43:44 +0100 Subject: [PATCH 01/19] First view of the new TaintableDb --- .../bytebuddy/iast/TaintableDbVisitor.java | 190 ++++++++++++++++++ .../bytebuddy/matcher/ignored_class_name.trie | 1 + .../jdbc/IastResultSetInstrumentation.java | 88 ++++++++ .../datadog/trace/api/iast/SourceTypes.java | 4 +- .../datadog/trace/api/iast/TaintableDb.java | 49 +++++ 5 files changed, 331 insertions(+), 1 deletion(-) create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableDbVisitor.java create mode 100644 dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java create mode 100644 internal-api/src/main/java/datadog/trace/api/iast/TaintableDb.java diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableDbVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableDbVisitor.java new file mode 100644 index 00000000000..4dab8d959de --- /dev/null +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableDbVisitor.java @@ -0,0 +1,190 @@ +package datadog.trace.agent.tooling.bytebuddy.iast; + +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; +import javax.annotation.Nullable; +import net.bytebuddy.asm.AsmVisitorWrapper; +import net.bytebuddy.description.field.FieldDescription; +import net.bytebuddy.description.field.FieldList; +import net.bytebuddy.description.method.MethodList; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.implementation.Implementation; +import net.bytebuddy.jar.asm.ClassVisitor; +import net.bytebuddy.jar.asm.FieldVisitor; +import net.bytebuddy.jar.asm.Opcodes; +import net.bytebuddy.jar.asm.Type; +import net.bytebuddy.pool.TypePool; +import net.bytebuddy.utility.OpenedClassReader; + +public class TaintableDbVisitor implements AsmVisitorWrapper { + + public static volatile boolean DEBUG = false; + static volatile boolean ENABLED = true; + + private static final String INTERFACE_NAME = "datadog/trace/api/iast/TaintableDb"; + private static final String FIELD_NAME = "$$DD$recordsRead"; + private static final String GETTER_NAME = "$$DD$getRecordsRead"; + private static final String SETTER_NAME = "$$DD$setRecordsRead"; + + private final Set types; + + public TaintableDbVisitor(final String... classNames) { + types = new HashSet<>(Arrays.asList(classNames)); + } + + @Override + public int mergeWriter(final int flags) { + return flags; + } + + @Override + public int mergeReader(int flags) { + return flags; + } + + @Override + public ClassVisitor wrap( + final TypeDescription instrumentedType, + final ClassVisitor classVisitor, + final Implementation.Context implementationContext, + final TypePool typePool, + final FieldList fields, + final MethodList methods, + final int writerFlags, + final int readerFlags) { + if (ENABLED) { + return types.contains(instrumentedType.getName()) + ? new AddTaintableDbInterfaceVisitor(classVisitor) + : classVisitor; + } else { + return NoOp.INSTANCE.wrap( + instrumentedType, + classVisitor, + implementationContext, + typePool, + fields, + methods, + writerFlags, + readerFlags); + } + } + + private static class AddTaintableDbInterfaceVisitor extends ClassVisitor { + + private String owner; + + private boolean addTaintable = true; + + protected AddTaintableDbInterfaceVisitor(final ClassVisitor classVisitor) { + super(OpenedClassReader.ASM_API, classVisitor); + } + + @Override + public void visit( + final int version, + final int access, + final String name, + final String signature, + final String superName, + final String[] interfaces) { + owner = name; + if (interfaces != null) { + for (final String iface : interfaces) { + if (INTERFACE_NAME.equals(iface)) { + addTaintable = false; + break; + } + } + } + super.visit( + version, + access, + name, + signature, + superName, + addTaintable ? addInterface(interfaces) : interfaces); + } + + @Override + public void visitEnd() { + if (addTaintable) { + addField(); + // addGetter(); + // if (!DEBUG) { + // addSetter(); + // } else { + // addSetterDebug(); + // } + } + } + + private String[] addInterface(@Nullable final String[] interfaces) { + if (interfaces == null || interfaces.length == 0) { + return new String[] {INTERFACE_NAME}; + } else { + final String[] newInterfaces = new String[interfaces.length + 1]; + System.arraycopy(interfaces, 0, newInterfaces, 0, interfaces.length); + newInterfaces[newInterfaces.length - 1] = INTERFACE_NAME; + return newInterfaces; + } + } + + private void addField() { + final FieldVisitor fv = + cv.visitField( + Opcodes.ACC_PRIVATE | Opcodes.ACC_TRANSIENT | Opcodes.ACC_VOLATILE, + FIELD_NAME, + Type.INT_TYPE.getDescriptor(), + null, + null); + fv.visitEnd(); + } + + // private void addGetter() { + // final MethodVisitor mv = + // cv.visitMethod(Opcodes.ACC_PUBLIC, GETTER_NAME, Type.INT_TYPE.getDescriptor(), null, + // null); + // mv.visitCode(); + // mv.visitVarInsn(Opcodes.ALOAD, 0); + // mv.visitFieldInsn(Opcodes.GETSTATIC , owner, FIELD_NAME, Type.INT_TYPE.getDescriptor()); + // mv.visitInsn(Opcodes.ARETURN); + // mv.visitMaxs(1, 1); + // mv.visitEnd(); + // } + // + // private void addSetter() { + // final MethodVisitor mv = + // cv.visitMethod( + // Opcodes.ACC_PUBLIC, SETTER_NAME, Type.INT_TYPE.getDescriptor(), null, null); + // mv.visitCode(); + // mv.visitVarInsn(Opcodes.ALOAD, 0); + // mv.visitVarInsn(Opcodes.ILOAD, 1); + // mv.visitFieldInsn(Opcodes.PUTFIELD, owner, FIELD_NAME, Type.INT_TYPE.getDescriptor()); + // mv.visitInsn(Opcodes.RETURN); + // mv.visitMaxs(2, 2); + // mv.visitEnd(); + // } + // + // private void addSetterDebug() { + // final MethodVisitor mv = + // cv.visitMethod( + // Opcodes.ACC_PUBLIC, SETTER_NAME, Type.INT_TYPE.getDescriptor(), null, null); + // mv.visitCode(); + // mv.visitVarInsn(Opcodes.ALOAD, 0); + // mv.visitVarInsn(Opcodes.ILOAD, 1); + // mv.visitFieldInsn(Opcodes.PUTFIELD, owner, FIELD_NAME, Type.INT_TYPE.getDescriptor()); + // + // mv.visitVarInsn(Opcodes.ALOAD, 0); + // mv.visitMethodInsn( + // Opcodes.INVOKESTATIC, + // Type.getInternalName(TaintableDb.DebugLogger.class), + // "logTaint", + // "(Ldatadog/trace/api/iast/TaintableDb;)V", + // false); + // mv.visitInsn(Opcodes.RETURN); + // mv.visitMaxs(2, 2); + // mv.visitEnd(); + // } + } +} diff --git a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie index 9372507c6b9..7212f21c352 100644 --- a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie +++ b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie @@ -42,6 +42,7 @@ 1 io.opentelemetry.javaagent.* 1 java.* 0 java.lang.ClassLoader +0 java.sql.ResultSet # allow exception profiling instrumentation 0 java.lang.Exception 0 java.lang.Error diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java new file mode 100644 index 00000000000..72f3cebec22 --- /dev/null +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -0,0 +1,88 @@ +package datadog.trace.instrumentation.jdbc; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +import com.google.auto.service.AutoService; +import datadog.trace.advice.ActiveRequestContext; +import datadog.trace.advice.RequiresRequestContext; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.agent.tooling.bytebuddy.iast.TaintableDbVisitor; +import datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers; +import datadog.trace.api.gateway.RequestContext; +import datadog.trace.api.gateway.RequestContextSlot; +import datadog.trace.api.iast.IastContext; +import datadog.trace.api.iast.InstrumentationBridge; +import datadog.trace.api.iast.Source; +import datadog.trace.api.iast.SourceTypes; +import datadog.trace.api.iast.TaintableDb; +import datadog.trace.api.iast.propagation.PropagationModule; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +@AutoService(InstrumenterModule.class) +public class IastResultSetInstrumentation extends InstrumenterModule.Iast + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasTypeAdvice { + + public IastResultSetInstrumentation() { + super("jdbc", "resultset"); + } + + // @Override + // public String instrumentedType() { + // return "java.sql.ResultSet"; + // } + + @Override + public String hierarchyMarkerType() { + return "java.sql.ResultSet"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(NameMatchers.named("java.sql.ResultSet")); + } + + @Override + public void typeAdvice(TypeTransformer transformer) { + transformer.applyAdvice(new TaintableDbVisitor(hierarchyMarkerType())); + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(named("getInt")).and(takesArguments(int.class)), + IastResultSetInstrumentation.class.getName() + "$GetParameterAdvice"); + } + + @RequiresRequestContext(RequestContextSlot.IAST) + public static class GetParameterAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + @Source(SourceTypes.SQL_TABLE) + public static void onExit( + @Advice.Argument(0) final int columnIndex, + @Advice.Return final Object value, + @Advice.This final TaintableDb resultSet, + @ActiveRequestContext RequestContext reqCtx) { + // int recordsRead = resultSet.$$DD$RecordsRead; + // int recordsRead = resultSet.$$DD$getRecordsRead(); + // resultSet.$$DD$setRecordsRead(recordsRead + 1); + // if (recordsRead > 1) { + // return; + // } + if (value == null) { + return; + } + final PropagationModule module = InstrumentationBridge.PROPAGATION; + if (module == null) { + return; + } + IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); + module.taintString(ctx, String.valueOf(value), SourceTypes.SQL_TABLE); + } + } +} diff --git a/internal-api/src/main/java/datadog/trace/api/iast/SourceTypes.java b/internal-api/src/main/java/datadog/trace/api/iast/SourceTypes.java index 1c3dc208d2a..515cab6f252 100644 --- a/internal-api/src/main/java/datadog/trace/api/iast/SourceTypes.java +++ b/internal-api/src/main/java/datadog/trace/api/iast/SourceTypes.java @@ -24,6 +24,7 @@ private SourceTypes() {} public static final byte GRPC_BODY = 13; public static final byte KAFKA_MESSAGE_KEY = 14; public static final byte KAFKA_MESSAGE_VALUE = 15; + public static final byte SQL_TABLE = 16; /** Array indexed with all source types, the index should match the source types values */ public static final String[] STRINGS = { @@ -42,7 +43,8 @@ private SourceTypes() {} "http.request.path", "grpc.request.body", "kafka.message.key", - "kafka.message.value" + "kafka.message.value", + "sql.row.value" }; public static String toString(final byte source) { diff --git a/internal-api/src/main/java/datadog/trace/api/iast/TaintableDb.java b/internal-api/src/main/java/datadog/trace/api/iast/TaintableDb.java new file mode 100644 index 00000000000..1fd8fcee017 --- /dev/null +++ b/internal-api/src/main/java/datadog/trace/api/iast/TaintableDb.java @@ -0,0 +1,49 @@ +package datadog.trace.api.iast; + +import de.thetaphi.forbiddenapis.SuppressForbidden; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public interface TaintableDb { + + int $$DD$getRecordsRead(); + + void $$DD$setRecordsRead(final int recordsRead); + + @SuppressForbidden + class DebugLogger { + private static final Logger LOGGER; + + static { + try { + LOGGER = LoggerFactory.getLogger("TaintableDb tainted objects"); + Class levelCls = Class.forName("ch.qos.logback.classic.Level"); + Method setLevel = LOGGER.getClass().getMethod("setLevel", levelCls); + Object debugLevel = levelCls.getField("DEBUG").get(null); + setLevel.invoke(LOGGER, debugLevel); + } catch (IllegalAccessException + | NoSuchFieldException + | ClassNotFoundException + | NoSuchMethodException + | InvocationTargetException e) { + throw new RuntimeException(e); + } + } + + public static void logTaint(TaintableDb t) { + String content; + if (t.getClass().getName().startsWith("java.")) { + content = t.toString(); + } else { + content = "(value not shown)"; // toString() may trigger tainting + } + LOGGER.debug( + "taint: {}[{}] {}", + t.getClass().getSimpleName(), + Integer.toHexString(System.identityHashCode(t)), + content); + } + } +} From 5c87acda4329c63b638fb538060514fe9993ab65 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Wed, 11 Dec 2024 15:50:58 +0100 Subject: [PATCH 02/19] Add tainted db values with contextStore --- .../com/datadog/iast/sink/SinkModuleBase.java | 19 +- .../java/com/datadog/iast/taint/Ranges.java | 18 ++ .../com/datadog/iast/taint/RangesTest.groovy | 16 ++ .../bytebuddy/iast/TaintableDbVisitor.java | 190 ------------------ .../jdbc/IastResultSetInstrumentation.java | 50 ++--- .../IastResultSetInstrumentationTest.groovy | 136 +++++++++++++ .../datadog/trace/api/iast/TaintableDb.java | 49 ----- 7 files changed, 214 insertions(+), 264 deletions(-) delete mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableDbVisitor.java create mode 100644 dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy delete mode 100644 internal-api/src/main/java/datadog/trace/api/iast/TaintableDb.java diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java index a735ebd4a13..ee9a2fac8ad 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java @@ -21,13 +21,16 @@ import datadog.trace.api.Config; import datadog.trace.api.Pair; import datadog.trace.api.iast.IastContext; +import datadog.trace.api.iast.SourceTypes; import datadog.trace.api.iast.Taintable; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.instrumentation.iastinstrumenter.IastExclusionTrie; import datadog.trace.instrumentation.iastinstrumenter.SourceMapperImpl; import datadog.trace.util.stacktrace.StackWalker; +import java.util.Arrays; import java.util.Iterator; +import java.util.List; import java.util.function.Predicate; import java.util.stream.Stream; import javax.annotation.Nonnull; @@ -39,6 +42,8 @@ public abstract class SinkModuleBase { private static final int MAX_EVIDENCE_LENGTH = Config.get().getIastTruncationMaxValueLength(); + private static final List DB_INJECTION_TYPES = + Arrays.asList(VulnerabilityType.SQL_INJECTION, VulnerabilityType.XSS); protected final OverheadController overheadController; protected final Reporter reporter; @@ -143,9 +148,21 @@ protected final Evidence checkInjection( return null; } + // filter out ranges that are not SQL_TABLE for types that are not DB_INJECTION_TYPES + final Range[] filteredRanges; + if (!DB_INJECTION_TYPES.contains(type)) { + filteredRanges = Ranges.excludeRangesBySource(valueRanges, SourceTypes.SQL_TABLE); + } else { + filteredRanges = valueRanges; + } + + if (filteredRanges == null || filteredRanges.length == 0) { + return null; + } + final StringBuilder evidence = new StringBuilder(); final RangeBuilder ranges = new RangeBuilder(); - addToEvidence(type, evidence, ranges, value, valueRanges, evidenceBuilder); + addToEvidence(type, evidence, ranges, value, filteredRanges, evidenceBuilder); // check if finally we have an injection if (ranges.isEmpty()) { diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java index a7b7c2df88d..483307f9e0d 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java @@ -426,4 +426,22 @@ public static Range[] splitRanges( return splittedRanges; } + + /** + * Remove the ranges that have the same origin as the input source. + * + * @param ranges the ranges to filter + * @param source the byte value of the source to exclude (see {@link SourceTypes}) + */ + public static Range[] excludeRangesBySource(Range[] ranges, byte source) { + RangeBuilder newRanges = new RangeBuilder(ranges.length); + + for (Range range : ranges) { + if (range.getSource().getOrigin() != source) { + newRanges.add(range); + } + } + + return newRanges.toArray(); + } } diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy index 7d39fa1cad6..748e11cb7c1 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy @@ -12,6 +12,7 @@ import static com.datadog.iast.util.HttpHeader.LOCATION import static com.datadog.iast.util.HttpHeader.REFERER import static datadog.trace.api.iast.SourceTypes.GRPC_BODY import static datadog.trace.api.iast.SourceTypes.REQUEST_HEADER_VALUE +import static datadog.trace.api.iast.SourceTypes.SQL_TABLE import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED import static com.datadog.iast.taint.Ranges.mergeRanges import static datadog.trace.api.iast.SourceTypes.REQUEST_HEADER_NAME @@ -378,6 +379,21 @@ class RangesTest extends DDSpecification { 1 | 3 | 2 | range(8, 8) | 0 | 0 | [] } + void 'test excludeRangesBySource method'() { + when: + final result = Ranges.excludeRangesBySource(ranges as Range[], source) + + then: + final expectedArray = expected as Range[] + result == expectedArray + + where: + ranges | source | expected + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | SQL_TABLE | [range(5, 3)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | REQUEST_HEADER_NAME | [rangeWithSource(0, 5, SQL_TABLE)] + [] | SQL_TABLE | [] + } + Range[] rangesFromSpec(List> spec) { def ranges = new Range[spec.size()] int j = 0 diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableDbVisitor.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableDbVisitor.java deleted file mode 100644 index 4dab8d959de..00000000000 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/iast/TaintableDbVisitor.java +++ /dev/null @@ -1,190 +0,0 @@ -package datadog.trace.agent.tooling.bytebuddy.iast; - -import java.util.Arrays; -import java.util.HashSet; -import java.util.Set; -import javax.annotation.Nullable; -import net.bytebuddy.asm.AsmVisitorWrapper; -import net.bytebuddy.description.field.FieldDescription; -import net.bytebuddy.description.field.FieldList; -import net.bytebuddy.description.method.MethodList; -import net.bytebuddy.description.type.TypeDescription; -import net.bytebuddy.implementation.Implementation; -import net.bytebuddy.jar.asm.ClassVisitor; -import net.bytebuddy.jar.asm.FieldVisitor; -import net.bytebuddy.jar.asm.Opcodes; -import net.bytebuddy.jar.asm.Type; -import net.bytebuddy.pool.TypePool; -import net.bytebuddy.utility.OpenedClassReader; - -public class TaintableDbVisitor implements AsmVisitorWrapper { - - public static volatile boolean DEBUG = false; - static volatile boolean ENABLED = true; - - private static final String INTERFACE_NAME = "datadog/trace/api/iast/TaintableDb"; - private static final String FIELD_NAME = "$$DD$recordsRead"; - private static final String GETTER_NAME = "$$DD$getRecordsRead"; - private static final String SETTER_NAME = "$$DD$setRecordsRead"; - - private final Set types; - - public TaintableDbVisitor(final String... classNames) { - types = new HashSet<>(Arrays.asList(classNames)); - } - - @Override - public int mergeWriter(final int flags) { - return flags; - } - - @Override - public int mergeReader(int flags) { - return flags; - } - - @Override - public ClassVisitor wrap( - final TypeDescription instrumentedType, - final ClassVisitor classVisitor, - final Implementation.Context implementationContext, - final TypePool typePool, - final FieldList fields, - final MethodList methods, - final int writerFlags, - final int readerFlags) { - if (ENABLED) { - return types.contains(instrumentedType.getName()) - ? new AddTaintableDbInterfaceVisitor(classVisitor) - : classVisitor; - } else { - return NoOp.INSTANCE.wrap( - instrumentedType, - classVisitor, - implementationContext, - typePool, - fields, - methods, - writerFlags, - readerFlags); - } - } - - private static class AddTaintableDbInterfaceVisitor extends ClassVisitor { - - private String owner; - - private boolean addTaintable = true; - - protected AddTaintableDbInterfaceVisitor(final ClassVisitor classVisitor) { - super(OpenedClassReader.ASM_API, classVisitor); - } - - @Override - public void visit( - final int version, - final int access, - final String name, - final String signature, - final String superName, - final String[] interfaces) { - owner = name; - if (interfaces != null) { - for (final String iface : interfaces) { - if (INTERFACE_NAME.equals(iface)) { - addTaintable = false; - break; - } - } - } - super.visit( - version, - access, - name, - signature, - superName, - addTaintable ? addInterface(interfaces) : interfaces); - } - - @Override - public void visitEnd() { - if (addTaintable) { - addField(); - // addGetter(); - // if (!DEBUG) { - // addSetter(); - // } else { - // addSetterDebug(); - // } - } - } - - private String[] addInterface(@Nullable final String[] interfaces) { - if (interfaces == null || interfaces.length == 0) { - return new String[] {INTERFACE_NAME}; - } else { - final String[] newInterfaces = new String[interfaces.length + 1]; - System.arraycopy(interfaces, 0, newInterfaces, 0, interfaces.length); - newInterfaces[newInterfaces.length - 1] = INTERFACE_NAME; - return newInterfaces; - } - } - - private void addField() { - final FieldVisitor fv = - cv.visitField( - Opcodes.ACC_PRIVATE | Opcodes.ACC_TRANSIENT | Opcodes.ACC_VOLATILE, - FIELD_NAME, - Type.INT_TYPE.getDescriptor(), - null, - null); - fv.visitEnd(); - } - - // private void addGetter() { - // final MethodVisitor mv = - // cv.visitMethod(Opcodes.ACC_PUBLIC, GETTER_NAME, Type.INT_TYPE.getDescriptor(), null, - // null); - // mv.visitCode(); - // mv.visitVarInsn(Opcodes.ALOAD, 0); - // mv.visitFieldInsn(Opcodes.GETSTATIC , owner, FIELD_NAME, Type.INT_TYPE.getDescriptor()); - // mv.visitInsn(Opcodes.ARETURN); - // mv.visitMaxs(1, 1); - // mv.visitEnd(); - // } - // - // private void addSetter() { - // final MethodVisitor mv = - // cv.visitMethod( - // Opcodes.ACC_PUBLIC, SETTER_NAME, Type.INT_TYPE.getDescriptor(), null, null); - // mv.visitCode(); - // mv.visitVarInsn(Opcodes.ALOAD, 0); - // mv.visitVarInsn(Opcodes.ILOAD, 1); - // mv.visitFieldInsn(Opcodes.PUTFIELD, owner, FIELD_NAME, Type.INT_TYPE.getDescriptor()); - // mv.visitInsn(Opcodes.RETURN); - // mv.visitMaxs(2, 2); - // mv.visitEnd(); - // } - // - // private void addSetterDebug() { - // final MethodVisitor mv = - // cv.visitMethod( - // Opcodes.ACC_PUBLIC, SETTER_NAME, Type.INT_TYPE.getDescriptor(), null, null); - // mv.visitCode(); - // mv.visitVarInsn(Opcodes.ALOAD, 0); - // mv.visitVarInsn(Opcodes.ILOAD, 1); - // mv.visitFieldInsn(Opcodes.PUTFIELD, owner, FIELD_NAME, Type.INT_TYPE.getDescriptor()); - // - // mv.visitVarInsn(Opcodes.ALOAD, 0); - // mv.visitMethodInsn( - // Opcodes.INVOKESTATIC, - // Type.getInternalName(TaintableDb.DebugLogger.class), - // "logTaint", - // "(Ldatadog/trace/api/iast/TaintableDb;)V", - // false); - // mv.visitInsn(Opcodes.RETURN); - // mv.visitMaxs(2, 2); - // mv.visitEnd(); - // } - } -} diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index 72f3cebec22..a7636ac9b25 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -1,6 +1,7 @@ package datadog.trace.instrumentation.jdbc; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; @@ -10,7 +11,6 @@ import datadog.trace.advice.RequiresRequestContext; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.agent.tooling.bytebuddy.iast.TaintableDbVisitor; import datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; @@ -18,25 +18,23 @@ import datadog.trace.api.iast.InstrumentationBridge; import datadog.trace.api.iast.Source; import datadog.trace.api.iast.SourceTypes; -import datadog.trace.api.iast.TaintableDb; import datadog.trace.api.iast.propagation.PropagationModule; +import datadog.trace.bootstrap.ContextStore; +import datadog.trace.bootstrap.InstrumentationContext; +import java.sql.ResultSet; +import java.util.Map; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @AutoService(InstrumenterModule.class) public class IastResultSetInstrumentation extends InstrumenterModule.Iast - implements Instrumenter.ForTypeHierarchy, Instrumenter.HasTypeAdvice { + implements Instrumenter.ForTypeHierarchy { public IastResultSetInstrumentation() { super("jdbc", "resultset"); } - // @Override - // public String instrumentedType() { - // return "java.sql.ResultSet"; - // } - @Override public String hierarchyMarkerType() { return "java.sql.ResultSet"; @@ -47,33 +45,37 @@ public ElementMatcher hierarchyMatcher() { return implementsInterface(NameMatchers.named("java.sql.ResultSet")); } - @Override - public void typeAdvice(TypeTransformer transformer) { - transformer.applyAdvice(new TaintableDbVisitor(hierarchyMarkerType())); - } - @Override public void methodAdvice(MethodTransformer transformer) { transformer.applyAdvice( - isMethod().and(named("getInt")).and(takesArguments(int.class)), + isMethod() + .and(named("getString")) + .and(takesArguments(int.class).or(takesArguments(String.class))), IastResultSetInstrumentation.class.getName() + "$GetParameterAdvice"); } + @Override + public Map contextStore() { + return singletonMap("java.sql.ResultSet", Integer.class.getName()); + } + @RequiresRequestContext(RequestContextSlot.IAST) public static class GetParameterAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @Source(SourceTypes.SQL_TABLE) public static void onExit( - @Advice.Argument(0) final int columnIndex, - @Advice.Return final Object value, - @Advice.This final TaintableDb resultSet, + @Advice.Return final String value, + @Advice.This final ResultSet resultSet, @ActiveRequestContext RequestContext reqCtx) { - // int recordsRead = resultSet.$$DD$RecordsRead; - // int recordsRead = resultSet.$$DD$getRecordsRead(); - // resultSet.$$DD$setRecordsRead(recordsRead + 1); - // if (recordsRead > 1) { - // return; - // } + ContextStore contextStore = + InstrumentationContext.get(ResultSet.class, Integer.class); + if (contextStore.get(resultSet) != null) { + contextStore.put(resultSet, contextStore.get(resultSet) + 1); + return; + } else { + // first time + contextStore.put(resultSet, 1); + } if (value == null) { return; } @@ -82,7 +84,7 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taintString(ctx, String.valueOf(value), SourceTypes.SQL_TABLE); + module.taintString(ctx, value, SourceTypes.SQL_TABLE); } } } diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy new file mode 100644 index 00000000000..8e44757abb2 --- /dev/null +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy @@ -0,0 +1,136 @@ +import com.datadog.iast.taint.TaintedObject +import com.datadog.iast.test.IastAgentTestRunner +import datadog.trace.api.iast.SourceTypes +import foo.bar.IastInstrumentedConnection +import spock.lang.Shared + +import javax.sql.DataSource +import java.sql.Connection + +class IastResultSetInstrumentationTest extends IastAgentTestRunner { + @Shared + DataSource dataSource = new org.apache.tomcat.jdbc.pool.DataSource( + url: 'jdbc:h2:mem:iastUnitTest', + driverClassName: 'org.h2.Driver', + password: '', + maxActive: 1, + ) + + private Connection connection + Connection getConnection() { + connection = (connection ?: dataSource.connection) + } + + void cleanup() { + connection?.close() + } + + void setupSpec() { + dataSource.connection.withCloseable { + it.createStatement().withCloseable {stmt -> + stmt.executeUpdate('CREATE TABLE TEST (id integer NOT NULL, name VARCHAR NOT NULL)') + stmt.executeUpdate('INSERT INTO TEST VALUES(42, \'foo\')') + stmt.executeUpdate('INSERT INTO TEST VALUES(43, \'bar\')') + } + } + } + + void cleanupSpec() { + dataSource.close(true) + } + + void 'returned string is tainted with source values as sql table'() { + when: + def valueRead + TaintedObject taintedObject + runUnderIastTrace { + String sql = 'SELECT name FROM TEST LIMIT 1' + Connection constWrapper = new IastInstrumentedConnection(conn: connection) + + constWrapper.prepareStatement(sql).withCloseable { stmt -> + stmt.executeQuery().withCloseable { rs -> + if (rs.next()) { + valueRead = rs.getString(1) + } + } + } + + taintedObject = localTaintedObjects.get(valueRead) + } + + then: + valueRead == "foo" + taintedObject != null + with(taintedObject) { + with(ranges[0]) { + source.origin == SourceTypes.SQL_TABLE + source.value == valueRead + } + } + } + + void 'returned string is tainted with source values as sql table but second value is not tainted'() { + when: + def firstValue + def secondValue + TaintedObject firstTaintedObject + TaintedObject secondTaintedObject + boolean isFirst = true + runUnderIastTrace { + String sql = 'SELECT name FROM TEST LIMIT 2' + Connection constWrapper = new IastInstrumentedConnection(conn: connection) + + constWrapper.prepareStatement(sql).withCloseable { stmt -> + stmt.executeQuery().withCloseable { rs -> + while (rs.next()) { + if (isFirst) { + firstValue = rs.getString("name") + isFirst = false + } else { + secondValue = rs.getString("name") + } + } + } + } + + firstTaintedObject = localTaintedObjects.get(firstValue) + secondTaintedObject = localTaintedObjects.get(secondValue) + } + + then: + firstValue == "foo" + firstTaintedObject != null + with(firstTaintedObject) { + with(ranges[0]) { + source.origin == SourceTypes.SQL_TABLE + source.value == firstValue + } + } + secondValue == "bar" + secondTaintedObject == null + } + + void 'when returned value is not a string does not taint the result' () { + when: + def valueRead + TaintedObject taintedObject + runUnderIastTrace { + String sql = 'SELECT id FROM TEST LIMIT 1' + Connection constWrapper = new IastInstrumentedConnection(conn: connection) + + constWrapper.prepareStatement(sql).withCloseable { stmt -> + stmt.executeQuery().withCloseable { rs -> + if (rs.next()) { + valueRead = rs.getInt(1) + } + } + } + + taintedObject = localTaintedObjects.get(valueRead) + } + + then: + valueRead == 42 + taintedObject == null + } +} diff --git a/internal-api/src/main/java/datadog/trace/api/iast/TaintableDb.java b/internal-api/src/main/java/datadog/trace/api/iast/TaintableDb.java deleted file mode 100644 index 1fd8fcee017..00000000000 --- a/internal-api/src/main/java/datadog/trace/api/iast/TaintableDb.java +++ /dev/null @@ -1,49 +0,0 @@ -package datadog.trace.api.iast; - -import de.thetaphi.forbiddenapis.SuppressForbidden; -import java.lang.reflect.InvocationTargetException; -import java.lang.reflect.Method; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -public interface TaintableDb { - - int $$DD$getRecordsRead(); - - void $$DD$setRecordsRead(final int recordsRead); - - @SuppressForbidden - class DebugLogger { - private static final Logger LOGGER; - - static { - try { - LOGGER = LoggerFactory.getLogger("TaintableDb tainted objects"); - Class levelCls = Class.forName("ch.qos.logback.classic.Level"); - Method setLevel = LOGGER.getClass().getMethod("setLevel", levelCls); - Object debugLevel = levelCls.getField("DEBUG").get(null); - setLevel.invoke(LOGGER, debugLevel); - } catch (IllegalAccessException - | NoSuchFieldException - | ClassNotFoundException - | NoSuchMethodException - | InvocationTargetException e) { - throw new RuntimeException(e); - } - } - - public static void logTaint(TaintableDb t) { - String content; - if (t.getClass().getName().startsWith("java.")) { - content = t.toString(); - } else { - content = "(value not shown)"; // toString() may trigger tainting - } - LOGGER.debug( - "taint: {}[{}] {}", - t.getClass().getSimpleName(), - Integer.toHexString(System.identityHashCode(t)), - content); - } - } -} From 1e156894d62158478057560c82a4245a5dbbf8dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Wed, 11 Dec 2024 16:36:31 +0100 Subject: [PATCH 03/19] Add config file for maximum number of taints per ResultSet --- .../jdbc/IastResultSetInstrumentation.java | 10 ++++++---- .../main/java/datadog/trace/api/ConfigDefaults.java | 1 + .../main/java/datadog/trace/api/config/IastConfig.java | 1 + .../src/main/java/datadog/trace/api/Config.java | 7 +++++++ 4 files changed, 15 insertions(+), 4 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index a7636ac9b25..0694a427b14 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -1,9 +1,9 @@ package datadog.trace.instrumentation.jdbc; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static java.util.Collections.singletonMap; import static net.bytebuddy.matcher.ElementMatchers.isMethod; -import static net.bytebuddy.matcher.ElementMatchers.named; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; @@ -11,7 +11,7 @@ import datadog.trace.advice.RequiresRequestContext; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; -import datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers; +import datadog.trace.api.Config; import datadog.trace.api.gateway.RequestContext; import datadog.trace.api.gateway.RequestContextSlot; import datadog.trace.api.iast.IastContext; @@ -42,7 +42,7 @@ public String hierarchyMarkerType() { @Override public ElementMatcher hierarchyMatcher() { - return implementsInterface(NameMatchers.named("java.sql.ResultSet")); + return implementsInterface(named("java.sql.ResultSet")); } @Override @@ -71,7 +71,9 @@ public static void onExit( InstrumentationContext.get(ResultSet.class, Integer.class); if (contextStore.get(resultSet) != null) { contextStore.put(resultSet, contextStore.get(resultSet) + 1); - return; + if (contextStore.get(resultSet) > Config.get().getIastDbRowsToTaint()) { + return; + } } else { // first time contextStore.put(resultSet, 1); diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java index 9c12e7dda8c..1c72955e252 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java @@ -134,6 +134,7 @@ public final class ConfigDefaults { static final boolean DEFAULT_IAST_ANONYMOUS_CLASSES_ENABLED = true; static final boolean DEFAULT_IAST_STACK_TRACE_ENABLED = true; + static final int DEFAULT_IAST_DB_ROWS_TO_TAINT = 1; static final boolean DEFAULT_USM_ENABLED = false; diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java index 73f0b93a4c9..06e24bf13a1 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java @@ -27,6 +27,7 @@ public final class IastConfig { public static final String IAST_SOURCE_MAPPING_MAX_SIZE = "iast.source-mapping.max-size"; public static final String IAST_EXPERIMENTAL_PROPAGATION_ENABLED = "iast.experimental.propagation.enabled"; + public static final String IAST_DB_ROWS_TO_TAINT = "iast.db.rows.to.taint"; public static final String IAST_STACK_TRACE_ENABLED = "iast.stacktrace.enabled"; diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index e925bd5aaf1..f21d6b9fa1e 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -306,6 +306,7 @@ public static String getHostName() { private final int iastSourceMappingMaxSize; private final boolean iastStackTraceEnabled; private final boolean iastExperimentalPropagationEnabled; + private final int iastDbRowsToTaint; private final boolean ciVisibilityTraceSanitationEnabled; private final boolean ciVisibilityAgentlessEnabled; @@ -1318,6 +1319,8 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment()) configProvider.getBoolean(IAST_STACK_TRACE_ENABLED, DEFAULT_IAST_STACK_TRACE_ENABLED); iastExperimentalPropagationEnabled = configProvider.getBoolean(IAST_EXPERIMENTAL_PROPAGATION_ENABLED, false); + iastDbRowsToTaint = + configProvider.getInteger(IAST_DB_ROWS_TO_TAINT, DEFAULT_IAST_DB_ROWS_TO_TAINT); ciVisibilityTraceSanitationEnabled = configProvider.getBoolean(CIVISIBILITY_TRACE_SANITATION_ENABLED, true); @@ -2606,6 +2609,10 @@ public boolean isIastExperimentalPropagationEnabled() { return iastExperimentalPropagationEnabled; } + public int getIastDbRowsToTaint() { + return iastDbRowsToTaint; + } + public boolean isCiVisibilityEnabled() { return instrumenterConfig.isCiVisibilityEnabled(); } From fe595d2f2ca4eface18807a1517f6340f1d9be2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Wed, 11 Dec 2024 16:37:19 +0100 Subject: [PATCH 04/19] Improve tests for IastResultSetInstrumentation --- .../IastResultSetInstrumentationTest.groovy | 72 ++++++++++++++----- 1 file changed, 53 insertions(+), 19 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy index 8e44757abb2..261e27763ea 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy @@ -71,11 +71,8 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { void 'returned string is tainted with source values as sql table but second value is not tainted'() { when: - def firstValue - def secondValue - TaintedObject firstTaintedObject - TaintedObject secondTaintedObject - boolean isFirst = true + List valuesRead = [] + List taintedObjects = [] runUnderIastTrace { String sql = 'SELECT name FROM TEST LIMIT 2' Connection constWrapper = new IastInstrumentedConnection(conn: connection) @@ -83,31 +80,68 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { constWrapper.prepareStatement(sql).withCloseable { stmt -> stmt.executeQuery().withCloseable { rs -> while (rs.next()) { - if (isFirst) { - firstValue = rs.getString("name") - isFirst = false - } else { - secondValue = rs.getString("name") - } + valuesRead.add(rs.getString("name")) } } } - firstTaintedObject = localTaintedObjects.get(firstValue) - secondTaintedObject = localTaintedObjects.get(secondValue) + taintedObjects.add(localTaintedObjects.get(valuesRead[0])) + taintedObjects.add(localTaintedObjects.get(valuesRead[1])) } then: - firstValue == "foo" - firstTaintedObject != null - with(firstTaintedObject) { + valuesRead[0] == "foo" + taintedObjects[0] != null + with(taintedObjects[0]) { with(ranges[0]) { source.origin == SourceTypes.SQL_TABLE - source.value == firstValue + source.value == valuesRead[0] + } + } + valuesRead[1] == "bar" + taintedObjects[1] == null + } + + void 'returned string is tainted with source values as sql table and can taint up to two values'() { + given: + injectSysConfig('dd.iast.db.rows.to.taint', "2") + + when: + List valuesRead = [] + List taintedObjects = [] + runUnderIastTrace { + String sql = 'SELECT name FROM TEST LIMIT 2' + Connection constWrapper = new IastInstrumentedConnection(conn: connection) + + constWrapper.prepareStatement(sql).withCloseable { stmt -> + stmt.executeQuery().withCloseable { rs -> + while (rs.next()) { + valuesRead.add(rs.getString("name")) + } + } + } + + taintedObjects.add(localTaintedObjects.get(valuesRead[0])) + taintedObjects.add(localTaintedObjects.get(valuesRead[1])) + } + + then: + valuesRead[0] == "foo" + taintedObjects[0] != null + with(taintedObjects[0]) { + with(ranges[0]) { + source.origin == SourceTypes.SQL_TABLE + source.value == valuesRead[0] + } + } + valuesRead[1] == "bar" + taintedObjects[1] != null + with(taintedObjects[1]) { + with(ranges[0]) { + source.origin == SourceTypes.SQL_TABLE + source.value == valuesRead[1] } } - secondValue == "bar" - secondTaintedObject == null } void 'when returned value is not a string does not taint the result' () { From 41ae7eb2d4b9040e5d00c83cfc196cff8b73e5e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Thu, 12 Dec 2024 11:02:08 +0100 Subject: [PATCH 05/19] Improve instrumentation and tests --- .../jdbc/IastResultSetInstrumentation.java | 30 +++++++++---- .../IastResultSetInstrumentationTest.groovy | 42 ++++++++++++++++++- .../datadog/trace/api/config/IastConfig.java | 2 +- 3 files changed, 63 insertions(+), 11 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index 0694a427b14..d0aebe42a6d 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -47,9 +47,12 @@ public ElementMatcher hierarchyMatcher() { @Override public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod().and(named("next").and(takesArguments(0))), + IastResultSetInstrumentation.class.getName() + "$NextAdvice"); transformer.applyAdvice( isMethod() - .and(named("getString")) + .and(named("getString").or(named("getNString"))) .and(takesArguments(int.class).or(takesArguments(String.class))), IastResultSetInstrumentation.class.getName() + "$GetParameterAdvice"); } @@ -59,6 +62,21 @@ public Map contextStore() { return singletonMap("java.sql.ResultSet", Integer.class.getName()); } + public static class NextAdvice { + @Advice.OnMethodExit(suppress = Throwable.class) + @Source(SourceTypes.SQL_TABLE) + public static void onExit(@Advice.This final ResultSet resultSet) { + ContextStore contextStore = + InstrumentationContext.get(ResultSet.class, Integer.class); + if (contextStore.get(resultSet) != null) { + contextStore.put(resultSet, contextStore.get(resultSet) + 1); + } else { + // first time + contextStore.put(resultSet, 1); + } + } + } + @RequiresRequestContext(RequestContextSlot.IAST) public static class GetParameterAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @@ -69,14 +87,8 @@ public static void onExit( @ActiveRequestContext RequestContext reqCtx) { ContextStore contextStore = InstrumentationContext.get(ResultSet.class, Integer.class); - if (contextStore.get(resultSet) != null) { - contextStore.put(resultSet, contextStore.get(resultSet) + 1); - if (contextStore.get(resultSet) > Config.get().getIastDbRowsToTaint()) { - return; - } - } else { - // first time - contextStore.put(resultSet, 1); + if (contextStore.get(resultSet) > Config.get().getIastDbRowsToTaint()) { + return; } if (value == null) { return; diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy index 261e27763ea..0b0a9038065 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy @@ -102,9 +102,49 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { taintedObjects[1] == null } + void 'returned string is tainted with source values as sql table and second value in the same row is tainted'() { + when: + List valuesRead = [] + List taintedObjects = [] + runUnderIastTrace { + String sql = 'SELECT name FROM TEST LIMIT 1' + Connection constWrapper = new IastInstrumentedConnection(conn: connection) + + constWrapper.prepareStatement(sql).withCloseable { stmt -> + stmt.executeQuery().withCloseable { rs -> + while (rs.next()) { + valuesRead.add(rs.getString("name")) + valuesRead.add(rs.getString("name")) + } + } + } + + taintedObjects.add(localTaintedObjects.get(valuesRead[0])) + taintedObjects.add(localTaintedObjects.get(valuesRead[1])) + } + + then: + valuesRead[0] == "foo" + taintedObjects[0] != null + with(taintedObjects[0]) { + with(ranges[0]) { + source.origin == SourceTypes.SQL_TABLE + source.value == valuesRead[0] + } + } + valuesRead[1] == "foo" + taintedObjects[1] != null + with(taintedObjects[1]) { + with(ranges[0]) { + source.origin == SourceTypes.SQL_TABLE + source.value == valuesRead[1] + } + } + } + void 'returned string is tainted with source values as sql table and can taint up to two values'() { given: - injectSysConfig('dd.iast.db.rows.to.taint', "2") + injectSysConfig('dd.iast.db.rows-to-taint', "2") when: List valuesRead = [] diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java index 06e24bf13a1..4108d0a796e 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java @@ -27,7 +27,7 @@ public final class IastConfig { public static final String IAST_SOURCE_MAPPING_MAX_SIZE = "iast.source-mapping.max-size"; public static final String IAST_EXPERIMENTAL_PROPAGATION_ENABLED = "iast.experimental.propagation.enabled"; - public static final String IAST_DB_ROWS_TO_TAINT = "iast.db.rows.to.taint"; + public static final String IAST_DB_ROWS_TO_TAINT = "iast.db.rows-to-taint"; public static final String IAST_STACK_TRACE_ENABLED = "iast.stacktrace.enabled"; From a12b842db6e65dd52a436478b79be4519191f595 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Thu, 12 Dec 2024 16:55:40 +0100 Subject: [PATCH 06/19] Remove exclusion --- .../agent/tooling/bytebuddy/matcher/ignored_class_name.trie | 1 - 1 file changed, 1 deletion(-) diff --git a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie index 7212f21c352..9372507c6b9 100644 --- a/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie +++ b/dd-java-agent/agent-tooling/src/main/resources/datadog/trace/agent/tooling/bytebuddy/matcher/ignored_class_name.trie @@ -42,7 +42,6 @@ 1 io.opentelemetry.javaagent.* 1 java.* 0 java.lang.ClassLoader -0 java.sql.ResultSet # allow exception profiling instrumentation 0 java.lang.Exception 0 java.lang.Error From bfac899a662fd1e5f7548a37895d952ffc9c0dd9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Thu, 12 Dec 2024 16:57:04 +0100 Subject: [PATCH 07/19] Add additional names to the instrumentation --- .../instrumentation/jdbc/IastResultSetInstrumentation.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index d0aebe42a6d..33bffd1b9bc 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -32,7 +32,7 @@ public class IastResultSetInstrumentation extends InstrumenterModule.Iast implements Instrumenter.ForTypeHierarchy { public IastResultSetInstrumentation() { - super("jdbc", "resultset"); + super("jdbc", "jdbc-resultset", "iast-resultset"); } @Override From da1e53cdcd44952e7f7eda4112ceb4b02dc6e73c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Mon, 16 Dec 2024 15:41:29 +0100 Subject: [PATCH 08/19] Add new test --- .../src/test/groovy/com/datadog/iast/taint/RangesTest.groovy | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy index 748e11cb7c1..9f695f9d155 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy @@ -12,6 +12,7 @@ import static com.datadog.iast.util.HttpHeader.LOCATION import static com.datadog.iast.util.HttpHeader.REFERER import static datadog.trace.api.iast.SourceTypes.GRPC_BODY import static datadog.trace.api.iast.SourceTypes.REQUEST_HEADER_VALUE +import static datadog.trace.api.iast.SourceTypes.REQUEST_QUERY import static datadog.trace.api.iast.SourceTypes.SQL_TABLE import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED import static com.datadog.iast.taint.Ranges.mergeRanges @@ -391,6 +392,7 @@ class RangesTest extends DDSpecification { ranges | source | expected [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | SQL_TABLE | [range(5, 3)] [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | REQUEST_HEADER_NAME | [rangeWithSource(0, 5, SQL_TABLE)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | REQUEST_QUERY | [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] [] | SQL_TABLE | [] } From d1abb9d92b26dee7e56750e7f8d63705d18d79f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Tue, 17 Dec 2024 11:41:27 +0100 Subject: [PATCH 09/19] Update env variable name --- .../src/test/groovy/IastResultSetInstrumentationTest.groovy | 2 +- .../src/main/java/datadog/trace/api/config/IastConfig.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy index 0b0a9038065..3caab122e8b 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy @@ -144,7 +144,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { void 'returned string is tainted with source values as sql table and can taint up to two values'() { given: - injectSysConfig('dd.iast.db.rows-to-taint', "2") + injectSysConfig('dd.iast.dbRowsToTaint', "2") when: List valuesRead = [] diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java index 4108d0a796e..125e789efd2 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java @@ -27,7 +27,7 @@ public final class IastConfig { public static final String IAST_SOURCE_MAPPING_MAX_SIZE = "iast.source-mapping.max-size"; public static final String IAST_EXPERIMENTAL_PROPAGATION_ENABLED = "iast.experimental.propagation.enabled"; - public static final String IAST_DB_ROWS_TO_TAINT = "iast.db.rows-to-taint"; + public static final String IAST_DB_ROWS_TO_TAINT = "iast.dbRowsToTaint"; public static final String IAST_STACK_TRACE_ENABLED = "iast.stacktrace.enabled"; From 0e3795c0961dd0d1e1084422b2b6f768580de6ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Thu, 26 Dec 2024 10:34:01 +0100 Subject: [PATCH 10/19] Add source name to tainted objects --- .../jdbc/IastResultSetInstrumentation.java | 7 ++++++- .../IastResultSetInstrumentationTest.groovy | 17 +++++++++++++---- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index 33bffd1b9bc..8939fd01aed 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -82,6 +82,7 @@ public static class GetParameterAdvice { @Advice.OnMethodExit(suppress = Throwable.class) @Source(SourceTypes.SQL_TABLE) public static void onExit( + @Advice.Argument(0) Object argument, @Advice.Return final String value, @Advice.This final ResultSet resultSet, @ActiveRequestContext RequestContext reqCtx) { @@ -98,7 +99,11 @@ public static void onExit( return; } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); - module.taintString(ctx, value, SourceTypes.SQL_TABLE); + if (argument instanceof String) { + module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument); + } else { + module.taintString(ctx, value, SourceTypes.SQL_TABLE); + } } } } diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy index 3caab122e8b..fe23548efb1 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy @@ -64,6 +64,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { with(taintedObject) { with(ranges[0]) { source.origin == SourceTypes.SQL_TABLE + source.name == null source.value == valueRead } } @@ -72,6 +73,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { void 'returned string is tainted with source values as sql table but second value is not tainted'() { when: List valuesRead = [] + String column = "name" List taintedObjects = [] runUnderIastTrace { String sql = 'SELECT name FROM TEST LIMIT 2' @@ -80,7 +82,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { constWrapper.prepareStatement(sql).withCloseable { stmt -> stmt.executeQuery().withCloseable { rs -> while (rs.next()) { - valuesRead.add(rs.getString("name")) + valuesRead.add(rs.getString(column)) } } } @@ -95,6 +97,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { with(taintedObjects[0]) { with(ranges[0]) { source.origin == SourceTypes.SQL_TABLE + source.name == column source.value == valuesRead[0] } } @@ -105,6 +108,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { void 'returned string is tainted with source values as sql table and second value in the same row is tainted'() { when: List valuesRead = [] + String column = "name" List taintedObjects = [] runUnderIastTrace { String sql = 'SELECT name FROM TEST LIMIT 1' @@ -113,8 +117,8 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { constWrapper.prepareStatement(sql).withCloseable { stmt -> stmt.executeQuery().withCloseable { rs -> while (rs.next()) { - valuesRead.add(rs.getString("name")) - valuesRead.add(rs.getString("name")) + valuesRead.add(rs.getString(column)) + valuesRead.add(rs.getString(column)) } } } @@ -129,6 +133,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { with(taintedObjects[0]) { with(ranges[0]) { source.origin == SourceTypes.SQL_TABLE + source.name == column source.value == valuesRead[0] } } @@ -137,6 +142,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { with(taintedObjects[1]) { with(ranges[0]) { source.origin == SourceTypes.SQL_TABLE + source.name == column source.value == valuesRead[1] } } @@ -148,6 +154,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { when: List valuesRead = [] + String column = "name" List taintedObjects = [] runUnderIastTrace { String sql = 'SELECT name FROM TEST LIMIT 2' @@ -156,7 +163,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { constWrapper.prepareStatement(sql).withCloseable { stmt -> stmt.executeQuery().withCloseable { rs -> while (rs.next()) { - valuesRead.add(rs.getString("name")) + valuesRead.add(rs.getString(column)) } } } @@ -171,6 +178,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { with(taintedObjects[0]) { with(ranges[0]) { source.origin == SourceTypes.SQL_TABLE + source.name == column source.value == valuesRead[0] } } @@ -179,6 +187,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { with(taintedObjects[1]) { with(ranges[0]) { source.origin == SourceTypes.SQL_TABLE + source.name == column source.value == valuesRead[1] } } From f808e6ac513f0ea5e34df92393c2e54a37713e33 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Fri, 27 Dec 2024 11:48:41 +0100 Subject: [PATCH 11/19] Cover edge case of nested calls of getString --- .../propagation/PropagationModuleImpl.java | 60 +++++++++++++++++++ .../java/com/datadog/iast/taint/Ranges.java | 14 +++++ .../jdbc/IastResultSetInstrumentation.java | 9 ++- .../iast/propagation/PropagationModule.java | 7 +++ 4 files changed, 87 insertions(+), 3 deletions(-) diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java index 1503fd12909..4c65c314295 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java @@ -1,6 +1,7 @@ package com.datadog.iast.propagation; import static com.datadog.iast.model.Source.PROPAGATION_PLACEHOLDER; +import static com.datadog.iast.taint.Ranges.changeHighestPriorityRange; import static com.datadog.iast.taint.Ranges.highestPriorityRange; import static com.datadog.iast.util.ObjectVisitor.State.CONTINUE; import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED; @@ -668,6 +669,39 @@ public void markIfTainted(@Nullable Object target, int mark) { } } + @Override + public void changeSource(@Nullable Object target, byte origin) { + if (target == null) { + return; + } + changeSource(target, origin, null); + } + + @Override + public void changeSource(@Nullable Object target, byte origin, @Nullable CharSequence name) { + if (target == null) { + return; + } + final IastContext ctx = IastContext.Provider.get(); + if (ctx == null) { + return; + } + changeSource(ctx, target, origin, name); + } + + @Override + public void changeSource( + @Nullable final IastContext ctx, + @Nullable Object target, + byte origin, + @Nullable CharSequence name) { + if (ctx == null || target == null) { + return; + } + final TaintedObjects to = ctx.getTaintedObjects(); + changeHighestPrioritySource(to, target, origin, name); + } + @Override public boolean isTainted(@Nullable final Object target) { if (target == null) { @@ -784,6 +818,32 @@ private static Source highestPrioritySource( } } + private static void changeHighestPrioritySource( + final @Nonnull TaintedObjects to, + final @Nonnull Object object, + final byte origin, + @Nullable final CharSequence name) { + Source previousValue = highestPrioritySource(to, object); + if (previousValue == null) { + return; + } + Source newSource = newSource(object, origin, name, previousValue.getValue()); + if (object instanceof Taintable) { + ((Taintable) object).$$DD$setSource(newSource); + } else { + TaintedObject taintedObject = to.get(object); + if (taintedObject == null) { + return; + } + final Range[] ranges = getRanges(to, object); + if (ranges == null || ranges.length == 0) { + return; + } + changeHighestPriorityRange(ranges, newSource); + taintedObject.setRanges(ranges); + } + } + private static void internalTaint( @Nonnull final TaintedObjects to, @Nonnull final Object value, diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java index 483307f9e0d..fa04e083a53 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java @@ -144,6 +144,20 @@ public static Range highestPriorityRange(@Nonnull final Range[] ranges) { return ranges[0]; } + public static void changeHighestPriorityRange( + @Nonnull final Range[] ranges, @Nonnull final Source source) { + for (int i = 0; i < ranges.length; i++) { + if (ranges[i].getMarks() == NOT_MARKED) { + Range newRange = + new Range(ranges[i].getStart(), ranges[i].getLength(), source, ranges[i].getMarks()); + ranges[i] = newRange; + } + } + Range newRange = + new Range(ranges[0].getStart(), ranges[0].getLength(), source, ranges[0].getMarks()); + ranges[0] = newRange; + } + /** * Checks if all ranges are coming from any header, in case no ranges are provided it will return * {@code true} diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index 8939fd01aed..44249dee15b 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -64,7 +64,6 @@ public Map contextStore() { public static class NextAdvice { @Advice.OnMethodExit(suppress = Throwable.class) - @Source(SourceTypes.SQL_TABLE) public static void onExit(@Advice.This final ResultSet resultSet) { ContextStore contextStore = InstrumentationContext.get(ResultSet.class, Integer.class); @@ -100,9 +99,13 @@ public static void onExit( } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); if (argument instanceof String) { - module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument); + if (module.isTainted(value)) { + module.changeSource(ctx, value, SourceTypes.SQL_TABLE, (String) argument); + } else { + module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument); + } } else { - module.taintString(ctx, value, SourceTypes.SQL_TABLE); + module.taintString(ctx, value, SourceTypes.SQL_TABLE, String.valueOf(argument)); } } } diff --git a/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java b/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java index 371c3618835..30285dc8069 100644 --- a/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java +++ b/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java @@ -347,4 +347,11 @@ int taintObjectDeeply( Source findSource(@Nullable IastContext ctx, @Nullable Object target); void markIfTainted(@Nullable Object target, int mark); + + void changeSource(@Nullable Object target, byte origin); + + void changeSource(@Nullable Object target, byte origin, @Nullable CharSequence name); + + void changeSource( + @Nullable IastContext ctx, @Nullable Object target, byte origin, @Nullable CharSequence name); } From 75b3a8b07f2898ed49c088fb340ece97553f3909 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Fri, 27 Dec 2024 16:23:39 +0100 Subject: [PATCH 12/19] Remove unnecesary naming --- .../instrumentation/jdbc/IastResultSetInstrumentation.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index 44249dee15b..f3801de0ed7 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -105,7 +105,7 @@ public static void onExit( module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument); } } else { - module.taintString(ctx, value, SourceTypes.SQL_TABLE, String.valueOf(argument)); + module.taintString(ctx, value, SourceTypes.SQL_TABLE); } } } From 6ea517de20a24b4431f569df805f11089e60d350 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Tue, 7 Jan 2025 12:20:33 +0100 Subject: [PATCH 13/19] Replace changing source for CallDepthThreadLocalMap approach --- .../propagation/PropagationModuleImpl.java | 60 ------------------- .../java/com/datadog/iast/taint/Ranges.java | 14 ----- .../jdbc/IastResultSetInstrumentation.java | 15 +++-- .../iast/propagation/PropagationModule.java | 7 --- 4 files changed, 10 insertions(+), 86 deletions(-) diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java index 4c65c314295..1503fd12909 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/PropagationModuleImpl.java @@ -1,7 +1,6 @@ package com.datadog.iast.propagation; import static com.datadog.iast.model.Source.PROPAGATION_PLACEHOLDER; -import static com.datadog.iast.taint.Ranges.changeHighestPriorityRange; import static com.datadog.iast.taint.Ranges.highestPriorityRange; import static com.datadog.iast.util.ObjectVisitor.State.CONTINUE; import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED; @@ -669,39 +668,6 @@ public void markIfTainted(@Nullable Object target, int mark) { } } - @Override - public void changeSource(@Nullable Object target, byte origin) { - if (target == null) { - return; - } - changeSource(target, origin, null); - } - - @Override - public void changeSource(@Nullable Object target, byte origin, @Nullable CharSequence name) { - if (target == null) { - return; - } - final IastContext ctx = IastContext.Provider.get(); - if (ctx == null) { - return; - } - changeSource(ctx, target, origin, name); - } - - @Override - public void changeSource( - @Nullable final IastContext ctx, - @Nullable Object target, - byte origin, - @Nullable CharSequence name) { - if (ctx == null || target == null) { - return; - } - final TaintedObjects to = ctx.getTaintedObjects(); - changeHighestPrioritySource(to, target, origin, name); - } - @Override public boolean isTainted(@Nullable final Object target) { if (target == null) { @@ -818,32 +784,6 @@ private static Source highestPrioritySource( } } - private static void changeHighestPrioritySource( - final @Nonnull TaintedObjects to, - final @Nonnull Object object, - final byte origin, - @Nullable final CharSequence name) { - Source previousValue = highestPrioritySource(to, object); - if (previousValue == null) { - return; - } - Source newSource = newSource(object, origin, name, previousValue.getValue()); - if (object instanceof Taintable) { - ((Taintable) object).$$DD$setSource(newSource); - } else { - TaintedObject taintedObject = to.get(object); - if (taintedObject == null) { - return; - } - final Range[] ranges = getRanges(to, object); - if (ranges == null || ranges.length == 0) { - return; - } - changeHighestPriorityRange(ranges, newSource); - taintedObject.setRanges(ranges); - } - } - private static void internalTaint( @Nonnull final TaintedObjects to, @Nonnull final Object value, diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java index fa04e083a53..483307f9e0d 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java @@ -144,20 +144,6 @@ public static Range highestPriorityRange(@Nonnull final Range[] ranges) { return ranges[0]; } - public static void changeHighestPriorityRange( - @Nonnull final Range[] ranges, @Nonnull final Source source) { - for (int i = 0; i < ranges.length; i++) { - if (ranges[i].getMarks() == NOT_MARKED) { - Range newRange = - new Range(ranges[i].getStart(), ranges[i].getLength(), source, ranges[i].getMarks()); - ranges[i] = newRange; - } - } - Range newRange = - new Range(ranges[0].getStart(), ranges[0].getLength(), source, ranges[0].getMarks()); - ranges[0] = newRange; - } - /** * Checks if all ranges are coming from any header, in case no ranges are provided it will return * {@code true} diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index f3801de0ed7..8a843569a54 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -19,6 +19,7 @@ import datadog.trace.api.iast.Source; import datadog.trace.api.iast.SourceTypes; import datadog.trace.api.iast.propagation.PropagationModule; +import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.InstrumentationContext; import java.sql.ResultSet; @@ -78,6 +79,11 @@ public static void onExit(@Advice.This final ResultSet resultSet) { @RequiresRequestContext(RequestContextSlot.IAST) public static class GetParameterAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static void onEnter() { + CallDepthThreadLocalMap.incrementCallDepth(ResultSet.class); + } + @Advice.OnMethodExit(suppress = Throwable.class) @Source(SourceTypes.SQL_TABLE) public static void onExit( @@ -85,6 +91,9 @@ public static void onExit( @Advice.Return final String value, @Advice.This final ResultSet resultSet, @ActiveRequestContext RequestContext reqCtx) { + if (CallDepthThreadLocalMap.decrementCallDepth(ResultSet.class) > 0) { + return; + } ContextStore contextStore = InstrumentationContext.get(ResultSet.class, Integer.class); if (contextStore.get(resultSet) > Config.get().getIastDbRowsToTaint()) { @@ -99,11 +108,7 @@ public static void onExit( } IastContext ctx = reqCtx.getData(RequestContextSlot.IAST); if (argument instanceof String) { - if (module.isTainted(value)) { - module.changeSource(ctx, value, SourceTypes.SQL_TABLE, (String) argument); - } else { - module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument); - } + module.taintString(ctx, value, SourceTypes.SQL_TABLE, (String) argument); } else { module.taintString(ctx, value, SourceTypes.SQL_TABLE); } diff --git a/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java b/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java index 30285dc8069..371c3618835 100644 --- a/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java +++ b/internal-api/src/main/java/datadog/trace/api/iast/propagation/PropagationModule.java @@ -347,11 +347,4 @@ int taintObjectDeeply( Source findSource(@Nullable IastContext ctx, @Nullable Object target); void markIfTainted(@Nullable Object target, int mark); - - void changeSource(@Nullable Object target, byte origin); - - void changeSource(@Nullable Object target, byte origin, @Nullable CharSequence name); - - void changeSource( - @Nullable IastContext ctx, @Nullable Object target, byte origin, @Nullable CharSequence name); } From f250faf0205faa8257ada09a56656d0012c39584 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Tue, 7 Jan 2025 12:35:38 +0100 Subject: [PATCH 14/19] Fix build error --- .../instrumentation/jdbc/IastResultSetInstrumentation.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java index 8a843569a54..5a6447613ba 100644 --- a/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java +++ b/dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/IastResultSetInstrumentation.java @@ -30,7 +30,7 @@ @AutoService(InstrumenterModule.class) public class IastResultSetInstrumentation extends InstrumenterModule.Iast - implements Instrumenter.ForTypeHierarchy { + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { public IastResultSetInstrumentation() { super("jdbc", "jdbc-resultset", "iast-resultset"); From ed39adf7830abcf91ed421166b5dfa3a39b5dc75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Fri, 10 Jan 2025 12:43:24 +0100 Subject: [PATCH 15/19] Rename ENV Variable and add more coverage to tests --- .../IastResultSetInstrumentationTest.groovy | 37 ++++++++++++++++++- .../datadog/trace/api/config/IastConfig.java | 2 +- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy index fe23548efb1..bcad0b535f8 100644 --- a/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/jdbc/src/test/groovy/IastResultSetInstrumentationTest.groovy @@ -1,11 +1,15 @@ import com.datadog.iast.taint.TaintedObject import com.datadog.iast.test.IastAgentTestRunner import datadog.trace.api.iast.SourceTypes +import datadog.trace.bootstrap.CallDepthThreadLocalMap import foo.bar.IastInstrumentedConnection import spock.lang.Shared import javax.sql.DataSource import java.sql.Connection +import java.sql.ResultSet + +import static datadog.trace.api.config.IastConfig.IAST_DB_ROWS_TO_TAINT class IastResultSetInstrumentationTest extends IastAgentTestRunner { @Shared @@ -59,6 +63,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { } then: + CallDepthThreadLocalMap.getCallDepth(ResultSet) == 0 valueRead == "foo" taintedObject != null with(taintedObject) { @@ -92,6 +97,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { } then: + CallDepthThreadLocalMap.getCallDepth(ResultSet) == 0 valuesRead[0] == "foo" taintedObjects[0] != null with(taintedObjects[0]) { @@ -128,6 +134,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { } then: + CallDepthThreadLocalMap.getCallDepth(ResultSet) == 0 valuesRead[0] == "foo" taintedObjects[0] != null with(taintedObjects[0]) { @@ -150,7 +157,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { void 'returned string is tainted with source values as sql table and can taint up to two values'() { given: - injectSysConfig('dd.iast.dbRowsToTaint', "2") + injectSysConfig(IAST_DB_ROWS_TO_TAINT, "2") when: List valuesRead = [] @@ -173,6 +180,7 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { } then: + CallDepthThreadLocalMap.getCallDepth(ResultSet) == 0 valuesRead[0] == "foo" taintedObjects[0] != null with(taintedObjects[0]) { @@ -213,7 +221,34 @@ class IastResultSetInstrumentationTest extends IastAgentTestRunner { } then: + CallDepthThreadLocalMap.getCallDepth(ResultSet) == 0 valueRead == 42 taintedObject == null } + + void 'when CallDepthThreadLocalMap is greater than 0 does not taint the values' () { + when: + def valueRead + TaintedObject taintedObject + CallDepthThreadLocalMap.incrementCallDepth(ResultSet) + runUnderIastTrace { + String sql = 'SELECT name FROM TEST LIMIT 1' + Connection constWrapper = new IastInstrumentedConnection(conn: connection) + + constWrapper.prepareStatement(sql).withCloseable { stmt -> + stmt.executeQuery().withCloseable { rs -> + if (rs.next()) { + valueRead = rs.getString(1) + } + } + } + + taintedObject = localTaintedObjects.get(valueRead) + } + + then: + CallDepthThreadLocalMap.getCallDepth(ResultSet) == 1 + valueRead == "foo" + taintedObject == null + } } diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java index 42d257d0c42..6abe05a046f 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java +++ b/dd-trace-api/src/main/java/datadog/trace/api/config/IastConfig.java @@ -30,7 +30,7 @@ public final class IastConfig { public static final String IAST_SECURITY_CONTROLS_ENABLED = "iast.security-controls.enabled"; public static final String IAST_SECURITY_CONTROLS_CONFIGURATION = "iast.security-controls.configuration"; - public static final String IAST_DB_ROWS_TO_TAINT = "iast.dbRowsToTaint"; + public static final String IAST_DB_ROWS_TO_TAINT = "iast.db.rows-to-taint"; private IastConfig() {} } From d9b5d81b0eb7249ddc56915fe74a2d522ca03de7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Fri, 10 Jan 2025 14:42:35 +0100 Subject: [PATCH 16/19] Refactor excluded sources to be more generic --- .../datadog/iast/model/VulnerabilityType.java | 126 ++++++++++++++---- .../com/datadog/iast/sink/SinkModuleBase.java | 11 +- .../java/com/datadog/iast/taint/Ranges.java | 17 ++- .../com/datadog/iast/taint/RangesTest.groovy | 14 +- 4 files changed, 128 insertions(+), 40 deletions(-) diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java index f9a05fec280..25dc91d2bd8 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java @@ -15,6 +15,7 @@ import static datadog.trace.api.iast.VulnerabilityMarks.XPATH_INJECTION_MARK; import static datadog.trace.api.iast.VulnerabilityMarks.XSS_MARK; +import datadog.trace.api.iast.SourceTypes; import datadog.trace.api.iast.VulnerabilityTypes; import java.io.File; import java.util.function.BiFunction; @@ -23,90 +24,145 @@ public interface VulnerabilityType { - VulnerabilityType WEAK_CIPHER = type(VulnerabilityTypes.WEAK_CIPHER).build(); - VulnerabilityType WEAK_HASH = type(VulnerabilityTypes.WEAK_HASH).build(); + VulnerabilityType WEAK_CIPHER = + type(VulnerabilityTypes.WEAK_CIPHER).excludedSources(SourceTypes.SQL_TABLE).build(); + VulnerabilityType WEAK_HASH = + type(VulnerabilityTypes.WEAK_HASH).excludedSources(SourceTypes.SQL_TABLE).build(); VulnerabilityType INSECURE_COOKIE = - type(VulnerabilityTypes.INSECURE_COOKIE).hash(VulnerabilityType::evidenceHash).build(); + type(VulnerabilityTypes.INSECURE_COOKIE) + .hash(VulnerabilityType::evidenceHash) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType NO_HTTPONLY_COOKIE = - type(VulnerabilityTypes.NO_HTTPONLY_COOKIE).hash(VulnerabilityType::evidenceHash).build(); + type(VulnerabilityTypes.NO_HTTPONLY_COOKIE) + .hash(VulnerabilityType::evidenceHash) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType HSTS_HEADER_MISSING = - type(VulnerabilityTypes.HSTS_HEADER_MISSING).hash(VulnerabilityType::serviceHash).build(); + type(VulnerabilityTypes.HSTS_HEADER_MISSING) + .hash(VulnerabilityType::serviceHash) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType XCONTENTTYPE_HEADER_MISSING = type(VulnerabilityTypes.XCONTENTTYPE_HEADER_MISSING) .hash(VulnerabilityType::serviceHash) + .excludedSources(SourceTypes.SQL_TABLE) .build(); VulnerabilityType NO_SAMESITE_COOKIE = - type(VulnerabilityTypes.NO_SAMESITE_COOKIE).hash(VulnerabilityType::evidenceHash).build(); + type(VulnerabilityTypes.NO_SAMESITE_COOKIE) + .hash(VulnerabilityType::evidenceHash) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType SQL_INJECTION = type(VulnerabilityTypes.SQL_INJECTION).mark(SQL_INJECTION_MARK).build(); VulnerabilityType COMMAND_INJECTION = - type(VulnerabilityTypes.COMMAND_INJECTION).mark(COMMAND_INJECTION_MARK).build(); + type(VulnerabilityTypes.COMMAND_INJECTION) + .mark(COMMAND_INJECTION_MARK) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType PATH_TRAVERSAL = type(VulnerabilityTypes.PATH_TRAVERSAL) .separator(File.separatorChar) .mark(PATH_TRAVERSAL_MARK) + .excludedSources(SourceTypes.SQL_TABLE) .build(); VulnerabilityType LDAP_INJECTION = - type(VulnerabilityTypes.LDAP_INJECTION).mark(LDAP_INJECTION_MARK).build(); - VulnerabilityType SSRF = type(VulnerabilityTypes.SSRF).mark(SSRF_MARK).build(); + type(VulnerabilityTypes.LDAP_INJECTION) + .mark(LDAP_INJECTION_MARK) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); + VulnerabilityType SSRF = + type(VulnerabilityTypes.SSRF).mark(SSRF_MARK).excludedSources(SourceTypes.SQL_TABLE).build(); VulnerabilityType UNVALIDATED_REDIRECT = - type(VulnerabilityTypes.UNVALIDATED_REDIRECT).mark(UNVALIDATED_REDIRECT_MARK).build(); - VulnerabilityType WEAK_RANDOMNESS = type(VulnerabilityTypes.WEAK_RANDOMNESS).build(); + type(VulnerabilityTypes.UNVALIDATED_REDIRECT) + .mark(UNVALIDATED_REDIRECT_MARK) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); + VulnerabilityType WEAK_RANDOMNESS = + type(VulnerabilityTypes.WEAK_RANDOMNESS).excludedSources(SourceTypes.SQL_TABLE).build(); VulnerabilityType XPATH_INJECTION = - type(VulnerabilityTypes.XPATH_INJECTION).mark(XPATH_INJECTION_MARK).build(); + type(VulnerabilityTypes.XPATH_INJECTION) + .mark(XPATH_INJECTION_MARK) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType TRUST_BOUNDARY_VIOLATION = - type(VulnerabilityTypes.TRUST_BOUNDARY_VIOLATION).mark(TRUST_BOUNDARY_VIOLATION_MARK).build(); + type(VulnerabilityTypes.TRUST_BOUNDARY_VIOLATION) + .mark(TRUST_BOUNDARY_VIOLATION_MARK) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType XSS = type(VulnerabilityTypes.XSS).mark(XSS_MARK).build(); VulnerabilityType HEADER_INJECTION = - type(VulnerabilityTypes.HEADER_INJECTION).mark(HEADER_INJECTION_MARK).build(); + type(VulnerabilityTypes.HEADER_INJECTION) + .mark(HEADER_INJECTION_MARK) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); - VulnerabilityType STACKTRACE_LEAK = type(VulnerabilityTypes.STACKTRACE_LEAK).build(); + VulnerabilityType STACKTRACE_LEAK = + type(VulnerabilityTypes.STACKTRACE_LEAK).excludedSources(SourceTypes.SQL_TABLE).build(); - VulnerabilityType VERB_TAMPERING = type(VulnerabilityTypes.VERB_TAMPERING).build(); + VulnerabilityType VERB_TAMPERING = + type(VulnerabilityTypes.VERB_TAMPERING).excludedSources(SourceTypes.SQL_TABLE).build(); VulnerabilityType ADMIN_CONSOLE_ACTIVE = type(VulnerabilityTypes.ADMIN_CONSOLE_ACTIVE) .deduplicable(false) .hash(VulnerabilityType::serviceHash) + .excludedSources(SourceTypes.SQL_TABLE) .build(); VulnerabilityType DEFAULT_HTML_ESCAPE_INVALID = - type(VulnerabilityTypes.DEFAULT_HTML_ESCAPE_INVALID).build(); + type(VulnerabilityTypes.DEFAULT_HTML_ESCAPE_INVALID) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); - VulnerabilityType SESSION_TIMEOUT = type(VulnerabilityTypes.SESSION_TIMEOUT).build(); + VulnerabilityType SESSION_TIMEOUT = + type(VulnerabilityTypes.SESSION_TIMEOUT).excludedSources(SourceTypes.SQL_TABLE).build(); VulnerabilityType DIRECTORY_LISTING_LEAK = - type(VulnerabilityTypes.DIRECTORY_LISTING_LEAK).build(); - VulnerabilityType INSECURE_JSP_LAYOUT = type(VulnerabilityTypes.INSECURE_JSP_LAYOUT).build(); + type(VulnerabilityTypes.DIRECTORY_LISTING_LEAK) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); + VulnerabilityType INSECURE_JSP_LAYOUT = + type(VulnerabilityTypes.INSECURE_JSP_LAYOUT).excludedSources(SourceTypes.SQL_TABLE).build(); - VulnerabilityType HARDCODED_SECRET = type(VulnerabilityTypes.HARDCODED_SECRET).build(); + VulnerabilityType HARDCODED_SECRET = + type(VulnerabilityTypes.HARDCODED_SECRET).excludedSources(SourceTypes.SQL_TABLE).build(); VulnerabilityType INSECURE_AUTH_PROTOCOL = - type(VulnerabilityTypes.INSECURE_AUTH_PROTOCOL).hash(VulnerabilityType::evidenceHash).build(); + type(VulnerabilityTypes.INSECURE_AUTH_PROTOCOL) + .hash(VulnerabilityType::evidenceHash) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType REFLECTION_INJECTION = - type(VulnerabilityTypes.REFLECTION_INJECTION).mark(REFLECTION_INJECTION_MARK).build(); + type(VulnerabilityTypes.REFLECTION_INJECTION) + .mark(REFLECTION_INJECTION_MARK) + .excludedSources(SourceTypes.SQL_TABLE) + .build(); VulnerabilityType SESSION_REWRITING = type(VulnerabilityTypes.SESSION_REWRITING) .deduplicable(false) .hash(VulnerabilityType::serviceHash) + .excludedSources(SourceTypes.SQL_TABLE) .build(); VulnerabilityType DEFAULT_APP_DEPLOYED = type(VulnerabilityTypes.DEFAULT_APP_DEPLOYED) .deduplicable(false) .hash(VulnerabilityType::serviceHash) + .excludedSources(SourceTypes.SQL_TABLE) .build(); VulnerabilityType UNTRUSTED_DESERIALIZATION = type(VulnerabilityTypes.UNTRUSTED_DESERIALIZATION) .mark(UNTRUSTED_DESERIALIZATION_MARK) + .excludedSources(SourceTypes.SQL_TABLE) .build(); /* All vulnerability types that have a mark. Should be updated if new vulnerabilityType with mark is added */ @@ -137,6 +193,8 @@ public interface VulnerabilityType { /** A flag to indicate if the vulnerability is deduplicable. */ boolean isDeduplicable(); + byte[] excludedSources(); + static Builder type(final byte type) { return new Builder(type); } @@ -151,6 +209,8 @@ class VulnerabilityTypeImpl implements VulnerabilityType { private final boolean deduplicable; + private final byte[] excludedSources; + private final BiFunction hash; public VulnerabilityTypeImpl( @@ -158,11 +218,13 @@ public VulnerabilityTypeImpl( final char separator, final int mark, final boolean deduplicable, + final byte[] excludedSources, final BiFunction hash) { this.type = type; this.separator = separator; this.mark = mark; this.deduplicable = deduplicable; + this.excludedSources = excludedSources; this.hash = hash; } @@ -191,6 +253,11 @@ public boolean isDeduplicable() { return deduplicable; } + @Override + public byte[] excludedSources() { + return excludedSources; + } + /** Useful for troubleshooting issues when vulns are serialized without moshi */ public String getName() { return name(); @@ -202,6 +269,7 @@ class Builder { private char separator = ' '; private int mark = NOT_MARKED; private boolean deduplicable = true; + private byte[] excludedSources = new byte[0]; private BiFunction hash = VulnerabilityType::fileAndLineHash; @@ -224,13 +292,23 @@ public Builder deduplicable(final boolean deduplicable) { return this; } + public Builder excludedSources(final byte[] excludedSources) { + this.excludedSources = excludedSources; + return this; + } + + public Builder excludedSources(final byte excludedSources) { + this.excludedSources = new byte[] {excludedSources}; + return this; + } + public Builder hash(final BiFunction hash) { this.hash = hash; return this; } public VulnerabilityType build() { - return new VulnerabilityTypeImpl(type, separator, mark, deduplicable, hash); + return new VulnerabilityTypeImpl(type, separator, mark, deduplicable, excludedSources, hash); } } diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java index ee9a2fac8ad..0fef8a16824 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java @@ -21,16 +21,13 @@ import datadog.trace.api.Config; import datadog.trace.api.Pair; import datadog.trace.api.iast.IastContext; -import datadog.trace.api.iast.SourceTypes; import datadog.trace.api.iast.Taintable; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.instrumentation.iastinstrumenter.IastExclusionTrie; import datadog.trace.instrumentation.iastinstrumenter.SourceMapperImpl; import datadog.trace.util.stacktrace.StackWalker; -import java.util.Arrays; import java.util.Iterator; -import java.util.List; import java.util.function.Predicate; import java.util.stream.Stream; import javax.annotation.Nonnull; @@ -42,8 +39,6 @@ public abstract class SinkModuleBase { private static final int MAX_EVIDENCE_LENGTH = Config.get().getIastTruncationMaxValueLength(); - private static final List DB_INJECTION_TYPES = - Arrays.asList(VulnerabilityType.SQL_INJECTION, VulnerabilityType.XSS); protected final OverheadController overheadController; protected final Reporter reporter; @@ -148,10 +143,10 @@ protected final Evidence checkInjection( return null; } - // filter out ranges that are not SQL_TABLE for types that are not DB_INJECTION_TYPES + // filter excluded ranges final Range[] filteredRanges; - if (!DB_INJECTION_TYPES.contains(type)) { - filteredRanges = Ranges.excludeRangesBySource(valueRanges, SourceTypes.SQL_TABLE); + if (type.excludedSources().length > 0) { + filteredRanges = Ranges.excludeRangesBySource(valueRanges, type.excludedSources()); } else { filteredRanges = valueRanges; } diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java index 483307f9e0d..d8bbec33ce8 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java @@ -433,15 +433,28 @@ public static Range[] splitRanges( * @param ranges the ranges to filter * @param source the byte value of the source to exclude (see {@link SourceTypes}) */ - public static Range[] excludeRangesBySource(Range[] ranges, byte source) { + public static Range[] excludeRangesBySource(Range[] ranges, byte[] source) { RangeBuilder newRanges = new RangeBuilder(ranges.length); for (Range range : ranges) { - if (range.getSource().getOrigin() != source) { + boolean exclude = false; + + for (byte origin : source) { + if (range.getSource().getOrigin() == origin) { + exclude = true; + break; + } + } + + if (!exclude) { newRanges.add(range); } } return newRanges.toArray(); } + + public static Range[] excludeRangesBySource(Range[] ranges, byte source) { + return excludeRangesBySource(ranges, new byte[] {source}); + } } diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy index 9f695f9d155..fc63d979a13 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy @@ -382,18 +382,20 @@ class RangesTest extends DDSpecification { void 'test excludeRangesBySource method'() { when: - final result = Ranges.excludeRangesBySource(ranges as Range[], source) + final result = Ranges.excludeRangesBySource(ranges as Range[], source as byte[]) then: final expectedArray = expected as Range[] result == expectedArray where: - ranges | source | expected - [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | SQL_TABLE | [range(5, 3)] - [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | REQUEST_HEADER_NAME | [rangeWithSource(0, 5, SQL_TABLE)] - [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | REQUEST_QUERY | [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] - [] | SQL_TABLE | [] + ranges | source | expected + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [SQL_TABLE] | [range(5, 3)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [SQL_TABLE, REQUEST_QUERY] | [range(5, 3)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [REQUEST_HEADER_NAME] | [rangeWithSource(0, 5, SQL_TABLE)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [REQUEST_QUERY] | [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [REQUEST_QUERY, REQUEST_HEADER_NAME] | [rangeWithSource(0, 5, SQL_TABLE)] + [] | [SQL_TABLE] | [] } Range[] rangesFromSpec(List> spec) { From 58997111187ce8c1de4a3807eadfacd98aa9342c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Fri, 10 Jan 2025 14:48:57 +0100 Subject: [PATCH 17/19] Fix Code Quality Violation --- .../main/java/com/datadog/iast/model/VulnerabilityType.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java index 25dc91d2bd8..7a613770980 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java @@ -255,7 +255,7 @@ public boolean isDeduplicable() { @Override public byte[] excludedSources() { - return excludedSources; + return excludedSources.clone(); } /** Useful for troubleshooting issues when vulns are serialized without moshi */ @@ -293,7 +293,7 @@ public Builder deduplicable(final boolean deduplicable) { } public Builder excludedSources(final byte[] excludedSources) { - this.excludedSources = excludedSources; + this.excludedSources = excludedSources.clone(); return this; } From cc3f5efacdd91154f615ad40500fdb369f68a004 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Wed, 15 Jan 2025 12:40:57 +0100 Subject: [PATCH 18/19] Modify byte[] to BitSet --- .../datadog/iast/model/VulnerabilityType.java | 86 +++++++++---------- .../com/datadog/iast/sink/SinkModuleBase.java | 2 +- .../java/com/datadog/iast/taint/Ranges.java | 18 +--- 3 files changed, 44 insertions(+), 62 deletions(-) diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java index 7a613770980..73fc8a6147d 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java @@ -18,40 +18,43 @@ import datadog.trace.api.iast.SourceTypes; import datadog.trace.api.iast.VulnerabilityTypes; import java.io.File; +import java.util.BitSet; import java.util.function.BiFunction; import java.util.zip.CRC32; import javax.annotation.Nonnull; public interface VulnerabilityType { + BitSet DB_EXCLUDED = new BitSet(SourceTypes.SQL_TABLE); + VulnerabilityType WEAK_CIPHER = - type(VulnerabilityTypes.WEAK_CIPHER).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.WEAK_CIPHER).excludedSources(DB_EXCLUDED).build(); VulnerabilityType WEAK_HASH = - type(VulnerabilityTypes.WEAK_HASH).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.WEAK_HASH).excludedSources(DB_EXCLUDED).build(); VulnerabilityType INSECURE_COOKIE = type(VulnerabilityTypes.INSECURE_COOKIE) .hash(VulnerabilityType::evidenceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType NO_HTTPONLY_COOKIE = type(VulnerabilityTypes.NO_HTTPONLY_COOKIE) .hash(VulnerabilityType::evidenceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType HSTS_HEADER_MISSING = type(VulnerabilityTypes.HSTS_HEADER_MISSING) .hash(VulnerabilityType::serviceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType XCONTENTTYPE_HEADER_MISSING = type(VulnerabilityTypes.XCONTENTTYPE_HEADER_MISSING) .hash(VulnerabilityType::serviceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType NO_SAMESITE_COOKIE = type(VulnerabilityTypes.NO_SAMESITE_COOKIE) .hash(VulnerabilityType::evidenceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType SQL_INJECTION = @@ -59,39 +62,39 @@ public interface VulnerabilityType { VulnerabilityType COMMAND_INJECTION = type(VulnerabilityTypes.COMMAND_INJECTION) .mark(COMMAND_INJECTION_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType PATH_TRAVERSAL = type(VulnerabilityTypes.PATH_TRAVERSAL) .separator(File.separatorChar) .mark(PATH_TRAVERSAL_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType LDAP_INJECTION = type(VulnerabilityTypes.LDAP_INJECTION) .mark(LDAP_INJECTION_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType SSRF = - type(VulnerabilityTypes.SSRF).mark(SSRF_MARK).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.SSRF).mark(SSRF_MARK).excludedSources(DB_EXCLUDED).build(); VulnerabilityType UNVALIDATED_REDIRECT = type(VulnerabilityTypes.UNVALIDATED_REDIRECT) .mark(UNVALIDATED_REDIRECT_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType WEAK_RANDOMNESS = - type(VulnerabilityTypes.WEAK_RANDOMNESS).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.WEAK_RANDOMNESS).excludedSources(DB_EXCLUDED).build(); VulnerabilityType XPATH_INJECTION = type(VulnerabilityTypes.XPATH_INJECTION) .mark(XPATH_INJECTION_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType TRUST_BOUNDARY_VIOLATION = type(VulnerabilityTypes.TRUST_BOUNDARY_VIOLATION) .mark(TRUST_BOUNDARY_VIOLATION_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType XSS = type(VulnerabilityTypes.XSS).mark(XSS_MARK).build(); @@ -99,70 +102,66 @@ public interface VulnerabilityType { VulnerabilityType HEADER_INJECTION = type(VulnerabilityTypes.HEADER_INJECTION) .mark(HEADER_INJECTION_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType STACKTRACE_LEAK = - type(VulnerabilityTypes.STACKTRACE_LEAK).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.STACKTRACE_LEAK).excludedSources(DB_EXCLUDED).build(); VulnerabilityType VERB_TAMPERING = - type(VulnerabilityTypes.VERB_TAMPERING).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.VERB_TAMPERING).excludedSources(DB_EXCLUDED).build(); VulnerabilityType ADMIN_CONSOLE_ACTIVE = type(VulnerabilityTypes.ADMIN_CONSOLE_ACTIVE) .deduplicable(false) .hash(VulnerabilityType::serviceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType DEFAULT_HTML_ESCAPE_INVALID = - type(VulnerabilityTypes.DEFAULT_HTML_ESCAPE_INVALID) - .excludedSources(SourceTypes.SQL_TABLE) - .build(); + type(VulnerabilityTypes.DEFAULT_HTML_ESCAPE_INVALID).excludedSources(DB_EXCLUDED).build(); VulnerabilityType SESSION_TIMEOUT = - type(VulnerabilityTypes.SESSION_TIMEOUT).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.SESSION_TIMEOUT).excludedSources(DB_EXCLUDED).build(); VulnerabilityType DIRECTORY_LISTING_LEAK = - type(VulnerabilityTypes.DIRECTORY_LISTING_LEAK) - .excludedSources(SourceTypes.SQL_TABLE) - .build(); + type(VulnerabilityTypes.DIRECTORY_LISTING_LEAK).excludedSources(DB_EXCLUDED).build(); VulnerabilityType INSECURE_JSP_LAYOUT = - type(VulnerabilityTypes.INSECURE_JSP_LAYOUT).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.INSECURE_JSP_LAYOUT).excludedSources(DB_EXCLUDED).build(); VulnerabilityType HARDCODED_SECRET = - type(VulnerabilityTypes.HARDCODED_SECRET).excludedSources(SourceTypes.SQL_TABLE).build(); + type(VulnerabilityTypes.HARDCODED_SECRET).excludedSources(DB_EXCLUDED).build(); VulnerabilityType INSECURE_AUTH_PROTOCOL = type(VulnerabilityTypes.INSECURE_AUTH_PROTOCOL) .hash(VulnerabilityType::evidenceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType REFLECTION_INJECTION = type(VulnerabilityTypes.REFLECTION_INJECTION) .mark(REFLECTION_INJECTION_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType SESSION_REWRITING = type(VulnerabilityTypes.SESSION_REWRITING) .deduplicable(false) .hash(VulnerabilityType::serviceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType DEFAULT_APP_DEPLOYED = type(VulnerabilityTypes.DEFAULT_APP_DEPLOYED) .deduplicable(false) .hash(VulnerabilityType::serviceHash) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); VulnerabilityType UNTRUSTED_DESERIALIZATION = type(VulnerabilityTypes.UNTRUSTED_DESERIALIZATION) .mark(UNTRUSTED_DESERIALIZATION_MARK) - .excludedSources(SourceTypes.SQL_TABLE) + .excludedSources(DB_EXCLUDED) .build(); /* All vulnerability types that have a mark. Should be updated if new vulnerabilityType with mark is added */ @@ -193,7 +192,7 @@ public interface VulnerabilityType { /** A flag to indicate if the vulnerability is deduplicable. */ boolean isDeduplicable(); - byte[] excludedSources(); + BitSet excludedSources(); static Builder type(final byte type) { return new Builder(type); @@ -209,7 +208,7 @@ class VulnerabilityTypeImpl implements VulnerabilityType { private final boolean deduplicable; - private final byte[] excludedSources; + private final BitSet excludedSources; private final BiFunction hash; @@ -218,7 +217,7 @@ public VulnerabilityTypeImpl( final char separator, final int mark, final boolean deduplicable, - final byte[] excludedSources, + final BitSet excludedSources, final BiFunction hash) { this.type = type; this.separator = separator; @@ -254,8 +253,8 @@ public boolean isDeduplicable() { } @Override - public byte[] excludedSources() { - return excludedSources.clone(); + public BitSet excludedSources() { + return excludedSources; } /** Useful for troubleshooting issues when vulns are serialized without moshi */ @@ -269,7 +268,7 @@ class Builder { private char separator = ' '; private int mark = NOT_MARKED; private boolean deduplicable = true; - private byte[] excludedSources = new byte[0]; + private BitSet excludedSources = new BitSet(); private BiFunction hash = VulnerabilityType::fileAndLineHash; @@ -292,13 +291,8 @@ public Builder deduplicable(final boolean deduplicable) { return this; } - public Builder excludedSources(final byte[] excludedSources) { - this.excludedSources = excludedSources.clone(); - return this; - } - - public Builder excludedSources(final byte excludedSources) { - this.excludedSources = new byte[] {excludedSources}; + public Builder excludedSources(final BitSet excludedSources) { + this.excludedSources = excludedSources; return this; } diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java index 0fef8a16824..ce6efe67499 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/sink/SinkModuleBase.java @@ -145,7 +145,7 @@ protected final Evidence checkInjection( // filter excluded ranges final Range[] filteredRanges; - if (type.excludedSources().length > 0) { + if (!type.excludedSources().isEmpty()) { filteredRanges = Ranges.excludeRangesBySource(valueRanges, type.excludedSources()); } else { filteredRanges = valueRanges; diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java index d8bbec33ce8..91bfe870401 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/taint/Ranges.java @@ -10,6 +10,7 @@ import com.datadog.iast.util.Ranged; import com.datadog.iast.util.StringUtils; import datadog.trace.api.iast.SourceTypes; +import java.util.BitSet; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -433,28 +434,15 @@ public static Range[] splitRanges( * @param ranges the ranges to filter * @param source the byte value of the source to exclude (see {@link SourceTypes}) */ - public static Range[] excludeRangesBySource(Range[] ranges, byte[] source) { + public static Range[] excludeRangesBySource(Range[] ranges, BitSet source) { RangeBuilder newRanges = new RangeBuilder(ranges.length); for (Range range : ranges) { - boolean exclude = false; - - for (byte origin : source) { - if (range.getSource().getOrigin() == origin) { - exclude = true; - break; - } - } - - if (!exclude) { + if (!source.get(range.getSource().getOrigin())) { newRanges.add(range); } } return newRanges.toArray(); } - - public static Range[] excludeRangesBySource(Range[] ranges, byte source) { - return excludeRangesBySource(ranges, new byte[] {source}); - } } From 8c5291ee75d1491843e427f785afc5460e549674 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mario=20Vidal=20Dom=C3=ADnguez?= Date: Wed, 15 Jan 2025 14:20:12 +0100 Subject: [PATCH 19/19] Fix tests --- .../com/datadog/iast/taint/RangesTest.groovy | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy index fc63d979a13..42f6a97bad2 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/RangesTest.groovy @@ -382,20 +382,20 @@ class RangesTest extends DDSpecification { void 'test excludeRangesBySource method'() { when: - final result = Ranges.excludeRangesBySource(ranges as Range[], source as byte[]) + final result = Ranges.excludeRangesBySource(ranges as Range[], source as BitSet) then: final expectedArray = expected as Range[] result == expectedArray where: - ranges | source | expected - [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [SQL_TABLE] | [range(5, 3)] - [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [SQL_TABLE, REQUEST_QUERY] | [range(5, 3)] - [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [REQUEST_HEADER_NAME] | [rangeWithSource(0, 5, SQL_TABLE)] - [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [REQUEST_QUERY] | [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] - [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | [REQUEST_QUERY, REQUEST_HEADER_NAME] | [rangeWithSource(0, 5, SQL_TABLE)] - [] | [SQL_TABLE] | [] + ranges | source | expected + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | bitSetOf(SQL_TABLE) | [range(5, 3)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | bitSetOf(SQL_TABLE, REQUEST_QUERY) | [range(5, 3)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | bitSetOf(REQUEST_HEADER_NAME) | [rangeWithSource(0, 5, SQL_TABLE)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | bitSetOf(REQUEST_QUERY) | [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] + [rangeWithSource(0, 5, SQL_TABLE), range(5, 3)] | bitSetOf(REQUEST_QUERY, REQUEST_HEADER_NAME) | [rangeWithSource(0, 5, SQL_TABLE)] + [] | bitSetOf(SQL_TABLE) | [] } Range[] rangesFromSpec(List> spec) { @@ -437,4 +437,12 @@ class RangesTest extends DDSpecification { Range rangeWithSource(final int start, final int length, final byte source, final String name = 'name', final String value = 'value') { return new Range(start, length, new Source(source, name, value), NOT_MARKED) } + + BitSet bitSetOf(byte... values) { + BitSet bitSet = new BitSet() + for (byte value : values) { + bitSet.set(value) + } + return bitSet + } }