Add missing checks for an active request context

DataDog · Mar 14, 2024 · e78a54c · e78a54c
1 parent bf8602f
commit e78a54c
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 6 deletions.
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
@@ -8,6 +8,7 @@
 import datadog.trace.api.iast.Source;
 import datadog.trace.api.iast.SourceTypes;
 import datadog.trace.api.iast.propagation.PropagationModule;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 
 /**
  * Detects when a header name is directly called from user code. This uses call site instrumentation
@@ -27,7 +28,7 @@ public static String after(@CallSite.This HttpHeader header, @CallSite.Return St
       return result;
     }
     try {
-      final IastContext ctx = IastContext.Provider.get();
+      final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan());
       if (ctx == null) {
         return result;
       }

diff --git a/...tp-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HeaderNameCallSite.java b/...tp-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/HeaderNameCallSite.java
@@ -7,6 +7,7 @@
 import datadog.trace.api.iast.Source;
 import datadog.trace.api.iast.SourceTypes;
 import datadog.trace.api.iast.propagation.PropagationModule;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import org.apache.pekko.http.javadsl.model.HttpHeader;
 
 /**
@@ -28,7 +29,7 @@ public static String after(@CallSite.This HttpHeader header, @CallSite.Return St
       return result;
     }
     try {
-      final IastContext ctx = IastContext.Provider.get();
+      final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan());
       if (ctx == null) {
         return result;
       }

diff --git a/...c/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestCallSite.java b/...c/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestCallSite.java
@@ -7,6 +7,7 @@
 import datadog.trace.api.iast.Source;
 import datadog.trace.api.iast.SourceTypes;
 import datadog.trace.api.iast.propagation.PropagationModule;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import jakarta.servlet.http.HttpServletRequest;
 
 /**
@@ -30,7 +31,7 @@ public static String afterPath(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         try {
-          final IastContext ctx = IastContext.Provider.get();
+          final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan());
           if (ctx != null) {
             module.taint(ctx, retValue, SourceTypes.REQUEST_PATH);
           }
@@ -52,7 +53,7 @@ public static StringBuffer afterGetRequestURL(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         try {
-          final IastContext ctx = IastContext.Provider.get();
+          final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan());
           if (ctx != null) {
             module.taint(ctx, retValue, SourceTypes.REQUEST_URI);
           }

diff --git a/.../src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestCallSite.java b/.../src/main/java/datadog/trace/instrumentation/servlet/http/HttpServletRequestCallSite.java
@@ -7,6 +7,7 @@
 import datadog.trace.api.iast.Source;
 import datadog.trace.api.iast.SourceTypes;
 import datadog.trace.api.iast.propagation.PropagationModule;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import javax.servlet.http.HttpServletRequest;
 
 /**
@@ -30,7 +31,7 @@ public static String afterPath(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         try {
-          final IastContext ctx = IastContext.Provider.get();
+          final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan());
           if (ctx != null) {
             module.taint(ctx, retValue, SourceTypes.REQUEST_PATH);
           }
@@ -52,7 +53,7 @@ public static StringBuffer afterGetRequestURL(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         try {
-          final IastContext ctx = IastContext.Provider.get();
+          final IastContext ctx = IastContext.Provider.get(AgentTracer.activeSpan());
           if (ctx != null) {
             module.taint(ctx, retValue, SourceTypes.REQUEST_URI);
           }