Use inner class inside IastContext to prevent issues in instrumentations

DataDog · Oct 22, 2023 · 0efacb4 · 0efacb4
1 parent b74e1b3
commit 0efacb4
Show file tree

Hide file tree

Showing 45 changed files with 90 additions and 81 deletions.
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
@@ -29,7 +29,7 @@ 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 IastContext iastCtx = IastContext.get(ctx);
+      final IastContext iastCtx = IastContext.Provider.get(ctx);
       module.taintDeeply(
           iastCtx, o, SourceTypes.GRPC_BODY, GrpcRequestMessageHandler::isProtobufArtifact);
     }

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
@@ -98,17 +98,17 @@ public IastMetricCollector getMetricCollector() {
 
   @Nullable
   public static IastRequestContext get() {
-    return asRequestContext(IastContext.get());
+    return asRequestContext(IastContext.Provider.get());
   }
 
   @Nullable
   public static IastRequestContext get(final AgentSpan span) {
-    return asRequestContext(IastContext.get(span));
+    return asRequestContext(IastContext.Provider.get(span));
   }
 
   @Nullable
   public static IastRequestContext get(final RequestContext reqCtx) {
-    return asRequestContext(IastContext.get(reqCtx));
+    return asRequestContext(IastContext.Provider.get(reqCtx));
   }
 
   private static IastRequestContext asRequestContext(final IastContext ctx) {

diff --git a/...rc/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java b/...rc/main/java/datadog/trace/instrumentation/akkahttp/iast/CookieHeaderInstrumentation.java
@@ -62,7 +62,7 @@ static void after(
         return;
       }
 
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       Iterator<HttpCookiePair> iterator = cookiePairs.iterator();
       while (iterator.hasNext()) {
         HttpCookiePair pair = iterator.next();

diff --git a/...src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpRequestInstrumentation.java b/...src/main/java/datadog/trace/instrumentation/akkahttp/iast/HttpRequestInstrumentation.java
@@ -71,7 +71,7 @@ static void onExit(
         return;
       }
 
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       Iterator<HttpHeader> iterator = headers.iterator();
       while (iterator.hasNext()) {
         HttpHeader h = iterator.next();

diff --git a/...tp-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/UriInstrumentation.java b/...tp-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/UriInstrumentation.java
@@ -88,7 +88,7 @@ static void after(@Advice.This /*Uri*/ Object uri, @Advice.Return Uri.Query ret)
         return;
       }
 
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       Iterator<Tuple2<String, String>> iterator = ret.iterator();
       while (iterator.hasNext()) {
         Tuple2<String, String> pair = iterator.next();

diff --git a/...0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMapFunction.java b/...0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMapFunction.java
@@ -23,7 +23,7 @@ public Tuple1<Map<String, String>> apply(Tuple1<Map<String, String>> v1) {
       return v1;
     }
 
-    final IastContext ctx = IastContext.get();
+    final IastContext ctx = IastContext.Provider.get();
     Iterator<Tuple2<String, String>> iterator = m.iterator();
     while (iterator.hasNext()) {
       Tuple2<String, String> e = iterator.next();

diff --git a/.../main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMultiMapFunction.java b/.../main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintMultiMapFunction.java
@@ -24,7 +24,7 @@ public Tuple1<Map<String, List<String>>> apply(Tuple1<Map<String, List<String>>>
       return v1;
     }
 
-    final IastContext ctx = IastContext.get();
+    final IastContext ctx = IastContext.Provider.get();
     Iterator<Tuple2<String, List<String>>> entriesIterator = m.iterator();
     while (entriesIterator.hasNext()) {
       Tuple2<String, List<String>> e = entriesIterator.next();

diff --git a/...0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSeqFunction.java b/...0/src/main/java/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSeqFunction.java
@@ -27,7 +27,7 @@ public Tuple1<Seq<Tuple2<String, String>>> apply(Tuple1<Seq<Tuple2<String, Strin
       return v1;
     }
 
-    final IastContext ctx = IastContext.get();
+    final IastContext ctx = IastContext.Provider.get();
     Iterator<Tuple2<String, String>> iterator = seq.iterator();
     Set<String> seenKeys = Collections.newSetFromMap(new IdentityHashMap<>());
     while (iterator.hasNext()) {

diff --git a/...ava/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSingleParameterFunction.java b/...ava/datadog/trace/instrumentation/akkahttp/iast/helpers/TaintSingleParameterFunction.java
@@ -41,7 +41,7 @@ public Tuple1<Object> apply(Tuple1<Object> v1) {
     }
 
     if (value instanceof Iterable) {
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       Iterator<?> iterator = ((Iterable<?>) value).iterator();
       while (iterator.hasNext()) {
         Object o = iterator.next();

diff --git a/.../java/datadog/trace/instrumentation/akkahttp102/iast/helpers/TaintParametersFunction.java b/.../java/datadog/trace/instrumentation/akkahttp102/iast/helpers/TaintParametersFunction.java
@@ -34,7 +34,7 @@ public Tuple1<T> apply(Tuple1<T> v1) {
     }
 
     if (value instanceof Iterable) {
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       Iterator<?> iterator = ((Iterable<?>) value).iterator();
       while (iterator.hasNext()) {
         Object o = iterator.next();

diff --git a/...c/main/java/datadog/trace/instrumentation/jersey/AbstractFormProviderInstrumentation.java b/...c/main/java/datadog/trace/instrumentation/jersey/AbstractFormProviderInstrumentation.java
@@ -43,7 +43,7 @@ public static void onExit(@Advice.Return Map<String, List<String>> result) {
       if (prop == null || result == null || result.isEmpty()) {
         return;
       }
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       for (Map.Entry<String, List<String>> entry : result.entrySet()) {
         final String name = entry.getKey();
         prop.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);

diff --git a/.../main/java/datadog/trace/instrumentation/jersey/InboundMessageContextInstrumentation.java b/.../main/java/datadog/trace/instrumentation/jersey/InboundMessageContextInstrumentation.java
@@ -46,7 +46,7 @@ public static class InstrumenterAdviceGetHeaders {
     public static void onExit(@Advice.Return Map<String, List<String>> headers) {
       final PropagationModule prop = InstrumentationBridge.PROPAGATION;
       if (prop != null && headers != null && !headers.isEmpty()) {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
           final String name = entry.getKey();
           prop.taint(ctx, name, SourceTypes.REQUEST_HEADER_NAME);
@@ -64,7 +64,7 @@ public static class InstrumenterAdviceGetRequestCookies {
     public static void onExit(@Advice.Return Map<String, Object> cookies) {
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null && cookies != null && !cookies.isEmpty()) {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (Map.Entry<String, Object> entry : cookies.entrySet()) {
           final String name = entry.getKey();
           module.taint(ctx, name, SourceTypes.REQUEST_COOKIE_NAME);

diff --git a/...c/main/java/datadog/trace/instrumentation/pekkohttp/iast/CookieHeaderInstrumentation.java b/...c/main/java/datadog/trace/instrumentation/pekkohttp/iast/CookieHeaderInstrumentation.java
@@ -58,7 +58,7 @@ static void after(
       if (prop == null || cookiePairs == null || cookiePairs.isEmpty()) {
         return;
       }
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       if (!prop.isTainted(ctx, cookie)) {
         return;
       }

diff --git a/...rc/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpRequestInstrumentation.java b/...rc/main/java/datadog/trace/instrumentation/pekkohttp/iast/HttpRequestInstrumentation.java
@@ -71,7 +71,7 @@ static void onExit(
         return;
       }
 
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       Iterator<HttpHeader> iterator = headers.iterator();
       while (iterator.hasNext()) {
         HttpHeader h = iterator.next();

diff --git a/...tp-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/UriInstrumentation.java b/...tp-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/iast/UriInstrumentation.java
@@ -84,7 +84,7 @@ static void after(@Advice.This /*Uri*/ Object uri, @Advice.Return Uri.Query ret)
         return;
       }
 
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       if (!prop.isTainted(ctx, uri)) {
         return;
       }

diff --git a/.../src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMapFunction.java b/.../src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMapFunction.java
@@ -23,7 +23,7 @@ public Tuple1<Map<String, String>> apply(Tuple1<Map<String, String>> v1) {
       return v1;
     }
 
-    final IastContext ctx = IastContext.get();
+    final IastContext ctx = IastContext.Provider.get();
     Iterator<Tuple2<String, String>> iterator = m.iterator();
     while (iterator.hasNext()) {
       Tuple2<String, String> e = iterator.next();

diff --git a/...main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMultiMapFunction.java b/...main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintMultiMapFunction.java
@@ -24,7 +24,7 @@ public Tuple1<Map<String, List<String>>> apply(Tuple1<Map<String, List<String>>>
       return v1;
     }
 
-    final IastContext ctx = IastContext.get();
+    final IastContext ctx = IastContext.Provider.get();
     Iterator<Tuple2<String, List<String>>> entriesIterator = m.iterator();
     while (entriesIterator.hasNext()) {
       Tuple2<String, List<String>> e = entriesIterator.next();

diff --git a/...in/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintParametersFunction.java b/...in/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintParametersFunction.java
@@ -34,7 +34,7 @@ public Tuple1<T> apply(Tuple1<T> v1) {
     }
 
     if (value instanceof Iterable) {
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       Iterator<?> iterator = ((Iterable<?>) value).iterator();
       while (iterator.hasNext()) {
         Object o = iterator.next();

diff --git a/.../src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSeqFunction.java b/.../src/main/java/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSeqFunction.java
@@ -27,7 +27,7 @@ public Tuple1<Seq<Tuple2<String, String>>> apply(Tuple1<Seq<Tuple2<String, Strin
       return v1;
     }
 
-    final IastContext ctx = IastContext.get();
+    final IastContext ctx = IastContext.Provider.get();
     Iterator<Tuple2<String, String>> iterator = seq.iterator();
     Set<String> seenKeys = Collections.newSetFromMap(new IdentityHashMap<>());
     while (iterator.hasNext()) {

diff --git a/...va/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSingleParameterFunction.java b/...va/datadog/trace/instrumentation/pekkohttp/iast/helpers/TaintSingleParameterFunction.java
@@ -41,7 +41,7 @@ public Tuple1<Object> apply(Tuple1<Object> v1) {
     }
 
     if (value instanceof Iterable) {
-      final IastContext ctx = IastContext.get();
+      final IastContext ctx = IastContext.Provider.get();
       Iterator<?> iterator = ((Iterable<?>) value).iterator();
       while (iterator.hasNext()) {
         Object o = iterator.next();

diff --git a/...ppsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorAdvice.java b/...ppsec/src/main/java/datadog/trace/instrumentation/resteasy/CookieParamInjectorAdvice.java
@@ -18,7 +18,7 @@ public static void onExit(
       if (module != null) {
         if (result instanceof Collection) {
           Collection<?> collection = (Collection<?>) result;
-          final IastContext ctx = IastContext.get();
+          final IastContext ctx = IastContext.Provider.get();
           for (Object o : collection) {
             if (o instanceof String) {
               module.taint(ctx, (String) o, SourceTypes.REQUEST_COOKIE_VALUE, paramName);

diff --git a/...-appsec/src/main/java/datadog/trace/instrumentation/resteasy/FormParamInjectorAdvice.java b/...-appsec/src/main/java/datadog/trace/instrumentation/resteasy/FormParamInjectorAdvice.java
@@ -17,7 +17,7 @@ public static void onExit(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         if (result instanceof Collection) {
-          final IastContext ctx = IastContext.get();
+          final IastContext ctx = IastContext.Provider.get();
           Collection<?> collection = (Collection<?>) result;
           for (Object o : collection) {
             if (o instanceof String) {

diff --git a/...ppsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorAdvice.java b/...ppsec/src/main/java/datadog/trace/instrumentation/resteasy/HeaderParamInjectorAdvice.java
@@ -19,7 +19,7 @@ public static void onExit(
         try {
           if (result instanceof Collection) {
             Collection<?> collection = (Collection<?>) result;
-            final IastContext ctx = IastContext.get();
+            final IastContext ctx = IastContext.Provider.get();
             for (Object o : collection) {
               if (o instanceof String) {
                 module.taint(ctx, (String) o, SourceTypes.REQUEST_HEADER_VALUE, paramName);

diff --git a/...-appsec/src/main/java/datadog/trace/instrumentation/resteasy/PathParamInjectorAdvice.java b/...-appsec/src/main/java/datadog/trace/instrumentation/resteasy/PathParamInjectorAdvice.java
@@ -18,7 +18,7 @@ public static void onExit(
       if (module != null) {
         if (result instanceof Collection) {
           Collection<?> collection = (Collection<?>) result;
-          final IastContext ctx = IastContext.get();
+          final IastContext ctx = IastContext.Provider.get();
           for (Object o : collection) {
             if (o instanceof String) {
               module.taint(ctx, o, SourceTypes.REQUEST_PATH_PARAMETER, paramName);

diff --git a/...appsec/src/main/java/datadog/trace/instrumentation/resteasy/QueryParamInjectorAdvice.java b/...appsec/src/main/java/datadog/trace/instrumentation/resteasy/QueryParamInjectorAdvice.java
@@ -18,7 +18,7 @@ public static void onExit(
       if (module != null) {
         if (result instanceof Collection) {
           Collection<?> collection = (Collection<?>) result;
-          final IastContext ctx = IastContext.get();
+          final IastContext ctx = IastContext.Provider.get();
           for (Object o : collection) {
             if (o instanceof String) {
               module.taint(ctx, (String) o, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);

diff --git a/...uest-3/src/main/java/datadog/trace/instrumentation/servlet3/MultipartInstrumentation.java b/...uest-3/src/main/java/datadog/trace/instrumentation/servlet3/MultipartInstrumentation.java
@@ -91,7 +91,7 @@ public static void onExit(
       }
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String value : headerValues) {
           module.taint(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, headerName);
         }
@@ -108,7 +108,7 @@ public static void onExit(@Advice.Return final Collection<String> headerNames) {
       }
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String name : headerNames) {
           module.taint(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER);
         }

diff --git a/...ain/java/datadog/trace/instrumentation/servlet3/callsite/HttpServlet3RequestCallSite.java b/...ain/java/datadog/trace/instrumentation/servlet3/callsite/HttpServlet3RequestCallSite.java
@@ -22,7 +22,7 @@ public static java.util.Map<java.lang.String, java.lang.String[]> afterGetParame
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
     if (module != null && map != null && !map.isEmpty()) {
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final Map.Entry<String, String[]> entry : map.entrySet()) {
           final String name = entry.getKey();
           module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);

diff --git a/...rc/main/java/datadog/trace/instrumentation/servlet3/callsite/Servlet3RequestCallSite.java b/...rc/main/java/datadog/trace/instrumentation/servlet3/callsite/Servlet3RequestCallSite.java
@@ -21,7 +21,7 @@ public static java.util.Map<java.lang.String, java.lang.String[]> afterGetParame
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
     if (module != null && map != null && !map.isEmpty()) {
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final Map.Entry<String, String[]> entry : map.entrySet()) {
           final String name = entry.getKey();
           module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);

diff --git a/...c/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestCallSite.java b/...c/main/java/datadog/trace/instrumentation/servlet5/JakartaHttpServletRequestCallSite.java
@@ -61,7 +61,7 @@ public static Enumeration<String> afterGetParameterNames(
         parameterNames.add(paramName);
       }
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String name : parameterNames) {
           module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
         }
@@ -89,7 +89,7 @@ public static String[] afterGetParameterValues(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         try {
-          final IastContext ctx = IastContext.get();
+          final IastContext ctx = IastContext.Provider.get();
           for (final String value : parameterValues) {
             module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
           }
@@ -110,7 +110,7 @@ public static java.util.Map<java.lang.String, java.lang.String[]> afterGetParame
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
     if (module != null && map != null && !map.isEmpty()) {
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final Map.Entry<String, String[]> entry : map.entrySet()) {
           final String name = entry.getKey();
           module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);

diff --git a/...src/main/java/datadog/trace/instrumentation/servlet5/JakartaMultipartInstrumentation.java b/...src/main/java/datadog/trace/instrumentation/servlet5/JakartaMultipartInstrumentation.java
@@ -87,7 +87,7 @@ public static void onExit(
       }
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String value : headerValues) {
           module.taint(ctx, value, SourceTypes.REQUEST_MULTIPART_PARAMETER, headerName);
         }
@@ -104,7 +104,7 @@ public static void onExit(@Advice.Return final Collection<String> headerNames) {
       }
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String name : headerNames) {
           module.taint(ctx, name, SourceTypes.REQUEST_MULTIPART_PARAMETER);
         }

diff --git a/...5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletRequestCallSite.java b/...5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletRequestCallSite.java
@@ -59,7 +59,7 @@ public static Enumeration<String> afterGetParameterNames(
         parameterNames.add(paramName);
       }
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String name : parameterNames) {
           module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
         }
@@ -87,7 +87,7 @@ public static String[] afterGetParameterValues(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         try {
-          final IastContext ctx = IastContext.get();
+          final IastContext ctx = IastContext.Provider.get();
           for (final String value : parameterValues) {
             module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
           }
@@ -107,7 +107,7 @@ public static java.util.Map<java.lang.String, java.lang.String[]> afterGetParame
     final PropagationModule module = InstrumentationBridge.PROPAGATION;
     if (module != null && map != null && !map.isEmpty()) {
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final Map.Entry<String, String[]> entry : map.entrySet()) {
           final String name = entry.getKey();
           module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);

diff --git a/...c/main/java/datadog/trace/instrumentation/servlet/request/HttpServletRequestCallSite.java b/...c/main/java/datadog/trace/instrumentation/servlet/request/HttpServletRequestCallSite.java
@@ -66,7 +66,7 @@ public static Enumeration<String> afterGetHeaders(
         headerValues.add(headerValue);
       }
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String value : headerValues) {
           module.taint(ctx, value, SourceTypes.REQUEST_HEADER_VALUE, headerName);
         }
@@ -102,7 +102,7 @@ public static Enumeration<String> afterGetHeaderNames(
         headerNames.add(headerName);
       }
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String name : headerNames) {
           module.taint(ctx, name, SourceTypes.REQUEST_HEADER_NAME);
         }
@@ -126,7 +126,7 @@ public static Cookie[] afterGetCookies(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         try {
-          final IastContext ctx = IastContext.get();
+          final IastContext ctx = IastContext.Provider.get();
           for (final Cookie cookie : cookies) {
             module.taint(ctx, cookie, SourceTypes.REQUEST_COOKIE_VALUE);
           }
@@ -193,7 +193,7 @@ public static Enumeration<String> afterGetParameterNames(
         parameterNames.add(paramName);
       }
       try {
-        final IastContext ctx = IastContext.get();
+        final IastContext ctx = IastContext.Provider.get();
         for (final String name : parameterNames) {
           module.taint(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME);
         }
@@ -221,7 +221,7 @@ public static String[] afterGetParameterValues(
       final PropagationModule module = InstrumentationBridge.PROPAGATION;
       if (module != null) {
         try {
-          final IastContext ctx = IastContext.get();
+          final IastContext ctx = IastContext.Provider.get();
           for (final String value : parameterValues) {
             module.taint(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, paramName);
           }