String builder taint tracking

dd-java-agent/agent-iast/README.MD

@@ -0,0 +1,70 @@
+# agent-iast
+
+## Benchmark results
+
+### String
+
+#### concat()
+
+| Benchmark                           | Mode | Cnt   | Score   | Error   | Units |
+|-------------------------------------|------|-------|---------|---------|-------|
+| StringConcatBenchmark.baseline      | ss   | 15000 |  50.977 | ± 1.561 | ns/op |
+| StringConcatBenchmark.iastDisabled  | ss   | 15000 |  52.963 | ± 0.748 | ns/op |
+| StringConcatBenchmark.notTainted    | ss   | 15000 |  77.463 | ± 0.990 | ns/op |
+| StringConcatBenchmark.stringTainted | ss   | 15000 | 132.859 | ± 3.565 | ns/op |
+| StringConcatBenchmark.paramTainted  | ss   | 15000 | 156.825 | ± 4.530 | ns/op |
+| StringConcatBenchmark.bothTainted   | ss   | 15000 | 145.427 | ± 3.138 | ns/op |
+
+### String Builder
+
+#### constructor()
+
+| Benchmark                               | Mode | Cnt   | Score   | Error   | Units |
+|-----------------------------------------|------|-------|---------|---------|-------|
+| StringBuilderInitBenchmark.baseline     | ss   | 15000 | 43.278  | ± 0.666 | ns/op |
+| StringBuilderInitBenchmark.iastDisabled | ss   | 15000 | 45.373  | ± 2.391 | ns/op |
+| StringBuilderInitBenchmark.notTainted   | ss   | 15000 | 66.833  | ± 1.292 | ns/op |
+| StringBuilderInitBenchmark.tainted      | ss   | 15000 | 100.316 | ± 2.767 | ns/op |
+
+#### append()
+
+| Benchmark                                         | Mode | Cnt   | Score   | Error   | Units |
+|---------------------------------------------------|------|-------|---------|---------|-------|
+| StringBuilderAppendBenchmark.baseline             | ss   | 15000 |  50.261 | ± 2.212 | ns/op |
+| StringBuilderAppendBenchmark.iastDisabled         | ss   | 15000 |  52.746 | ± 0.567 | ns/op |
+| StringBuilderAppendBenchmark.notTainted           | ss   | 15000 |  90.821 | ± 2.245 | ns/op |
+| StringBuilderAppendBenchmark.stringBuilderTainted | ss   | 15000 |  79.958 | ± 2.289 | ns/op |
+| StringBuilderAppendBenchmark.paramTainted         | ss   | 15000 | 116.093 | ± 3.961 | ns/op |
+| StringBuilderAppendBenchmark.bothTainted          | ss   | 15000 | 107.229 | ± 4.275 | ns/op |
+
+#### toString()
+
+| Benchmark                                   | Mode | Cnt   | Score  | Error   | Units |
+|---------------------------------------------|------|-------|--------|---------|-------|
+| StringBuilderToStringBenchmark.baseline     | ss   | 15000 | 29.817 | ± 2.493 | ns/op |
+| StringBuilderToStringBenchmark.iastDisabled | ss   | 15000 | 30.570 | ± 1.794 | ns/op |
+| StringBuilderToStringBenchmark.notTainted   | ss   | 15000 | 57.370 | ± 1.333 | ns/op |
+| StringBuilderToStringBenchmark.tainted      | ss   | 15000 | 92.077 | ± 1.775 | ns/op |
+
+### batch append operations
+
+| Benchmark                                 | (stringCount) | (taintedPct) | Mode | Cnt   | Score | Error | Units |
+|-------------------------------------------|---------------|--------------|------|-------|-------|-------|-------|
+| StringBuilderBatchBenchmark.baseline      | 10            | 0            | ss   | 15000 | 0.348 | 0.009 | us/op |
+| StringBuilderBatchBenchmark.baseline      | 10            | 50           | ss   | 15000 | 0.317 | 0.009 | us/op |
+| StringBuilderBatchBenchmark.baseline      | 10            | 100          | ss   | 15000 | 0.355 | 0.010 | us/op |
+| StringBuilderBatchBenchmark.iastDisabled  | 10            | 0            | ss   | 15000 | 0.355 | 0.009 | us/op |
+| StringBuilderBatchBenchmark.iastDisabled  | 10            | 50           | ss   | 15000 | 0.344 | 0.008 | us/op |
+| StringBuilderBatchBenchmark.iastDisabled  | 10            | 100          | ss   | 15000 | 0.370 | 0.013 | us/op |
+| StringBuilderBatchBenchmark.iastEnabled   | 10            | 0            | ss   | 15000 | 0.551 | 0.014 | us/op |
+| StringBuilderBatchBenchmark.iastEnabled   | 10            | 50           | ss   | 15000 | 0.794 | 0.014 | us/op |
+| StringBuilderBatchBenchmark.iastEnabled   | 10            | 100          | ss   | 15000 | 0.900 | 0.014 | us/op |
+| StringBuilderBatchBenchmark.baseline      | 100           | 0            | ss   | 15000 | 2.508 | 0.025 | us/op |                                                                
+| StringBuilderBatchBenchmark.baseline      | 100           | 50           | ss   | 15000 | 2.419 | 0.019 | us/op |                                                                
+| StringBuilderBatchBenchmark.baseline      | 100           | 100          | ss   | 15000 | 2.499 | 0.026 | us/op |  
+| StringBuilderBatchBenchmark.iastDisabled  | 100           | 0            | ss   | 15000 | 2.499 | 0.023 | us/op |
+| StringBuilderBatchBenchmark.iastDisabled  | 100           | 50           | ss   | 15000 | 2.608 | 0.180 | us/op |
+| StringBuilderBatchBenchmark.iastDisabled  | 100           | 100          | ss   | 15000 | 2.596 | 0.201 | us/op |
+| StringBuilderBatchBenchmark.iastEnabled   | 100           | 0            | ss   | 15000 | 3.426 | 0.028 | us/op |
+| StringBuilderBatchBenchmark.iastEnabled   | 100           | 50           | ss   | 15000 | 6.676 | 0.275 | us/op |
+| StringBuilderBatchBenchmark.iastEnabled   | 100           | 100          | ss   | 15000 | 7.539 | 0.164 | us/op |

dd-java-agent/agent-iast/build.gradle

@@ -17,6 +17,10 @@ dependencies {
   testImplementation project(':utils:test-utils')
 
   jmh project(':utils:test-utils')
+  jmh project(':dd-trace-core')
+  jmh project(':dd-java-agent:agent-builder')
+  jmh project(':dd-java-agent:instrumentation:iast-instrumenter')
+  jmh project(':dd-java-agent:instrumentation:java-lang')
 }
 
 sourceCompatibility = JavaVersion.VERSION_1_8

dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/AbstractBenchmark.java

@@ -0,0 +1,138 @@
+package com.datadog.iast.propagation;
+
+import static java.util.concurrent.TimeUnit.NANOSECONDS;
+
+import com.datadog.iast.IastRequestContext;
+import com.datadog.iast.IastSystem;
+import com.datadog.iast.model.Range;
+import com.datadog.iast.model.Source;
+import datadog.trace.api.Config;
+import datadog.trace.api.gateway.InstrumentationGateway;
+import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.ScopeSource;
+import datadog.trace.bootstrap.instrumentation.api.TagContext;
+import datadog.trace.common.writer.Writer;
+import datadog.trace.core.CoreTracer;
+import datadog.trace.core.DDSpan;
+import java.util.List;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@State(Scope.Thread)
+@OutputTimeUnit(NANOSECONDS)
+@BenchmarkMode(Mode.SingleShotTime)
+@Warmup(iterations = 50_000)
+@Measurement(iterations = 5_000)
+@Fork(value = 3)
+public abstract class AbstractBenchmark<C extends AbstractBenchmark.BenchmarkContext> {
+
+  private static final Logger LOG = LoggerFactory.getLogger(AbstractBenchmark.class);
+
+  private AgentSpan span;
+  private AgentScope scope;
+  protected C context;
+
+  @Setup(Level.Trial)
+  public void setup() {
+    final InstrumentationGateway gateway = new InstrumentationGateway();
+    IastSystem.start(gateway.getSubscriptionService(RequestContextSlot.IAST));
+    final CoreTracer tracer =
+        CoreTracer.builder().instrumentationGateway(gateway).writer(new NoOpWriter()).build();
+    AgentTracer.forceRegister(tracer);
+  }
+
+  @Setup(Level.Iteration)
+  public void start() {
+    context = initializeContext();
+    final TagContext tagContext = new TagContext();
+    if (Config.get().isIastEnabled()) {
+      tagContext.withRequestContextDataIast(context.getIastContext());
+    }
+    final AgentTracer.TracerAPI tracer = AgentTracer.get();
+    span = tracer.startSpan("benchmark", tagContext, true);
+    scope = tracer.activateSpan(span, ScopeSource.INSTRUMENTATION);
+  }
+
+  @TearDown(Level.Iteration)
+  public void stop() {
+    scope.close();
+    span.finish();
+  }
+
+  protected abstract C initializeContext();
+
+  protected <E> E tainted(final IastRequestContext context, final E value, final Range... ranges) {
+    final E result = notTainted(value);
+    context.getTaintedObjects().taint(result, ranges);
+    return result;
+  }
+
+  @SuppressWarnings({"StringOperationCanBeSimplified", "unchecked"})
+  protected <E> E notTainted(final E value) {
+    final E result;
+    if (value instanceof String) {
+      result = (E) new String((String) value);
+    } else {
+      result = value;
+    }
+    computeHash(result); // compute it before to ensure all tests compare the same
+    return result;
+  }
+
+  protected Source source() {
+    return new Source((byte) 0, "key", "value");
+  }
+
+  private static long computeHash(final Object value) {
+    final long hash = System.identityHashCode(value);
+    LOG.trace("{} hash: {}", value, hash);
+    return hash;
+  }
+
+  protected abstract static class BenchmarkContext {
+
+    private final IastRequestContext iastContext;
+
+    protected BenchmarkContext(final IastRequestContext iasContext) {
+      this.iastContext = iasContext;
+    }
+
+    public IastRequestContext getIastContext() {
+      return iastContext;
+    }
+  }
+
+  private static class NoOpWriter implements Writer {
+
+    @Override
+    public void write(final List<DDSpan> trace) {}
+
+    @Override
+    public void start() {}
+
+    @Override
+    public boolean flush() {
+      return false;
+    }
+
+    @Override
+    public void close() {}
+
+    @Override
+    public void incrementDropCounts(final int spanCount) {}
+  }
+}

dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderAppendBenchmark.java

@@ -0,0 +1,98 @@
+package com.datadog.iast.propagation;
+
+import com.datadog.iast.IastRequestContext;
+import com.datadog.iast.model.Range;
+import datadog.trace.api.iast.InstrumentationBridge;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.Fork;
+
+public class StringBuilderAppendBenchmark
+    extends AbstractBenchmark<StringBuilderAppendBenchmark.Context> {
+
+  @Override
+  protected Context initializeContext() {
+    final IastRequestContext context = new IastRequestContext();
+    final String notTainted = notTainted("I am not a tainted string");
+    final String tainted = tainted(context, "I am a tainted string", new Range(5, 6, source()));
+    final StringBuilder notTaintedBuilder =
+        notTainted(new StringBuilder("I am not a tainted string builder"));
+    final StringBuilder taintedBuilder =
+        tainted(
+            context, new StringBuilder("I am a tainted string builder"), new Range(5, 6, source()));
+    return new Context(context, notTainted, tainted, notTaintedBuilder, taintedBuilder);
+  }
+
+  @Benchmark
+  @Fork(jvmArgsAppend = {"-Ddd.iast.enabled=false"})
+  public StringBuilder baseline() {
+    return context.notTaintedBuilder.append(context.notTainted);
+  }
+
+  @Benchmark
+  @Fork(jvmArgsAppend = {"-Ddd.iast.enabled=false"})
+  public StringBuilder iastDisabled() {
+    final String param = context.notTainted;
+    final StringBuilder self = context.notTaintedBuilder.append(param);
+    InstrumentationBridge.onStringBuilderAppend(self, param);
+    return self;
+  }
+
+  @Benchmark
+  @Fork(jvmArgsAppend = {"-Ddd.iast.enabled=true"})
+  public StringBuilder notTainted() {
+    final String param = context.notTainted;
+    final StringBuilder self = context.notTaintedBuilder.append(param);
+    InstrumentationBridge.onStringBuilderAppend(self, param);
+    return self;
+  }
+
+  @Benchmark
+  @Fork(jvmArgsAppend = {"-Ddd.iast.enabled=true"})
+  public StringBuilder paramTainted() {
+    final String param = context.tainted;
+    final StringBuilder self = context.notTaintedBuilder.append(param);
+    InstrumentationBridge.onStringBuilderAppend(self, param);
+    return self;
+  }
+
+  @Benchmark
+  @Fork(jvmArgsAppend = {"-Ddd.iast.enabled=true"})
+  public StringBuilder stringBuilderTainted() {
+    final String param = context.notTainted;
+    final StringBuilder self = context.taintedBuilder.append(param);
+    InstrumentationBridge.onStringBuilderAppend(self, param);
+    return self;
+  }
+
+  @Benchmark
+  @Fork(jvmArgsAppend = {"-Ddd.iast.enabled=true"})
+  public StringBuilder bothTainted() {
+    final String param = context.tainted;
+    final StringBuilder self = context.taintedBuilder.append(param);
+    InstrumentationBridge.onStringBuilderAppend(self, param);
+    return self;
+  }
+
+  protected static class Context extends AbstractBenchmark.BenchmarkContext {
+
+    private final String notTainted;
+    private final String tainted;
+
+    private final StringBuilder notTaintedBuilder;
+
+    private final StringBuilder taintedBuilder;
+
+    protected Context(
+        final IastRequestContext context,
+        final String notTainted,
+        final String tainted,
+        final StringBuilder notTaintedBuilder,
+        final StringBuilder taintedBuilder) {
+      super(context);
+      this.tainted = tainted;
+      this.notTainted = notTainted;
+      this.notTaintedBuilder = notTaintedBuilder;
+      this.taintedBuilder = taintedBuilder;
+    }
+  }
+}