Improve IAST metric unwrapping logic

DataDog · Mar 22, 2024 · f3b71b7 · f3b71b7
1 parent 46fdeca
commit f3b71b7
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 32 deletions.
diff --git a/internal-api/src/main/java/datadog/trace/api/iast/telemetry/IastMetric.java b/internal-api/src/main/java/datadog/trace/api/iast/telemetry/IastMetric.java
@@ -3,7 +3,6 @@
 import datadog.trace.api.iast.SourceTypes;
 import datadog.trace.api.iast.VulnerabilityTypes;
 import java.util.Locale;
-import java.util.function.Function;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -132,34 +131,40 @@ public String getSpanTag(final byte tagValue) {
     return spanTags[tagValue];
   }
 
-  public static final class Tag {
+  public abstract static class Tag {
 
     public static final Tag VULNERABILITY_TYPE =
-        new Tag("vulnerability_type", VulnerabilityTypes.STRINGS, VulnerabilityTypes::unwrap);
+        new Tag("vulnerability_type", VulnerabilityTypes.STRINGS) {
+          @Nullable
+          @Override
+          public byte[] unwrap(byte tagValue) {
+            return VulnerabilityTypes.unwrap(tagValue);
+          }
+        };
 
     public static final Tag SOURCE_TYPE =
-        new Tag("source_type", SourceTypes.STRINGS, SourceTypes::unwrap);
+        new Tag("source_type", SourceTypes.STRINGS) {
 
-    private final String name;
+          @Nullable
+          @Override
+          public byte[] unwrap(byte tagValue) {
+            return SourceTypes.unwrap(tagValue);
+          }
+        };
 
-    private final String[] values;
+    protected final String name;
 
-    private final String[] telemetryTags;
+    protected final String[] values;
 
-    @Nullable private final Function<Byte, byte[]> unwrap;
+    protected final String[] telemetryTags;
 
     private Tag(final String name, final String[] values) {
-      this(name, values, null);
-    }
-
-    private Tag(final String name, final String[] values, final Function<Byte, byte[]> unwrap) {
       this.name = name;
       this.values = values;
       telemetryTags = new String[values.length];
       for (int i = 0; i < values.length; i++) {
         telemetryTags[i] = name + ":" + values[i];
       }
-      this.unwrap = unwrap;
     }
 
     public String getName() {
@@ -171,9 +176,7 @@ public int count() {
     }
 
     @Nullable
-    public byte[] unwrap(final byte tagValue) {
-      return unwrap == null ? null : unwrap.apply(tagValue);
-    }
+    public abstract byte[] unwrap(final byte tagValue);
 
     public String getTelemetryTag(final byte tagValue) {
       return telemetryTags[tagValue];

diff --git a/internal-api/src/main/java/datadog/trace/api/iast/telemetry/IastMetricCollector.java b/internal-api/src/main/java/datadog/trace/api/iast/telemetry/IastMetricCollector.java
@@ -103,25 +103,20 @@ public static void add(
   }
 
   public void addMetric(final IastMetric metric, final byte tagValue, final int value) {
-    final Tag tag = metric.getTag();
-    if (tag != null) {
-      final byte[] unwrapped = tag.unwrap(tagValue);
+    final int index = metric.getIndex(tagValue);
+    if (index >= 0) {
+      counters.getAndAdd(index, value);
+    } else if (tagValue < 0) {
+      final Tag tag = metric.getTag();
+      final byte[] unwrapped = tag == null ? null : tag.unwrap(tagValue);
       if (unwrapped != null) {
-        // e.g.: VulnerabilityTypes.RESPONSE_HEADER
-        for (final byte unwrappedValue : unwrapped) {
-          increment(metric.getIndex(unwrappedValue), value);
+        for (byte unwrappedValue : unwrapped) {
+          final int unwrappedIndex = metric.getIndex(unwrappedValue);
+          if (unwrappedIndex >= 0) {
+            counters.getAndAdd(unwrappedIndex, value);
+          }
         }
-      } else {
-        increment(metric.getIndex(tagValue), value);
       }
-    } else {
-      increment(metric.getIndex(tagValue), value);
-    }
-  }
-
-  private void increment(final int index, final int value) {
-    if (index >= 0) {
-      counters.getAndAdd(index, value);
     }
   }
 

diff --git a/internal-api/src/test/groovy/datadog/trace/api/iast/telemetry/IastMetricCollectorTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/iast/telemetry/IastMetricCollectorTest.groovy
@@ -226,9 +226,15 @@ class IastMetricCollectorTest extends DDSpecification {
     IastMetric.INSTRUMENTED_SINK   | VulnerabilityTypes.SPRING_RESPONSE // wrapped spring response
     IastMetric.EXECUTED_SINK       | VulnerabilityTypes.SPRING_RESPONSE
 
+    IastMetric.INSTRUMENTED_SINK   | VulnerabilityTypes.APPLICATION // wrapped application vuls
+    IastMetric.EXECUTED_SINK       | VulnerabilityTypes.APPLICATION
+
     IastMetric.INSTRUMENTED_SOURCE | SourceTypes.REQUEST_HEADER_NAME
     IastMetric.EXECUTED_SOURCE     | SourceTypes.REQUEST_HEADER_NAME
 
+    IastMetric.INSTRUMENTED_SOURCE | SourceTypes.KAFKA_MESSAGE // wrapped kafka sources
+    IastMetric.EXECUTED_SOURCE     | SourceTypes.KAFKA_MESSAGE
+
     IastMetric.EXECUTED_TAINTED    | null
   }
 }