Send IAST vulnerability secure marks to backend (#7645)

What Does This Do
Added the ability for ranges to return the types of vulnerabilities they are marked for.
A closed list of vulnerability types with assigned marks has been created.
The encoding and redaction of vulnerability evidence were updated to include a new secure_marks field in the IAST JSON, which holds an array of vulnerability types for which the evidence is marked.

Motivation
Send vulnerability secure marks to allow backend to recalculate vulnerability score

Author: jandro996

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Range.java

@@ -4,10 +4,13 @@
 
 import com.datadog.iast.model.json.SourceIndex;
 import com.datadog.iast.util.Ranged;
+import java.util.HashSet;
 import java.util.Objects;
+import java.util.Set;
 import java.util.StringJoiner;
 import javax.annotation.Nonnegative;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 public final class Range implements Ranged {
 
@@ -91,4 +94,17 @@ public Range consolidate() {
     return new Range(
         start, length, new Source(source.getOrigin(), source.getName(), source.getValue()), marks);
   }
+
+  public @Nullable Set<VulnerabilityType> getMarkedVulnerabilities() {
+    if (marks == NOT_MARKED) {
+      return null;
+    }
+    Set<VulnerabilityType> vulnerabilities = new HashSet<>();
+    for (VulnerabilityType type : VulnerabilityType.MARKED_VULNERABILITIES) {
+      if ((marks & type.mark()) != 0) {
+        vulnerabilities.add(type);
+      }
+    }
+    return vulnerabilities;
+  }
 }

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java

@@ -109,6 +109,22 @@ public interface VulnerabilityType {
           .mark(UNTRUSTED_DESERIALIZATION_MARK)
           .build();
 
+  /* All vulnerability types that have a mark. Should be updated if new vulnerabilityType with mark is added */
+  VulnerabilityType[] MARKED_VULNERABILITIES = {
+    SQL_INJECTION,
+    COMMAND_INJECTION,
+    PATH_TRAVERSAL,
+    LDAP_INJECTION,
+    SSRF,
+    UNVALIDATED_REDIRECT,
+    XPATH_INJECTION,
+    TRUST_BOUNDARY_VIOLATION,
+    XSS,
+    HEADER_INJECTION,
+    REFLECTION_INJECTION,
+    UNTRUSTED_DESERIALIZATION
+  };
+
   String name();
 
   /** A bit flag to ignore tainted ranges for this vulnerability. Set to 0 if none. */

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/json/AdapterFactory.java

@@ -184,6 +184,7 @@ public static class RedactionContext {
     private final boolean sensitive;
     private boolean sensitiveRanges;
     @Nullable private String redactedValue;
+    @Nullable private Set<VulnerabilityType> markedTypes;
 
     public RedactionContext(final Source source) {
       this.source = source;
@@ -192,6 +193,7 @@ public RedactionContext(final Source source) {
       if (this.sensitive) {
         this.redactedValue = handler.redactSource(source);
       }
+      this.markedTypes = null;
     }
 
     public Source getSource() {
@@ -217,5 +219,14 @@ public void markWithSensitiveRanges() {
         redactedValue = SensitiveHandler.get().redactSource(source);
       }
     }
+
+    public void setMarkedTypes(@Nullable Set<VulnerabilityType> markedTypes) {
+      this.markedTypes = markedTypes;
+    }
+
+    @Nullable
+    public Set<VulnerabilityType> getMarkedTypes() {
+      return markedTypes;
+    }
   }
 }

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/json/EvidenceAdapter.java

@@ -28,6 +28,7 @@
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Queue;
+import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -71,6 +72,20 @@ private String substring(final String value, final Ranged range) {
     return value.substring(range.getStart(), end);
   }
 
+  private static void writeSecureMarks(
+      final JsonWriter writer, final @Nullable Set<VulnerabilityType> markedVulnerabilities)
+      throws IOException {
+    if (markedVulnerabilities == null || markedVulnerabilities.isEmpty()) {
+      return;
+    }
+    writer.name("secure_marks");
+    writer.beginArray();
+    for (VulnerabilityType type : markedVulnerabilities) {
+      writer.value(type.name());
+    }
+    writer.endArray();
+  }
+
   private class DefaultEvidenceAdapter extends FormattingAdapter<Evidence> {
 
     @Override
@@ -128,6 +143,7 @@ private void writeValuePart(
       if (range != null) {
         writer.name("source");
         sourceAdapter.toJson(writer, range.getSource());
+        writeSecureMarks(writer, range.getMarkedVulnerabilities());
       }
       writer.endObject();
     }
@@ -389,6 +405,8 @@ static class RedactableTaintedValuePart implements ValuePart {
 
     private final List<Ranged> sensitiveRanges;
 
+    @Nullable private final Set<VulnerabilityType> markedTypes;
+
     private RedactableTaintedValuePart(
         final JsonAdapter<Source> adapter,
         final Range range,
@@ -403,11 +421,14 @@ private RedactableTaintedValuePart(
               .map(it -> shift(it, -range.getStart()))
               .sorted(Comparator.comparing(Ranged::getStart))
               .collect(Collectors.toList());
+
+      this.markedTypes = range.getMarkedVulnerabilities();
     }
 
     @Override
     public void write(final Context ctx, final JsonWriter writer) throws IOException {
       final RedactionContext redaction = ctx.getRedaction(source);
+      redaction.setMarkedTypes(markedTypes);
       if (redaction.shouldRedact()) {
         for (final ValuePart part : split(redaction)) {
           part.write(ctx, writer);
@@ -418,6 +439,7 @@ public void write(final Context ctx, final JsonWriter writer) throws IOException
         writeTruncableValue(writer, value);
         writer.name("source");
         adapter.toJson(writer, source);
+        writeSecureMarks(writer, markedTypes);
         writer.endObject();
       }
     }
@@ -454,9 +476,10 @@ private void addValuePart(
       if (start < end) {
         final Source source = ctx.getSource();
         final String chunk = value.substring(start, end);
+        final Set<VulnerabilityType> markedTypes = ctx.getMarkedTypes();
         if (!redact) {
           // append the value
-          valueParts.add(new TaintedValuePart(adapter, source, chunk, false));
+          valueParts.add(new TaintedValuePart(adapter, source, chunk, false, markedTypes));
         } else {
           final int length = chunk.length();
           final String sourceValue = source.getValue();
@@ -470,7 +493,7 @@ private void addValuePart(
             // otherwise redact the string
             pattern = SensitiveHandler.get().redactString(chunk);
           }
-          valueParts.add(new TaintedValuePart(adapter, source, pattern, true));
+          valueParts.add(new TaintedValuePart(adapter, source, pattern, true, markedTypes));
         }
       }
     }
@@ -489,15 +512,19 @@ static class TaintedValuePart implements ValuePart {
 
     private final boolean redacted;
 
+    @Nullable private final Set<VulnerabilityType> markedTypes;
+
     private TaintedValuePart(
         final JsonAdapter<Source> adapter,
         final Source source,
         final String value,
-        final boolean redacted) {
+        final boolean redacted,
+        final @Nullable Set<VulnerabilityType> markedTypes) {
       this.adapter = adapter;
       this.source = source;
       this.value = value;
       this.redacted = redacted;
+      this.markedTypes = markedTypes;
     }
 
     @Override
@@ -516,6 +543,7 @@ public void write(final Context ctx, final JsonWriter writer) throws IOException
         writer.name("value");
       }
       writeTruncableValue(writer, value);
+      writeSecureMarks(writer, markedTypes);
       writer.endObject();
     }
   }

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

@@ -1,15 +1,24 @@
 package com.datadog.iast.model
 
+import static com.datadog.iast.model.VulnerabilityType.SQL_INJECTION
+import static com.datadog.iast.model.VulnerabilityType.XSS
 import datadog.trace.api.iast.SourceTypes
 import datadog.trace.api.iast.VulnerabilityMarks
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED
+import static datadog.trace.api.iast.VulnerabilityMarks.SQL_INJECTION_MARK
+import static datadog.trace.api.iast.VulnerabilityMarks.XSS_MARK
 import datadog.trace.test.util.DDSpecification
+import spock.lang.Shared
 
 class RangeTest extends DDSpecification {
 
+  @Shared
+  int  multipleMarks  = SQL_INJECTION_MARK | XSS_MARK
+
   def 'shift'() {
     given:
     final source = new Source(SourceTypes.NONE, null, null)
-    final orig = new Range(start, length, source, VulnerabilityMarks.SQL_INJECTION_MARK)
+    final orig = new Range(start, length, source, SQL_INJECTION_MARK)
 
     when:
     final result = orig.shift(shift)
@@ -19,7 +28,7 @@ class RangeTest extends DDSpecification {
     result.source == source
     result.start == startResult
     result.length == lengthResult
-    result.marks == VulnerabilityMarks.SQL_INJECTION_MARK
+    result.marks == SQL_INJECTION_MARK
     result.isValid() == valid
 
     where:
@@ -43,4 +52,24 @@ class RangeTest extends DDSpecification {
     then:
     result.is(orig)
   }
+
+
+
+  void 'test getMarkedVulnerabilities'() {
+    given:
+    final source = new Source(SourceTypes.NONE, null, null)
+    final range = new Range(0, 1, source, marks)
+
+    when:
+    final vulnerabilities = range.getMarkedVulnerabilities()
+
+    then:
+    vulnerabilities == expected
+
+    where:
+    marks                           | expected
+    NOT_MARKED                      | null
+    SQL_INJECTION_MARK              | [SQL_INJECTION] as Set
+    multipleMarks | [SQL_INJECTION, XSS] as Set
+  }
 }

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

@@ -1,5 +1,6 @@
 package com.datadog.iast.model
 
+import datadog.trace.api.iast.VulnerabilityMarks
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.test.util.DDSpecification
 
@@ -52,6 +53,13 @@ class VulnerabilityTypeTest extends DDSpecification {
     INSECURE_AUTH_PROTOCOL      | getSpanAndStackLocation(123)           | new Evidence("Authorization : Digest") | 871205334
   }
 
+  void 'test marked vulnerabilities are not NOT_MARKED'() {
+    expect:
+    for (VulnerabilityType type : VulnerabilityType.MARKED_VULNERABILITIES) {
+      assert type.mark() != VulnerabilityMarks.NOT_MARKED
+    }
+  }
+
   private Location getSpanAndStackLocation(final long spanId) {
     final span = Stub(AgentSpan)
     span.getSpanId() >> spanId

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/model/json/EvidenceEncodingTest.groovy

@@ -6,6 +6,7 @@ import com.datadog.iast.model.Source
 import com.squareup.moshi.JsonAdapter
 import com.squareup.moshi.Moshi
 import datadog.trace.api.config.IastConfig
+import static datadog.trace.api.iast.VulnerabilityMarks.XSS_MARK
 import datadog.trace.test.util.DDSpecification
 import org.skyscreamer.jsonassert.JSONAssert
 import spock.lang.Shared
@@ -64,10 +65,15 @@ class EvidenceEncodingTest extends DDSpecification {
     'Hello World'               | [range(0, 11, source(0))]                        | '{"valueParts": [{"value": "Hello World", "source": 0}]}'
     'Hello World'               | [range(5, 1, source(0))]                         | '{"valueParts": [{"value": "Hello"}, {"value": " ", "source": 0}, {"value": "World"}]}'
     'java.lang.Object@1cb991da' | [range(0, Integer.MAX_VALUE, source(0))]         | '{"valueParts": [{"value": "java.lang.Object@1cb991da", "source": 0}]}'
+    'Hello World'               | [range(6, 5, source(0), XSS_MARK)]               | '{"valueParts": [{"value": "Hello "}, {"value": "World", "source": 0, "secure_marks": [XSS]}]}'
   }
 
   private static Range range(final int start, final int length, final Source source) {
-    return new Range(start, length, source, NOT_MARKED)
+    return range(start, length, source, NOT_MARKED)
+  }
+
+  private static Range range(final int start, final int length, final Source source, final int mark) {
+    return new Range(start, length, source, mark)
   }
 
   private static Source source(final int index) {