Refactor of the PropagationModule and removal of WebModule (#6033)

DataDog · Oct 27, 2023 · b68b8f0 · b68b8f0
1 parent 9db54a0
commit b68b8f0
Show file tree

Hide file tree

Showing 213 changed files with 2,252 additions and 2,097 deletions.
diff --git a/...nt/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderAppendBenchmark.java b/...nt/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderAppendBenchmark.java
@@ -1,6 +1,6 @@
 package com.datadog.iast.propagation;
 
-import static com.datadog.iast.model.Range.NOT_MARKED;
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.model.Range;

diff --git a/...ent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderBatchBenchmark.java b/...ent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderBatchBenchmark.java
@@ -1,5 +1,6 @@
 package com.datadog.iast.propagation;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 import static java.util.concurrent.TimeUnit.MICROSECONDS;
 
 import com.datadog.iast.IastRequestContext;
@@ -34,8 +35,7 @@ protected StringBuilderBatchBenchmark.Context initializeContext() {
       final String value;
       if (current < limit) {
         value =
-            tainted(
-                context, UUID.randomUUID().toString(), new Range(3, 6, source(), Range.NOT_MARKED));
+            tainted(context, UUID.randomUUID().toString(), new Range(3, 6, source(), NOT_MARKED));
       } else {
         value = notTainted(UUID.randomUUID().toString());
       }

diff --git a/...gent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderInitBenchmark.java b/...gent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderInitBenchmark.java
@@ -1,5 +1,7 @@
 package com.datadog.iast.propagation;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.model.Range;
 import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
@@ -14,7 +16,7 @@ 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(3, 6, source(), Range.NOT_MARKED));
+        tainted(context, "I am a tainted string", new Range(3, 6, source(), NOT_MARKED));
     return new Context(context, notTainted, tainted);
   }
 

diff --git a/.../agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderToStringBenchmark.java b/.../agent-iast/src/jmh/java/com/datadog/iast/propagation/StringBuilderToStringBenchmark.java
@@ -1,5 +1,7 @@
 package com.datadog.iast.propagation;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.model.Range;
 import datadog.trace.instrumentation.java.lang.StringBuilderCallSite;
@@ -18,7 +20,7 @@ protected Context initializeContext() {
         tainted(
             context,
             new StringBuilder("I am a tainted string builder"),
-            new Range(5, 7, source(), Range.NOT_MARKED));
+            new Range(5, 7, source(), NOT_MARKED));
     return new Context(context, notTaintedBuilder, taintedBuilder);
   }
 

diff --git a/...ava-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatBenchmark.java b/...ava-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatBenchmark.java
@@ -1,5 +1,7 @@
 package com.datadog.iast.propagation;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.model.Range;
 import datadog.trace.instrumentation.java.lang.StringCallSite;
@@ -13,7 +15,7 @@ protected StringConcatBenchmark.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(3, 5, source(), Range.NOT_MARKED));
+        tainted(context, "I am a tainted string", new Range(3, 5, source(), NOT_MARKED));
     return new StringConcatBenchmark.Context(context, notTainted, tainted);
   }
 

diff --git a/...ent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBatchBenchmark.java b/...ent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBatchBenchmark.java
@@ -1,5 +1,6 @@
 package com.datadog.iast.propagation;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 import static java.util.concurrent.TimeUnit.MICROSECONDS;
 
 import com.datadog.iast.IastRequestContext;
@@ -54,7 +55,7 @@ protected StringConcatFactoryBatchBenchmark.Context initializeContext() {
       double current = i / (double) stringCount;
       final String value;
       if (current < limit) {
-        value = tainted(context, "Yep, tainted", new Range(3, 5, source(), Range.NOT_MARKED));
+        value = tainted(context, "Yep, tainted", new Range(3, 5, source(), NOT_MARKED));
       } else {
         value = notTainted("Nop, tainted");
       }

diff --git a/...nt/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBenchmark.java b/...nt/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringConcatFactoryBenchmark.java
@@ -1,5 +1,7 @@
 package com.datadog.iast.propagation;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.model.Range;
 import datadog.trace.api.iast.InstrumentationBridge;
@@ -13,8 +15,7 @@ public class StringConcatFactoryBenchmark
   protected StringConcatFactoryBenchmark.Context initializeContext() {
     final IastRequestContext context = new IastRequestContext();
     final String notTainted = notTainted("Nop, tainted");
-    final String tainted =
-        tainted(context, "Yep, tainted", new Range(3, 5, source(), Range.NOT_MARKED));
+    final String tainted = tainted(context, "Yep, tainted", new Range(3, 5, source(), NOT_MARKED));
     return new StringConcatFactoryBenchmark.Context(context, notTainted, tainted);
   }
 

diff --git a/dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringJoinBenchmark.java b/dd-java-agent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringJoinBenchmark.java
@@ -1,6 +1,6 @@
 package com.datadog.iast.propagation;
 
-import static com.datadog.iast.model.Range.NOT_MARKED;
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.model.Range;

diff --git a/...gent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringSubsequenceBenchmark.java b/...gent/agent-iast/src/jmh/java/com/datadog/iast/propagation/StringSubsequenceBenchmark.java
@@ -1,6 +1,6 @@
 package com.datadog.iast.propagation;
 
-import static com.datadog.iast.model.Range.NOT_MARKED;
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import com.datadog.iast.IastRequestContext;
 import com.datadog.iast.model.Range;

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/GrpcRequestMessageHandler.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/GrpcRequestMessageHandler.java
@@ -2,7 +2,7 @@
 
 import datadog.trace.api.gateway.Flow;
 import datadog.trace.api.gateway.RequestContext;
-import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.InstrumentationBridge;
 import datadog.trace.api.iast.SourceTypes;
 import datadog.trace.api.iast.propagation.PropagationModule;
@@ -29,9 +29,9 @@ public class GrpcRequestMessageHandler implements BiFunction<RequestContext, Obj
   public Flow<Void> apply(final RequestContext ctx, final Object o) {
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
     if (module != null && o != null) {
-      final IastRequestContext iastCtx = ctx.getData(RequestContextSlot.IAST);
+      final IastContext iastCtx = IastContext.Provider.get(ctx);
       module.taintDeeply(
-          iastCtx, SourceTypes.GRPC_BODY, o, GrpcRequestMessageHandler::isProtobufArtifact);
+          iastCtx, o, SourceTypes.GRPC_BODY, GrpcRequestMessageHandler::isProtobufArtifact);
     }
     return Flow.ResultFlow.empty();
   }

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastRequestContext.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastRequestContext.java
@@ -4,15 +4,15 @@
 import com.datadog.iast.overhead.OverheadContext;
 import com.datadog.iast.taint.TaintedObjects;
 import datadog.trace.api.gateway.RequestContext;
-import datadog.trace.api.gateway.RequestContextSlot;
+import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.telemetry.IastMetricCollector;
 import datadog.trace.api.iast.telemetry.IastMetricCollector.HasMetricCollector;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import java.util.concurrent.atomic.AtomicBoolean;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
-public class IastRequestContext implements HasMetricCollector {
+public class IastRequestContext implements IastContext, HasMetricCollector {
 
   private final VulnerabilityBatch vulnerabilityBatch;
   private final AtomicBoolean spanDataIsSet;
@@ -85,6 +85,7 @@ public OverheadContext getOverheadContext() {
     return overheadContext;
   }
 
+  @Nonnull
   public TaintedObjects getTaintedObjects() {
     return taintedObjects;
   }
@@ -97,22 +98,20 @@ public IastMetricCollector getMetricCollector() {
 
   @Nullable
   public static IastRequestContext get() {
-    return get(AgentTracer.activeSpan());
+    return asRequestContext(IastContext.Provider.get());
   }
 
   @Nullable
   public static IastRequestContext get(final AgentSpan span) {
-    if (span == null) {
-      return null;
-    }
-    return get(span.getRequestContext());
+    return asRequestContext(IastContext.Provider.get(span));
   }
 
   @Nullable
   public static IastRequestContext get(final RequestContext reqCtx) {
-    if (reqCtx == null) {
-      return null;
-    }
-    return reqCtx.getData(RequestContextSlot.IAST);
+    return asRequestContext(IastContext.Provider.get(reqCtx));
+  }
+
+  private static IastRequestContext asRequestContext(final IastContext ctx) {
+    return ctx instanceof IastRequestContext ? (IastRequestContext) ctx : null;
   }
 }
diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastSystem.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/IastSystem.java
@@ -23,7 +23,6 @@
 import com.datadog.iast.sink.XContentTypeModuleImpl;
 import com.datadog.iast.sink.XPathInjectionModuleImpl;
 import com.datadog.iast.sink.XssModuleImpl;
-import com.datadog.iast.source.WebModuleImpl;
 import com.datadog.iast.telemetry.TelemetryRequestEndedHandler;
 import com.datadog.iast.telemetry.TelemetryRequestStartedHandler;
 import datadog.trace.api.Config;
@@ -90,7 +89,6 @@ private static Consumer<IastModule> registerModule(final Dependencies dependenci
 
   private static Stream<IastModule> iastModules() {
     return Stream.of(
-        new WebModuleImpl(),
         new StringModuleImpl(),
         new FastCodecModule(),
         new SqlInjectionModuleImpl(),

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Range.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/Range.java
@@ -1,5 +1,7 @@
 package com.datadog.iast.model;
 
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
+
 import com.datadog.iast.model.json.SourceIndex;
 import com.datadog.iast.util.Ranged;
 import java.util.Objects;
@@ -9,8 +11,6 @@
 
 public final class Range implements Ranged {
 
-  public static final int NOT_MARKED = 0;
-
   private final @Nonnegative int start;
   private final @Nonnegative int length;
   private final @Nonnull @SourceIndex Source source;

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
@@ -11,6 +11,10 @@ public final class Source implements Taintable.Source {
   private final String name;
   private final String value;
 
+  public Source(final byte origin, final CharSequence name, final CharSequence value) {
+    this(origin, name == null ? null : name.toString(), value == null ? null : value.toString());
+  }
+
   public Source(final byte origin, final String name, final String value) {
     this.origin = origin;
     this.name = name;

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/model/VulnerabilityType.java
@@ -1,6 +1,6 @@
 package com.datadog.iast.model;
 
-import static com.datadog.iast.model.Range.NOT_MARKED;
+import static datadog.trace.api.iast.VulnerabilityMarks.NOT_MARKED;
 
 import datadog.trace.api.iast.VulnerabilityMarks;
 import datadog.trace.api.iast.VulnerabilityTypes;