Send IAST vulnerability secure marks to backend

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 org.jetbrains.annotations.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,21 @@ public interface VulnerabilityType {
           .mark(UNTRUSTED_DESERIALIZATION_MARK)
           .build();
 
+  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) {

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

@@ -71,8 +71,8 @@ class EvidenceRedactionTest extends DDSpecification {
     new StringValuePart(null)                                  | _
     new StringValuePart('')                                    | _
     new RedactedValuePart(null)                                | _
-    new TaintedValuePart(Stub(JsonAdapter), null, null, true)  | _
-    new TaintedValuePart(Stub(JsonAdapter), null, null, false) | _
+    new TaintedValuePart(Stub(JsonAdapter), null, null, true, null)  | _
+    new TaintedValuePart(Stub(JsonAdapter), null, null, false, null) | _
   }
 
   void 'test #suite'() {
@@ -94,6 +94,25 @@ class EvidenceRedactionTest extends DDSpecification {
     suite << readTestSuite('redaction/evidence-redaction-suite.yml')
   }
 
+  void 'test secure_marks #suite'() {
+    given:
+    Assume.assumeFalse("Ignored test", suite.ignored)
+    final type =  VulnerabilityBatch
+    final adapter = VulnerabilityEncoding.MOSHI.adapter(type)
+
+    when:
+    final redacted = adapter.toJson(suite.input)
+
+    then:
+    final received = JsonOutput.prettyPrint(redacted)
+    final description = suite.description
+    final expected = suite.expected
+    JSONAssert.assertEquals(description, expected, received, JSONCompareMode.NON_EXTENSIBLE)
+
+    where:
+    suite << readTestSuite('redaction/evidence-redaction-suite-with-marks.yml')
+  }
+
   private Iterable<TestSuite> readTestSuite(final String fileName) {
     final file = ClassLoader.getSystemResource(fileName)
     final tests = yaml.parse(file.openStream())