Fix string format when escaped patterns are used (#6795)

DataDog · Mar 13, 2024 · ee9940d · ee9940d
1 parent cd18597
commit ee9940d
Show file tree

Hide file tree

Showing 100 changed files with 1,463 additions and 535 deletions.
diff --git a/...-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/IastHttpServerTest.groovy b/...-agent/agent-iast/src/testFixtures/groovy/com/datadog/iast/test/IastHttpServerTest.groovy
@@ -33,7 +33,7 @@ abstract class IastHttpServerTest<SERVER> extends WithHttpServer<SERVER> impleme
       if (iastRequestContext) {
         TaintedObjects taintedObjects = iastRequestContext.getTaintedObjects()
         TAINTED_OBJECTS.offer(new TaintedObjectCollection(taintedObjects))
-        List<Vulnerability> vulns = iastRequestContext.getVulnerabilityBatch().getVulnerabilities()
+        List<Vulnerability> vulns = iastRequestContext.getVulnerabilityBatch().getVulnerabilities() ?: []
         VULNERABILITIES.offer(vulns)
       }
     }

diff --git a/...nt/agent-tooling/src/main/java/datadog/trace/agent/tooling/iast/TaintableEnumeration.java b/...nt/agent-tooling/src/main/java/datadog/trace/agent/tooling/iast/TaintableEnumeration.java
@@ -11,8 +11,7 @@ public class TaintableEnumeration implements Enumeration<String> {
 
   private static final String CLASS_NAME = TaintableEnumeration.class.getName();
 
-  private volatile IastContext context;
-  private volatile boolean contextFetched;
+  private final IastContext context;
 
   private final PropagationModule module;
 
@@ -25,11 +24,13 @@ public class TaintableEnumeration implements Enumeration<String> {
   private final Enumeration<String> delegate;
 
   private TaintableEnumeration(
+      final IastContext ctx,
       @NonNull final Enumeration<String> delegate,
       @NonNull final PropagationModule module,
       final byte origin,
       @Nullable final CharSequence name,
       final boolean useValueAsName) {
+    this.context = ctx;
     this.delegate = delegate;
     this.module = module;
     this.origin = origin;
@@ -57,21 +58,13 @@ public String nextElement() {
       throw e;
     }
     try {
-      module.taint(context(), next, origin, name(next));
+      module.taint(context, next, origin, name(next));
     } catch (final Throwable e) {
       module.onUnexpectedException("Failed to taint enumeration", e);
     }
     return next;
   }
 
-  private IastContext context() {
-    if (!contextFetched) {
-      contextFetched = true;
-      context = IastContext.Provider.get();
-    }
-    return context;
-  }
-
   private CharSequence name(final String value) {
     if (name != null) {
       return name;
@@ -84,18 +77,20 @@ private static boolean nonTaintableEnumerationStack(final StackTraceElement elem
   }
 
   public static Enumeration<String> wrap(
+      final IastContext ctx,
       @NonNull final Enumeration<String> delegate,
       @NonNull final PropagationModule module,
       final byte origin,
       @Nullable final CharSequence name) {
-    return new TaintableEnumeration(delegate, module, origin, name, false);
+    return new TaintableEnumeration(ctx, delegate, module, origin, name, false);
   }
 
   public static Enumeration<String> wrap(
+      final IastContext ctx,
       @NonNull final Enumeration<String> delegate,
       @NonNull final PropagationModule module,
       final byte origin,
       boolean useValueAsName) {
-    return new TaintableEnumeration(delegate, module, origin, null, useValueAsName);
+    return new TaintableEnumeration(ctx, delegate, module, origin, null, useValueAsName);
   }
 }
diff --git a/...-tooling/src/test/groovy/datadog/trace/agent/tooling/iast/TaintableEnumerationTest.groovy b/...-tooling/src/test/groovy/datadog/trace/agent/tooling/iast/TaintableEnumerationTest.groovy
@@ -1,44 +1,26 @@
 package datadog.trace.agent.tooling.iast
 
-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
-import datadog.trace.bootstrap.instrumentation.api.AgentSpan
-import datadog.trace.bootstrap.instrumentation.api.AgentTracer
 import datadog.trace.test.util.DDSpecification
 import spock.lang.Shared
 
 class TaintableEnumerationTest extends DDSpecification {
 
   @Shared
-  protected static final AgentTracer.TracerAPI ORIGINAL_TRACER = AgentTracer.get()
-
-  protected AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI)
-
-  protected IastContext iastCtx = Mock(IastContext)
-
-  protected RequestContext reqCtx = Mock(RequestContext) {
-    getData(RequestContextSlot.IAST) >> iastCtx
-  }
-
-  protected AgentSpan span = Mock(AgentSpan) {
-    getRequestContext() >> reqCtx
-  }
+  protected IastContext iastCtx = Stub(IastContext)
 
   protected PropagationModule module
 
-
   void setup() {
-    AgentTracer.forceRegister(tracer)
     module = Mock(PropagationModule)
     InstrumentationBridge.registerIastModule(module)
   }
 
   void cleanup() {
-    AgentTracer.forceRegister(ORIGINAL_TRACER)
     InstrumentationBridge.clearIastModules()
   }
 
@@ -47,35 +29,34 @@ class TaintableEnumerationTest extends DDSpecification {
     final values = (1..10).collect { "value$it".toString() }
     final origin = SourceTypes.REQUEST_PARAMETER_NAME
     final name = 'test'
-    final enumeration = TaintableEnumeration.wrap(Collections.enumeration(values), module, origin, name)
+    final enumeration = TaintableEnumeration.wrap(iastCtx, Collections.enumeration(values), module, origin, name)
 
     when:
     final result = enumeration.collect()
 
     then:
     result == values
-    values.each { 1 * module.taint(_, it, origin, name) }
-    1 * tracer.activeSpan() >> span // only one access to the active context
+    values.each { 1 * module.taint(iastCtx, it, origin, name) }
   }
 
   void 'underlying enumerated values are tainted with the value as a name'() {
     given:
     final values = (1..10).collect { "value$it".toString() }
     final origin = SourceTypes.REQUEST_PARAMETER_NAME
-    final enumeration = TaintableEnumeration.wrap(Collections.enumeration(values), module, origin, true)
+    final enumeration = TaintableEnumeration.wrap(iastCtx, Collections.enumeration(values), module, origin, true)
 
     when:
     final result = enumeration.collect()
 
     then:
     result == values
-    values.each { 1 * module.taint(_, it, origin, it) }
+    values.each { 1 * module.taint(iastCtx, it, origin, it) }
   }
 
   void 'taintable enumeration leaves no trace in case of error'() {
     given:
     final origin = SourceTypes.REQUEST_PARAMETER_NAME
-    final enumeration = TaintableEnumeration.wrap(new BadEnumeration(), module, origin, true)
+    final enumeration = TaintableEnumeration.wrap(iastCtx, new BadEnumeration(), module, origin, true)
 
     when:
     enumeration.hasMoreElements()

diff --git a/...rc/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java b/...rc/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java
@@ -11,8 +11,12 @@
 import akka.http.scaladsl.model.headers.Cookie;
 import akka.http.scaladsl.model.headers.HttpCookiePair;
 import com.google.auto.service.AutoService;
+import datadog.trace.advice.ActiveRequestContext;
+import datadog.trace.advice.RequiresRequestContext;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+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.Source;
@@ -50,20 +54,23 @@ public void methodAdvice(MethodTransformer transformer) {
         CookieHeaderInstrumentation.class.getName() + "$TaintAllCookiesAdvice");
   }
 
+  @RequiresRequestContext(RequestContextSlot.IAST)
   static class TaintAllCookiesAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     @Source(SourceTypes.REQUEST_COOKIE_VALUE)
     static void after(
-        @Advice.This HttpHeader cookie, @Advice.Return Seq<HttpCookiePair> cookiePairs) {
+        @Advice.This HttpHeader cookie,
+        @Advice.Return Seq<HttpCookiePair> cookiePairs,
+        @ActiveRequestContext RequestContext reqCtx) {
       PropagationModule prop = InstrumentationBridge.PROPAGATION;
       if (prop == null || cookiePairs == null || cookiePairs.isEmpty()) {
         return;
       }
-      if (!prop.isTainted(cookie)) {
+      final IastContext ctx = IastContext.Provider.get(reqCtx);
+      if (!prop.isTainted(ctx, cookie)) {
         return;
       }
 
-      final IastContext ctx = IastContext.Provider.get();
       Iterator<HttpCookiePair> iterator = cookiePairs.iterator();
       while (iterator.hasNext()) {
         HttpCookiePair pair = iterator.next();

diff --git a/...tp-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HeaderNameCallSite.java b/...tp-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/HeaderNameCallSite.java
@@ -3,6 +3,7 @@
 import akka.http.javadsl.model.HttpHeader;
 import datadog.trace.agent.tooling.csi.CallSite;
 import datadog.trace.api.iast.IastCallSites;
+import datadog.trace.api.iast.IastContext;
 import datadog.trace.api.iast.InstrumentationBridge;
 import datadog.trace.api.iast.Source;
 import datadog.trace.api.iast.SourceTypes;
@@ -26,7 +27,11 @@ public static String after(@CallSite.This HttpHeader header, @CallSite.Return St
       return result;
     }
     try {
-      module.taintIfTainted(result, header, SourceTypes.REQUEST_HEADER_NAME, result);
+      final IastContext ctx = IastContext.Provider.get();
+      if (ctx == null) {
+        return result;
+      }
+      module.taintIfTainted(ctx, result, header, SourceTypes.REQUEST_HEADER_NAME, result);
     } catch (final Throwable e) {
       module.onUnexpectedException("onHeaderNames threw", e);
     }

diff --git a/...java/datadog/trace/instrumentation/akkahttp/iast/HttpHeaderSubclassesInstrumentation.java b/...java/datadog/trace/instrumentation/akkahttp/iast/HttpHeaderSubclassesInstrumentation.java
@@ -11,8 +11,13 @@
 import akka.http.scaladsl.model.HttpHeader;
 import akka.http.scaladsl.model.HttpRequest;
 import com.google.auto.service.AutoService;
+import datadog.trace.advice.ActiveRequestContext;
+import datadog.trace.advice.RequiresRequestContext;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+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.Propagation;
 import datadog.trace.api.iast.propagation.PropagationModule;
@@ -53,17 +58,21 @@ public void methodAdvice(MethodTransformer transformer) {
         HttpHeaderSubclassesInstrumentation.class.getName() + "$HttpHeaderSubclassesAdvice");
   }
 
+  @RequiresRequestContext(RequestContextSlot.IAST)
   static class HttpHeaderSubclassesAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     @Propagation
-    static void onExit(@Advice.This HttpHeader h, @Advice.Return String retVal) {
+    static void onExit(
+        @Advice.This HttpHeader h,
+        @Advice.Return String retVal,
+        @ActiveRequestContext RequestContext reqCtx) {
 
       PropagationModule propagation = InstrumentationBridge.PROPAGATION;
       if (propagation == null) {
         return;
       }
-
-      propagation.taintIfTainted(retVal, h);
+      IastContext ctx = IastContext.Provider.get(reqCtx);
+      propagation.taintIfTainted(ctx, retVal, h);
     }
   }
 }
diff --git a/...src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpRequestInstrumentation.java b/...src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpRequestInstrumentation.java
@@ -12,8 +12,12 @@
 import akka.http.scaladsl.model.HttpHeader;
 import akka.http.scaladsl.model.HttpRequest;
 import com.google.auto.service.AutoService;
+import datadog.trace.advice.ActiveRequestContext;
+import datadog.trace.advice.RequiresRequestContext;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
+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.Propagation;
@@ -58,25 +62,28 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   @SuppressFBWarnings("BC_IMPOSSIBLE_INSTANCEOF")
+  @RequiresRequestContext(RequestContextSlot.IAST)
   static class RequestHeadersAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     @Source(SourceTypes.REQUEST_HEADER_VALUE)
     static void onExit(
-        @Advice.This HttpRequest thiz, @Advice.Return(readOnly = false) Seq<HttpHeader> headers) {
+        @Advice.This HttpRequest thiz,
+        @Advice.Return(readOnly = false) Seq<HttpHeader> headers,
+        @ActiveRequestContext RequestContext reqCtx) {
       PropagationModule propagation = InstrumentationBridge.PROPAGATION;
       if (propagation == null || headers == null || headers.isEmpty()) {
         return;
       }
 
-      if (!propagation.isTainted(thiz)) {
+      final IastContext ctx = IastContext.Provider.get(reqCtx);
+      if (!propagation.isTainted(ctx, thiz)) {
         return;
       }
 
-      final IastContext ctx = IastContext.Provider.get();
       Iterator<HttpHeader> iterator = headers.iterator();
       while (iterator.hasNext()) {
         HttpHeader h = iterator.next();
-        if (propagation.isTainted(h)) {
+        if (propagation.isTainted(ctx, h)) {
           continue;
         }
         // unfortunately, the call to h.value() is instrumented, but
@@ -86,22 +93,26 @@ static void onExit(
     }
   }
 
+  @RequiresRequestContext(RequestContextSlot.IAST)
   static class EntityAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     @Propagation
     static void onExit(
         @Advice.This HttpRequest thiz,
-        @Advice.Return(readOnly = false, typing = DYNAMIC) Object entity) {
+        @Advice.Return(readOnly = false, typing = DYNAMIC) Object entity,
+        @ActiveRequestContext RequestContext reqCtx) {
+
       PropagationModule propagation = InstrumentationBridge.PROPAGATION;
       if (propagation == null || entity == null) {
         return;
       }
 
-      if (propagation.isTainted(entity)) {
+      IastContext ctx = IastContext.Provider.get(reqCtx);
+      if (propagation.isTainted(ctx, entity)) {
         return;
       }
 
-      propagation.taintIfTainted(entity, thiz);
+      propagation.taintIfTainted(ctx, entity, thiz);
     }
   }
 }
diff --git a/...src/main/java/datadog/trace/instrumentation/akkahttp/iast/PathMatcherInstrumentation.java b/...src/main/java/datadog/trace/instrumentation/akkahttp/iast/PathMatcherInstrumentation.java
@@ -64,13 +64,14 @@ static void onExit(
       scala.Tuple1 tuple = (scala.Tuple1) extractions;
       Object value = tuple._1();
 
+      final IastContext ctx = IastContext.Provider.get(reqCtx);
+
       // in the test, 4 instances of PathMatcher$Match are created, all with the same value
-      if (module.isTainted(value)) {
+      if (module.isTainted(ctx, value)) {
         return;
       }
 
       if (value instanceof String) {
-        final IastContext ctx = reqCtx.getData(RequestContextSlot.IAST);
         module.taint(ctx, value, SourceTypes.REQUEST_PATH_PARAMETER);
       }
     }