diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Source.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Source.java index 8c79cd025c6..361fa9b961d 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Source.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Source.java @@ -16,6 +16,9 @@ public final class Source implements Taintable.Source { private static final Logger LOGGER = LoggerFactory.getLogger(Source.class); + /** Placeholder for non char sequence objects */ + public static final Object PROPAGATION_PLACEHOLDER = new Object(); + // value to send in the rare case that the name/value have been garbage collected private static final String GARBAGE_COLLECTED_REF = "[unknown: original value was garbage collected]"; @@ -51,7 +54,9 @@ public String getValue() { @Nullable private String asString(@Nullable final Object target) { Object value = target; - if (value instanceof Reference) { + if (value == PROPAGATION_PLACEHOLDER) { + value = null; + } else if (value instanceof Reference) { value = ((Reference) value).get(); if (value == null) { value = GARBAGE_COLLECTED_REF; @@ -90,4 +95,11 @@ public boolean equals(Object o) { public int hashCode() { return Objects.hash(origin, getName(), getValue()); } + + public Source attachValue(final CharSequence result) { + if (value != PROPAGATION_PLACEHOLDER) { + return this; + } + return new Source(origin, name, result); + } } diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/BaseCodecModule.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/BaseCodecModule.java deleted file mode 100644 index 576d04de353..00000000000 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/BaseCodecModule.java +++ /dev/null @@ -1,111 +0,0 @@ -package com.datadog.iast.propagation; - -import static com.datadog.iast.taint.Tainteds.canBeTainted; - -import com.datadog.iast.model.Range; -import com.datadog.iast.taint.TaintedObject; -import com.datadog.iast.taint.TaintedObjects; -import datadog.trace.api.iast.IastContext; -import datadog.trace.api.iast.propagation.CodecModule; -import java.util.function.Function; -import javax.annotation.Nonnull; -import javax.annotation.Nullable; - -public abstract class BaseCodecModule implements CodecModule { - - @Override - public void onUrlDecode( - @Nonnull final String value, @Nullable final String encoding, @Nonnull final String result) { - if (!canBeTainted(result)) { - return; - } - taintIfInputIsTainted( - result, value, tainted -> urlDecodeRanges(value, encoding, result, tainted.getRanges())); - } - - @Override - public void onStringFromBytes( - @Nonnull final byte[] value, @Nullable final String charset, @Nonnull final String result) { - if (!canBeTainted(result)) { - return; - } - taintIfInputIsTainted( - result, value, tainted -> fromBytesRanges(value, charset, result, tainted.getRanges())); - } - - @Override - public void onStringGetBytes( - @Nonnull final String value, @Nullable final String charset, @Nonnull final byte[] result) { - if (result == null || result.length == 0) { - return; - } - taintIfInputIsTainted( - result, value, tainted -> getBytesRanges(value, charset, result, tainted.getRanges())); - } - - @Override - public void onBase64Encode(@Nullable byte[] value, @Nullable byte[] result) { - if (value == null || result == null || result.length == 0) { - return; - } - taintIfInputIsTainted( - result, value, tainted -> encodeBase64Ranges(value, result, tainted.getRanges())); - } - - @Override - public void onBase64Decode(@Nullable byte[] value, @Nullable byte[] result) { - if (value == null || result == null || result.length == 0) { - return; - } - taintIfInputIsTainted( - result, value, tainted -> decodeBase64Ranges(value, result, tainted.getRanges())); - } - - private static void taintIfInputIsTainted( - final Object value, final Object input, final Function mapper) { - final IastContext ctx = IastContext.Provider.get(); - if (ctx == null) { - return; - } - final TaintedObjects to = ctx.getTaintedObjects(); - final TaintedObject tainted = to.get(input); - if (hasRanges(tainted)) { - final Range[] ranges = mapper.apply(tainted); - if (hasRanges(ranges)) { - to.taint(value, ranges); - } - } - } - - private static boolean hasRanges(@Nullable final TaintedObject tainted) { - return tainted != null && hasRanges(tainted.getRanges()); - } - - private static boolean hasRanges(@Nullable final Range[] ranges) { - return ranges != null && ranges.length > 0; - } - - protected abstract Range[] urlDecodeRanges( - final @Nonnull String value, - final @Nullable String charset, - @Nonnull final String result, - @Nonnull final Range[] ranges); - - protected abstract Range[] fromBytesRanges( - final @Nonnull byte[] value, - final @Nullable String charset, - @Nonnull final String result, - @Nonnull final Range[] ranges); - - protected abstract Range[] getBytesRanges( - final @Nonnull String value, - final @Nullable String charset, - @Nonnull final byte[] result, - @Nonnull final Range[] ranges); - - protected abstract Range[] decodeBase64Ranges( - final @Nonnull byte[] value, @Nonnull final byte[] result, @Nonnull final Range[] ranges); - - protected abstract Range[] encodeBase64Ranges( - final @Nonnull byte[] value, @Nonnull final byte[] result, @Nonnull final Range[] ranges); -} diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/FastCodecModule.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/FastCodecModule.java index 919c4a20313..75e4317e4fe 100644 --- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/FastCodecModule.java +++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/propagation/FastCodecModule.java @@ -1,55 +1,36 @@ package com.datadog.iast.propagation; -import static com.datadog.iast.taint.Ranges.highestPriorityRange; - -import com.datadog.iast.model.Range; -import com.datadog.iast.taint.Ranges; +import datadog.trace.api.iast.propagation.CodecModule; import javax.annotation.Nonnull; import javax.annotation.Nullable; -public class FastCodecModule extends BaseCodecModule { +public class FastCodecModule extends PropagationModuleImpl implements CodecModule { @Override - protected Range[] urlDecodeRanges( - @Nonnull final String value, - @Nullable final String encoding, - @Nonnull final String result, - @Nonnull final Range[] ranges) { - final Range range = highestPriorityRange(ranges); - return new Range[] {Ranges.copyWithPosition(range, 0, result.length())}; + public void onUrlDecode( + @Nonnull final String value, @Nullable final String encoding, @Nonnull final String result) { + taintIfTainted(result, value); } @Override - protected Range[] fromBytesRanges( - @Nonnull final byte[] value, - @Nullable final String charset, - @Nonnull final String result, - @Nonnull final Range[] ranges) { - final Range range = highestPriorityRange(ranges); - return new Range[] {Ranges.copyWithPosition(range, 0, result.length())}; + public void onStringFromBytes( + @Nonnull final byte[] value, @Nullable final String charset, @Nonnull final String result) { + taintIfTainted(result, value); } @Override - protected Range[] getBytesRanges( - @Nonnull final String value, - @Nullable final String charset, - @Nonnull final byte[] result, - @Nonnull final Range[] ranges) { - final Range range = highestPriorityRange(ranges); - return new Range[] {Ranges.copyWithPosition(range, 0, result.length)}; + public void onStringGetBytes( + @Nonnull final String value, @Nullable final String charset, @Nonnull final byte[] result) { + taintIfTainted(result, value); } @Override - protected Range[] decodeBase64Ranges( - @Nonnull final byte[] value, @Nonnull final byte[] result, @Nonnull final Range[] ranges) { - final Range range = highestPriorityRange(ranges); - return new Range[] {Ranges.copyWithPosition(range, 0, result.length)}; + public void onBase64Encode(@Nullable byte[] value, @Nullable byte[] result) { + taintIfTainted(result, value); } @Override - protected Range[] encodeBase64Ranges( - @Nonnull final byte[] value, @Nonnull final byte[] result, @Nonnull final Range[] ranges) { - final Range range = highestPriorityRange(ranges); - return new Range[] {Ranges.copyWithPosition(range, 0, result.length)}; + public void onBase64Decode(@Nullable byte[] value, @Nullable byte[] result) { + taintIfTainted(result, value); } } 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 38278d5c58f..005b5912bd8 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,5 +1,6 @@ package com.datadog.iast.propagation; +import static com.datadog.iast.model.Source.PROPAGATION_PLACEHOLDER; 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; @@ -16,6 +17,7 @@ import datadog.trace.api.iast.Taintable; import datadog.trace.api.iast.propagation.PropagationModule; import java.lang.ref.WeakReference; +import java.lang.reflect.Array; import java.util.function.Predicate; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -313,7 +315,8 @@ private static Object sourceString(@Nullable final Object target, final boolean } return charSequence.toString(); } - return null; // ignore non char-sequence instances (e.g. byte buffers) + // ignore non char-sequence instances (e.g. byte buffers) + return value ? PROPAGATION_PLACEHOLDER : null; } @Contract("null -> false") @@ -324,6 +327,9 @@ private static boolean canBeTainted(@Nullable final Object target) { if (target instanceof CharSequence) { return Tainteds.canBeTainted((CharSequence) target); } + if (target.getClass().isArray()) { + return Array.getLength(target) > 0; + } return true; } @@ -396,7 +402,7 @@ private static Source highestPrioritySource( private static void internalTaint( @Nullable final IastContext ctx, @Nonnull final Object value, - @Nullable final Source source, + @Nullable Source source, int mark) { if (source == null) { return; @@ -409,6 +415,7 @@ private static void internalTaint( return; } if (value instanceof CharSequence) { + source = source.attachValue((CharSequence) value); to.taint(value, Ranges.forCharSequence((CharSequence) value, source, mark)); } else { to.taint(value, Ranges.forObject(source, mark)); @@ -419,7 +426,7 @@ private static void internalTaint( private static void internalTaint( @Nullable final IastContext ctx, @Nonnull final Object value, - @Nullable final Range[] ranges, + @Nullable Range[] ranges, final int mark) { if (ranges == null || ranges.length == 0) { return; @@ -429,8 +436,11 @@ private static void internalTaint( } else { final TaintedObjects to = getTaintedObjects(ctx); if (to != null) { - final Range[] markedRanges = markRanges(ranges, mark); - to.taint(value, markedRanges); + if (value instanceof CharSequence) { + ranges = attachSourceValue(ranges, (CharSequence) value); + } + ranges = markRanges(ranges, mark); + to.taint(value, ranges); } } } @@ -449,6 +459,20 @@ private static Range[] markRanges(@Nonnull final Range[] ranges, final int mark) return result; } + public static Range[] attachSourceValue( + @Nonnull final Range[] ranges, @Nonnull final CharSequence value) { + // unbound sources can only occur when there's a single range in the array + if (ranges.length != 1) { + return ranges; + } + final Range range = ranges[0]; + final Source source = range.getSource(); + final Source newSource = range.getSource().attachValue(value); + return newSource == source + ? ranges + : Ranges.forCharSequence(value, newSource, range.getMarks()); + } + private static class LazyContext implements IastContext { private boolean fetched; diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/FastCodecModuleTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/FastCodecModuleTest.groovy index 9e57b6fba97..61e146fc90d 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/FastCodecModuleTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/FastCodecModuleTest.groovy @@ -3,10 +3,7 @@ package com.datadog.iast.propagation import com.datadog.iast.taint.Ranges import com.datadog.iast.taint.TaintedObject import datadog.trace.api.iast.propagation.CodecModule -import groovy.transform.CompileDynamic - -@CompileDynamic class FastCodecModuleTest extends BaseCodecModuleTest { @Override @@ -40,37 +37,34 @@ class FastCodecModuleTest extends BaseCodecModuleTest { @Override protected void assertOnStringGetBytes(final String value, final String charset, final TaintedObject source, final TaintedObject target) { - final result = target.get() as byte[] assert target.ranges.size() == 1 final sourceRange = Ranges.highestPriorityRange(source.ranges) final range = target.ranges.first() assert range.start == 0 - assert range.length == result.length + assert range.length == Integer.MAX_VALUE // unbound for non char sequences assert range.source == sourceRange.source } @Override protected void assertBase64Decode(byte[] value, TaintedObject source, TaintedObject target) { - final result = target.get() as byte[] assert target.ranges.size() == 1 final sourceRange = Ranges.highestPriorityRange(source.ranges) final range = target.ranges.first() assert range.start == 0 - assert range.length == result.length + assert range.length == Integer.MAX_VALUE // unbound for non char sequences assert range.source == sourceRange.source } @Override protected void assertBase64Encode(byte[] value, TaintedObject source, TaintedObject target) { - final result = target.get() as byte[] assert target.ranges.size() == 1 final sourceRange = Ranges.highestPriorityRange(source.ranges) final range = target.ranges.first() assert range.start == 0 - assert range.length == result.length + assert range.length == Integer.MAX_VALUE // unbound for non char sequences assert range.source == sourceRange.source } } diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/PropagationModuleTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/PropagationModuleTest.groovy index 3586b8d8fce..1f88da8cad0 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/PropagationModuleTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/propagation/PropagationModuleTest.groovy @@ -4,6 +4,7 @@ import com.datadog.iast.IastModuleImplTestBase import com.datadog.iast.model.Range import com.datadog.iast.model.Source import com.datadog.iast.taint.Ranges + import com.datadog.iast.taint.TaintedObject import datadog.trace.api.Config import datadog.trace.api.iast.SourceTypes @@ -461,7 +462,7 @@ class PropagationModuleTest extends IastModuleImplTestBase { assert sourceValue == value.toString() break default: - assert sourceValue == null + assert sourceValue === Source.PROPAGATION_PLACEHOLDER break } if (name === value) { @@ -478,6 +479,35 @@ class PropagationModuleTest extends IastModuleImplTestBase { string('name') | date() } + void 'test propagation of the source value for non char sequences'() { + given: + final toTaint = 'hello' + final baos = toTaint.bytes + + when: 'tainting a non char sequence object' + module.taint(baos, SourceTypes.KAFKA_MESSAGE_KEY) + + then: + with(ctx.taintedObjects.get(baos)) { + assert ranges.length == 1 + final source = ranges.first().source + assert source.origin == SourceTypes.KAFKA_MESSAGE_KEY + assert source.@value === Source.PROPAGATION_PLACEHOLDER + assert source.value == null + } + + when: 'the object is propagated' + module.taintIfTainted(toTaint, baos) + + then: + with(ctx.taintedObjects.get(toTaint)) { + assert ranges.length == 1 + final source = ranges.first().source + assert source.origin == SourceTypes.KAFKA_MESSAGE_KEY + assert source.value == toTaint + } + } + private List> taintIfSuite() { return [ Tuple.tuple(string('string'), string('string')), diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintUtils.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintUtils.groovy index e82fe2c5392..0e963c63e12 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintUtils.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintUtils.groovy @@ -55,7 +55,9 @@ class TaintUtils { int start = pos int length = (upTo - i) - OPEN_MARK.length() assert length >= 0 - ranges.add(new Range(start, length, new Source(SourceTypes.NONE, null, null), mark)) + int from = i + OPEN_MARK.length() + String value = s.substring(from, from + length) + ranges.add(new Range(start, length, new Source(SourceTypes.NONE, null, value), mark)) pos += length i += OPEN_MARK.length() + length + CLOSE_MARK.length() - 1 } else { @@ -118,18 +120,19 @@ class TaintUtils { return result } - static Range toRange(List lst) { - toRange(lst.get(0), lst.get(1)) + static Range toRange(String string, List lst) { + toRange(string, lst.get(0), lst.get(1)) } - static Range toRange(int start, int length) { - new Range(start, length, new Source(SourceTypes.NONE, null, null), NOT_MARKED) + static Range toRange(String string, int start, int length) { + final value = string.substring(start, start + length) + new Range(start, length, new Source(SourceTypes.NONE, null, value), NOT_MARKED) } - static Range[] toRanges(List> lst) { + static Range[] toRanges(String string, List> lst) { if (lst == null) { return null } - lst.collect { toRange(it) } as Range[] + lst.collect { toRange(string, it) } as Range[] } } diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintUtilsTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintUtilsTest.groovy index 07c1c721874..0ef19c7d57b 100644 --- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintUtilsTest.groovy +++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/taint/TaintUtilsTest.groovy @@ -7,7 +7,7 @@ class TaintUtilsTest extends Specification { void 'taintFormat with empty ranges'() { expect: - taintFormat(s, toRanges(ranges)) == result + taintFormat(s, toRanges(s, ranges)) == result where: s | ranges | result @@ -16,10 +16,10 @@ class TaintUtilsTest extends Specification { void 'taintFormat amd fromTaintFormat'() { expect: - taintFormat(s, toRanges(ranges)) == result + taintFormat(s, toRanges(s, ranges)) == result and: - fromTaintFormat(result) == toRanges(ranges) + fromTaintFormat(result) == toRanges(s, ranges) where: s | ranges | result diff --git a/dd-smoke-tests/kafka/src/main/java/datadog/smoketest/kafka/iast/IastConfiguration.java b/dd-smoke-tests/kafka/src/main/java/datadog/smoketest/kafka/iast/IastConfiguration.java index 085f00adbaf..0964b417e7d 100644 --- a/dd-smoke-tests/kafka/src/main/java/datadog/smoketest/kafka/iast/IastConfiguration.java +++ b/dd-smoke-tests/kafka/src/main/java/datadog/smoketest/kafka/iast/IastConfiguration.java @@ -2,7 +2,6 @@ import java.util.HashMap; import java.util.Map; -import org.apache.kafka.clients.admin.NewTopic; import org.apache.kafka.clients.consumer.ConsumerConfig; import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.common.serialization.StringDeserializer; @@ -12,11 +11,10 @@ import org.springframework.context.annotation.Configuration; import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory; import org.springframework.kafka.config.TopicBuilder; -import org.springframework.kafka.core.ConsumerFactory; import org.springframework.kafka.core.DefaultKafkaConsumerFactory; import org.springframework.kafka.core.DefaultKafkaProducerFactory; +import org.springframework.kafka.core.KafkaAdmin; import org.springframework.kafka.core.KafkaTemplate; -import org.springframework.kafka.core.ProducerFactory; import org.springframework.kafka.support.serializer.JsonDeserializer; import org.springframework.kafka.support.serializer.JsonSerializer; @@ -24,46 +22,71 @@ public class IastConfiguration { public static final String GROUP_ID = "iast"; - public static final String TOPIC = "iast"; + + public static final String STRING_TOPIC = "iast_string"; + + public static final String JSON_TOPIC = "iast_json"; @Value("${spring.kafka.bootstrap-servers}") private String boostrapServers; @Bean - public NewTopic iastTopic() { - return TopicBuilder.name(TOPIC).partitions(1).replicas(1).compact().build(); + public KafkaAdmin.NewTopics iastTopics() { + return new KafkaAdmin.NewTopics( + TopicBuilder.name(STRING_TOPIC).partitions(1).replicas(1).compact().build(), + TopicBuilder.name(JSON_TOPIC).partitions(1).replicas(1).compact().build()); } @Bean - public ConsumerFactory iastConsumerFactory() { - Map config = new HashMap<>(); + public ConcurrentKafkaListenerContainerFactory iastStringListenerFactory() { + final Map config = new HashMap<>(); config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, boostrapServers); config.put(ConsumerConfig.GROUP_ID_CONFIG, GROUP_ID); config.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); - config.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, JsonDeserializer.class); - config.put(JsonDeserializer.TRUSTED_PACKAGES, "datadog.*"); - return new DefaultKafkaConsumerFactory<>(config); + config.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); + final DefaultKafkaConsumerFactory consumerFactory = + new DefaultKafkaConsumerFactory<>(config); + ConcurrentKafkaListenerContainerFactory factory = + new ConcurrentKafkaListenerContainerFactory<>(); + factory.setConsumerFactory(consumerFactory); + return factory; } @Bean - public ConcurrentKafkaListenerContainerFactory iastListenerFactory() { + public ConcurrentKafkaListenerContainerFactory iastJsonListenerFactory() { + final Map config = new HashMap<>(); + config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, boostrapServers); + config.put(ConsumerConfig.GROUP_ID_CONFIG, GROUP_ID); + config.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class); + config.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, JsonDeserializer.class); + config.put(JsonDeserializer.TRUSTED_PACKAGES, "datadog.*"); + final DefaultKafkaConsumerFactory consumerFactory = + new DefaultKafkaConsumerFactory<>(config); ConcurrentKafkaListenerContainerFactory factory = new ConcurrentKafkaListenerContainerFactory<>(); - factory.setConsumerFactory(iastConsumerFactory()); + factory.setConsumerFactory(consumerFactory); return factory; } @Bean - public ProducerFactory iastProducerFactory() { + public KafkaTemplate iastStringKafkaTemplate() { Map configProps = new HashMap<>(); configProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, boostrapServers); configProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class); - configProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, JsonSerializer.class); - return new DefaultKafkaProducerFactory<>(configProps); + configProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class); + final DefaultKafkaProducerFactory factory = + new DefaultKafkaProducerFactory<>(configProps); + return new KafkaTemplate<>(factory); } @Bean - public KafkaTemplate iastKafkaTemplate() { - return new KafkaTemplate<>(iastProducerFactory()); + public KafkaTemplate iastJsonKafkaTemplate() { + Map configProps = new HashMap<>(); + configProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, boostrapServers); + configProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class); + configProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, JsonSerializer.class); + final DefaultKafkaProducerFactory factory = + new DefaultKafkaProducerFactory<>(configProps); + return new KafkaTemplate<>(factory); } } diff --git a/dd-smoke-tests/kafka/src/main/java/datadog/smoketest/kafka/iast/IastController.java b/dd-smoke-tests/kafka/src/main/java/datadog/smoketest/kafka/iast/IastController.java index 1f82225bc66..d906761b5dd 100644 --- a/dd-smoke-tests/kafka/src/main/java/datadog/smoketest/kafka/iast/IastController.java +++ b/dd-smoke-tests/kafka/src/main/java/datadog/smoketest/kafka/iast/IastController.java @@ -1,7 +1,8 @@ package datadog.smoketest.kafka.iast; import static datadog.smoketest.kafka.iast.IastConfiguration.GROUP_ID; -import static datadog.smoketest.kafka.iast.IastConfiguration.TOPIC; +import static datadog.smoketest.kafka.iast.IastConfiguration.JSON_TOPIC; +import static datadog.smoketest.kafka.iast.IastConfiguration.STRING_TOPIC; import java.util.concurrent.CompletableFuture; import org.apache.kafka.clients.consumer.ConsumerRecord; @@ -21,36 +22,62 @@ public class IastController { private static final Logger LOGGER = LoggerFactory.getLogger(IastController.class); - private final KafkaTemplate iastKafkaTemplate; + private final KafkaTemplate stringTemplate; + private final KafkaTemplate jsonTemplate; public IastController( - @Qualifier("iastKafkaTemplate") final KafkaTemplate iastKafkaTemplate) { - this.iastKafkaTemplate = iastKafkaTemplate; + @Qualifier("iastStringKafkaTemplate") final KafkaTemplate stringTemplate, + @Qualifier("iastJsonKafkaTemplate") final KafkaTemplate jsonTemplate) { + this.stringTemplate = stringTemplate; + this.jsonTemplate = jsonTemplate; } - @GetMapping("/iast/kafka") - public CompletableFuture> vulnerability( - @RequestParam("type") final String type) { + @GetMapping("/iast/kafka/string") + public CompletableFuture> string(@RequestParam("type") final String type) { + return stringTemplate + .send(STRING_TOPIC, type, type) + .completable() + .thenApply(this::handleKafkaResponse); + } + + @GetMapping("/iast/kafka/json") + public CompletableFuture> json(@RequestParam("type") final String type) { final IastMessage message = new IastMessage(); message.setValue(type); - return iastKafkaTemplate - .send(TOPIC, type, message) + return jsonTemplate + .send(JSON_TOPIC, type, message) .completable() .thenApply(this::handleKafkaResponse); } - @KafkaListener(groupId = GROUP_ID, topics = TOPIC, containerFactory = "iastListenerFactory") - public void listen(final ConsumerRecord record) { - final String type = record.key(); + @KafkaListener( + groupId = GROUP_ID, + topics = STRING_TOPIC, + containerFactory = "iastStringListenerFactory") + public void listenString(final ConsumerRecord record) { + handle(record.key(), record.value()); + } + + @KafkaListener( + groupId = GROUP_ID, + topics = JSON_TOPIC, + containerFactory = "iastJsonListenerFactory") + public void listenJson(final ConsumerRecord record) { final IastMessage message = record.value(); - if ("source_key".equals(type)) { - LOGGER.info("Kafka tainted key: " + type); - } else if ("source_value".equals(type)) { - LOGGER.info("Kafka tainted value: " + message.getValue()); + handle(record.key(), message.getValue()); + } + + private void handle(final String key, final String value) { + if (key.endsWith("source_key")) { + LOGGER.info("Kafka tainted key: " + key); + } else if (key.endsWith("source_value")) { + LOGGER.info("Kafka tainted value: " + value); + } else { + throw new IllegalArgumentException("Non valid key " + key); } } - private ResponseEntity handleKafkaResponse(final SendResult result) { + private ResponseEntity handleKafkaResponse(final SendResult result) { if (result.getRecordMetadata().hasOffset()) { return ResponseEntity.ok("OK"); } else { diff --git a/dd-smoke-tests/kafka/src/test/groovy/IastKafkaSmokeTest.groovy b/dd-smoke-tests/kafka/src/test/groovy/IastKafkaSmokeTest.groovy index 1a5e4534900..23d9da8058c 100644 --- a/dd-smoke-tests/kafka/src/test/groovy/IastKafkaSmokeTest.groovy +++ b/dd-smoke-tests/kafka/src/test/groovy/IastKafkaSmokeTest.groovy @@ -48,9 +48,10 @@ class IastKafkaSmokeTest extends AbstractIastServerSmokeTest { return processBuilder } - void 'test kafka key source'() { + void 'test kafka #endpoint key source'() { setup: - final url = "http://localhost:${httpPort}/iast/kafka?type=source_key" + final type = "${endpoint}_source_key" + final url = "http://localhost:${httpPort}/iast/kafka/$endpoint?type=${type}" when: final response = client.newCall(new Request.Builder().url(url).get().build()).execute() @@ -58,14 +59,19 @@ class IastKafkaSmokeTest extends AbstractIastServerSmokeTest { then: response.body().string() == 'OK' hasTainted { tainted -> - tainted.value == 'Kafka tainted key: source_key' && - tainted.ranges[0].source.origin == 'kafka.message.key' + tainted.value == "Kafka tainted key: $type" && + tainted.ranges[0].source.origin == 'kafka.message.key' && + tainted.ranges[0].source.value == type } + + where: + endpoint << ['json', 'string'] } - void 'test kafka value source'() { + void 'test kafka #endpoint value source'() { setup: - final url = "http://localhost:${httpPort}/iast/kafka?type=source_value" + final type = "${endpoint}_source_value" + final url = "http://localhost:${httpPort}/iast/kafka/$endpoint?type=${type}" when: final response = client.newCall(new Request.Builder().url(url).get().build()).execute() @@ -73,8 +79,12 @@ class IastKafkaSmokeTest extends AbstractIastServerSmokeTest { then: response.body().string() == 'OK' hasTainted { tainted -> - tainted.value == 'Kafka tainted value: source_value' && - tainted.ranges[0].source.origin == 'kafka.message.value' + tainted.value == "Kafka tainted value: $type" && + tainted.ranges[0].source.origin == 'kafka.message.value' && + tainted.ranges[0].source.value == type } + + where: + endpoint << ['json', 'string'] } }