Initial version of the native agent with tagging

Author: manuel-alvarez-alvarez

benchmark/benchmarks.sh

@@ -8,6 +8,12 @@ export UTILS_DIR="${SCRIPT_DIR}/utils"
 export SHELL_UTILS_DIR="${UTILS_DIR}/shell"
 export K6_UTILS_DIR="${UTILS_DIR}/k6"
 export TRACER="${SCRIPT_DIR}/tracer/dd-java-agent.jar"
+NATIVE_LIBRARY=$(readlink --canonicalize "${SCRIPT_DIR}/../native-agent/libdd-java-agent.so")
+if [ -f "$NATIVE_LIBRARY" ]; then
+  export NATIVE_TRACER="-agentpath:$NATIVE_LIBRARY"
+else
+  export NATIVE_TRACER=""
+fi
 export NO_AGENT_VARIANT="no_agent"
 
 run_benchmarks() {

benchmark/dacapo/benchmark.json

@@ -35,6 +35,12 @@
         "JAVA_OPTS": "-javaagent:${TRACER} -Ddd.iast.enabled=true"
       }
     },
+    "iast_NATIVE": {
+      "env": {
+        "VARIANT": "iast_NATIVE",
+        "JAVA_OPTS": "${NATIVE_TRACER} -javaagent:${TRACER} -Ddd.iast.enabled=true"
+      }
+    },
     "iast_GLOBAL": {
       "env": {
         "VARIANT": "iast_GLOBAL",

benchmark/load/insecure-bank/benchmark.json

@@ -24,6 +24,12 @@
         "JAVA_OPTS": "-javaagent:${TRACER} -Ddd.iast.enabled=true"
       }
     },
+    "iast_NATIVE": {
+      "env": {
+        "VARIANT": "iast_NATIVE",
+        "JAVA_OPTS": "${NATIVE_TRACER} -javaagent:${TRACER} -Ddd.iast.enabled=true"
+      }
+    },
     "iast_GLOBAL": {
       "env": {
         "VARIANT": "iast_GLOBAL",
@@ -36,6 +42,12 @@
         "JAVA_OPTS": "-javaagent:${TRACER} -Ddd.iast.enabled=true -Ddd.iast.detection.mode=FULL"
       }
     },
+    "iast_NATIVE_FULL": {
+      "env": {
+        "VARIANT": "iast_NATIVE_FULL",
+        "JAVA_OPTS": "${NATIVE_TRACER} -javaagent:${TRACER} -Ddd.iast.enabled=true -Ddd.iast.detection.mode=FULL"
+      }
+    },
     "iast_INACTIVE": {
       "env": {
         "VARIANT": "iast_INACTIVE",

benchmark/load/petclinic/benchmark.json

@@ -41,6 +41,12 @@
         "VARIANT": "iast",
         "JAVA_OPTS": "-javaagent:${TRACER} -Ddd.iast.enabled=true"
       }
+    },
+    "iast_NATIVE": {
+      "env": {
+        "VARIANT": "iast_NATIVE",
+        "JAVA_OPTS": "${NATIVE_TRACER} -javaagent:${TRACER} -Ddd.iast.enabled=true"
+      }
     }
   }
 }

dd-smoke-tests/iast-util/src/testFixtures/groovy/datadog/smoketest/AbstractIastServerSmokeTest.groovy

@@ -12,6 +12,8 @@ import org.spockframework.runtime.SpockTimeoutError
 import spock.lang.Shared
 import spock.util.concurrent.PollingConditions
 
+import java.nio.file.Files
+import java.nio.file.Paths
 import java.util.concurrent.TimeoutException
 
 @CompileDynamic
@@ -22,6 +24,21 @@ abstract class AbstractIastServerSmokeTest extends AbstractServerSmokeTest {
   @Shared
   private final JsonSlurper jsonSlurper = new JsonSlurper()
 
+  @Override
+  def javaProperties() {
+    def source = Paths.get(AbstractIastServerSmokeTest.protectionDomain.codeSource.location.toURI())
+    while (source != null && !source.endsWith('dd-trace-java')) {
+      source = source.parent
+    }
+    final nativeLib = source.resolve('native-agent').resolve('libdd-java-agent.so')
+    if (!Files.exists(nativeLib)) {
+      throw new IllegalArgumentException('Native agent missing')
+    }
+    final props = ["-agentpath:$nativeLib".toString()]
+    super.javaProperties().each { props.add(it) }
+    return props.toArray(new String[0])
+  }
+
   @Override
   String logLevel() {
     return 'debug'

internal-api/build.gradle

@@ -211,6 +211,9 @@ dependencies {
   // it contains annotations that are also present in the instrumented application classes
   api "com.datadoghq:dd-javac-plugin-client:0.2.2"
 
+  // native agent that should be loaded by the system classloader
+  api files("$rootDir/native-agent/dd-native-java-agent-0.1.0-SNAPSHOT.jar")
+
   testImplementation project(":utils:test-utils")
   testImplementation("org.assertj:assertj-core:3.20.2")
   testImplementation libs.bundles.junit5

internal-api/src/main/java/datadog/trace/api/iast/IastContext.java

@@ -2,7 +2,9 @@
 
 import datadog.trace.api.gateway.RequestContext;
 import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.api.iast.taint.NativeTaintedObjectsAdapter;
 import datadog.trace.api.iast.taint.TaintedObjects;
+import datadog.trace.api.nagent.iast.NativeTaintedUtils;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.io.Closeable;
@@ -50,14 +52,27 @@ public static void register(@Nonnull final Provider instance) {
     /** Release the current request context, e.g. free resources, add objects to pools, ... */
     public abstract void releaseRequestContext(@Nonnull IastContext context);
 
+    /**
+     * Returns the current active tainted objects dictionary reusing the native stub if available
+     *
+     * @see #taintedObjectsFn()
+     */
+    @Nullable
+    public static TaintedObjects taintedObjects() {
+      if (NativeTaintedUtils.isEnabled()) {
+        return new NativeTaintedObjectsAdapter(Provider::taintedObjectsFn);
+      } else {
+        return taintedObjectsFn();
+      }
+    }
+
     /**
      * Gets the current active tainted objects dictionary, if no provider is configured this method
      * defaults to fetching the dictionary included in the current request context
      *
      * @see Provider#taintedObjects()
      */
-    @Nullable
-    public static TaintedObjects taintedObjects() {
+    private static TaintedObjects taintedObjectsFn() {
       if (INSTANCE == null) {
         final IastContext ctx = get();
         return ctx == null ? null : ctx.getTaintedObjects();

internal-api/src/main/java/datadog/trace/api/iast/taint/NativeTaintedObjectsAdapter.java

@@ -0,0 +1,65 @@
+package datadog.trace.api.iast.taint;
+
+import datadog.trace.api.nagent.iast.NativeTaintedUtils;
+import java.util.Iterator;
+import java.util.function.Supplier;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class NativeTaintedObjectsAdapter implements TaintedObjects {
+
+  private final Supplier<TaintedObjects> delegateSupplier;
+  private volatile TaintedObjects delegate;
+
+  public NativeTaintedObjectsAdapter(@Nonnull final Supplier<TaintedObjects> delegateSupplier) {
+    this.delegateSupplier = delegateSupplier;
+  }
+
+  @Nullable
+  @Override
+  public TaintedObject taint(@Nonnull final Object obj, @Nonnull final Range[] ranges) {
+    NativeTaintedUtils.setTainted(obj);
+    return delegate().taint(obj, ranges);
+  }
+
+  @Nullable
+  @Override
+  public TaintedObject get(@Nonnull final Object obj) {
+    if (isTainted(obj)) {
+      return delegate().get(obj);
+    }
+    return null;
+  }
+
+  @Override
+  public boolean isTainted(@Nonnull final Object obj) {
+    return NativeTaintedUtils.isTainted(obj);
+  }
+
+  @Override
+  public void clear() {
+    delegate().clear();
+  }
+
+  @Override
+  public int count() {
+    return delegate().count();
+  }
+
+  @Nonnull
+  @Override
+  public Iterator<TaintedObject> iterator() {
+    return delegate().iterator();
+  }
+
+  private TaintedObjects delegate() {
+    if (delegate == null) {
+      synchronized (this) {
+        if (delegate == null) {
+          delegate = delegateSupplier.get();
+        }
+      }
+    }
+    return delegate;
+  }
+}