Skip to content

Commit

Permalink
Improve IAST metric unwrapping logic
Browse files Browse the repository at this point in the history
  • Loading branch information
manuel-alvarez-alvarez committed Mar 22, 2024
1 parent 46fdeca commit f3b71b7
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 32 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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() {
Expand All @@ -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];
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}

0 comments on commit f3b71b7

Please sign in to comment.