Ensure binding as map works

DataDog · Sep 13, 2024 · ecde2b8 · ecde2b8
1 parent a311e59
commit ecde2b8
Show file tree

Hide file tree

Showing 7 changed files with 63 additions and 16 deletions.
diff --git a/...y2JettyTest/groovy/datadog/trace/instrumentation/jersey2/iast/IastJersey2JettyTest.groovy b/...y2JettyTest/groovy/datadog/trace/instrumentation/jersey2/iast/IastJersey2JettyTest.groovy
@@ -150,7 +150,7 @@ class IastJersey2JettyTest extends IastRequestTestRunner {
 
   void 'all form params'() {
     when:
-    String url = buildUrl 'iast/all_form'
+    String url = buildUrl "iast/$variant"
     def body = new FormBody.Builder()
       .add('var1', 'foo')
       .add('var1', 'bar')
@@ -187,6 +187,9 @@ class IastJersey2JettyTest extends IastRequestTestRunner {
       value 'a b c'
       range 0, 5, source(SourceTypes.REQUEST_PARAMETER_VALUE, 'var2', 'a b c')
     }
+
+    where:
+    variant << ['all_form', 'all_form_map']
   }
 
   void 'cookie'() {

diff --git a/...rc/jersey2JettyTest/groovy/datadog/trace/instrumentation/jersey2/iast/IastResource.groovy b/...rc/jersey2JettyTest/groovy/datadog/trace/instrumentation/jersey2/iast/IastResource.groovy
@@ -98,6 +98,14 @@ class IastResource {
     return "IAST: ${pairList}"
   }
 
+  @Path("/all_form_map")
+  @POST
+  @Produces(MediaType.TEXT_PLAIN)
+  String form(final MultivaluedMap<String, String> form) {
+    def pairList = collectMultiMap(form)
+    return "IAST: ${pairList}"
+  }
+
   private static collectMultiMap(final MultivaluedMap<String, String> map) {
     return map.keySet().sort().collect {key ->
       final values = map[key]

diff --git a/...y3JettyTest/groovy/datadog/trace/instrumentation/jersey3/iast/IastJersey3JettyTest.groovy b/...y3JettyTest/groovy/datadog/trace/instrumentation/jersey3/iast/IastJersey3JettyTest.groovy
@@ -150,7 +150,7 @@ class IastJersey3JettyTest extends IastRequestTestRunner {
 
   void 'all form params'() {
     when:
-    String url = buildUrl 'iast/all_form'
+    String url = buildUrl "iast/$variant"
     def body = new FormBody.Builder()
       .add('var1', 'foo')
       .add('var1', 'bar')
@@ -187,6 +187,9 @@ class IastJersey3JettyTest extends IastRequestTestRunner {
       value 'a b c'
       range 0, 5, source(SourceTypes.REQUEST_PARAMETER_VALUE, 'var2', 'a b c')
     }
+
+    where:
+    variant << ['all_form', 'all_form_map']
   }
 
   void 'cookie'() {

diff --git a/...rc/jersey3JettyTest/groovy/datadog/trace/instrumentation/jersey3/iast/IastResource.groovy b/...rc/jersey3JettyTest/groovy/datadog/trace/instrumentation/jersey3/iast/IastResource.groovy
@@ -98,6 +98,14 @@ class IastResource {
     return "IAST: ${pairList}"
   }
 
+  @Path("/all_form_map")
+  @POST
+  @Produces(MediaType.TEXT_PLAIN)
+  String form(final MultivaluedMap<String, String> form) {
+    def pairList = collectMultiMap(form)
+    return "IAST: ${pairList}"
+  }
+
   private static collectMultiMap(final MultivaluedMap<String, String> map) {
     return map.keySet().sort().collect {key ->
       final values = map[key]

diff --git a/.../main/java/datadog/trace/instrumentation/jersey/InboundMessageContextInstrumentation.java b/.../main/java/datadog/trace/instrumentation/jersey/InboundMessageContextInstrumentation.java
@@ -2,6 +2,7 @@
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
+import static datadog.trace.instrumentation.jersey.JerseyTaintHelper.taintMap;
 import static datadog.trace.instrumentation.jersey.JerseyTaintHelper.taintMultiValuedMap;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
@@ -140,6 +141,9 @@ public static void onExit(
         return;
       }
       module.taintObject(ctx, result, SourceTypes.REQUEST_BODY);
+      if (result instanceof Map) {
+        taintMap(ctx, module, SourceTypes.REQUEST_PARAMETER_VALUE, (Map<?, ?>) result);
+      }
     }
   }
 }
diff --git a/...entation/jersey/src/main/java/datadog/trace/instrumentation/jersey/JerseyTaintHelper.java b/...entation/jersey/src/main/java/datadog/trace/instrumentation/jersey/JerseyTaintHelper.java
@@ -16,12 +16,43 @@ public static void taintMultiValuedMap(
       final byte type,
       final Map<String, List<String>> target) {
     final byte nameType = SourceTypes.namedSource(type);
+    final boolean reportName = nameType != type;
     for (Map.Entry<String, List<String>> entry : target.entrySet()) {
       final String name = entry.getKey();
-      module.taintString(ctx, name, nameType, name);
+      if (reportName) {
+        module.taintString(ctx, name, nameType, name);
+      }
       for (String value : entry.getValue()) {
         module.taintString(ctx, value, type, name);
       }
     }
   }
+
+  public static void taintMap(
+      final IastContext ctx,
+      final PropagationModule module,
+      final byte type,
+      final Map<?, ?> target) {
+    final byte nameType = SourceTypes.namedSource(type);
+    final boolean reportName = nameType != type;
+    for (final Map.Entry<?, ?> entry : target.entrySet()) {
+      final Object key = entry.getKey();
+      if (key instanceof String) {
+        final String name = (String) key;
+        if (reportName) {
+          module.taintString(ctx, name, nameType, name);
+        }
+        final Object value = entry.getValue();
+        if (value instanceof String) {
+          module.taintString(ctx, (String) value, type, name);
+        } else if (value instanceof List) {
+          for (final Object item : (List<?>) value) {
+            if (item instanceof String) {
+              module.taintString(ctx, (String) item, type, name);
+            }
+          }
+        }
+      }
+    }
+  }
 }
diff --git a/.../src/main/java/datadog/trace/instrumentation/jersey/UriRoutingContextInstrumentation.java b/.../src/main/java/datadog/trace/instrumentation/jersey/UriRoutingContextInstrumentation.java
@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.jersey;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.instrumentation.jersey.JerseyTaintHelper.taintMultiValuedMap;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
@@ -70,12 +71,7 @@ public static void onExit(
         return;
       }
       prop.taintObject(ctx, pathParams, SourceTypes.REQUEST_PATH_PARAMETER);
-      for (Map.Entry<String, List<String>> entry : pathParams.entrySet()) {
-        final String name = entry.getKey();
-        for (String value : entry.getValue()) {
-          prop.taintString(ctx, value, SourceTypes.REQUEST_PATH_PARAMETER, name);
-        }
-      }
+      taintMultiValuedMap(ctx, prop, SourceTypes.REQUEST_PATH_PARAMETER, pathParams);
     }
   }
 
@@ -98,13 +94,7 @@ public static void onExit(
         return;
       }
       prop.taintObject(ctx, queryParams, SourceTypes.REQUEST_PARAMETER_VALUE);
-      for (Map.Entry<String, List<String>> entry : queryParams.entrySet()) {
-        final String name = entry.getKey();
-        prop.taintString(ctx, name, SourceTypes.REQUEST_PARAMETER_NAME, name);
-        for (String value : entry.getValue()) {
-          prop.taintString(ctx, value, SourceTypes.REQUEST_PARAMETER_VALUE, name);
-        }
-      }
+      taintMultiValuedMap(ctx, prop, SourceTypes.REQUEST_PARAMETER_VALUE, queryParams);
     }
   }
 }