[Java] Support local taint tracking through JDK Collections lambdas

java/ql/src/semmle/code/java/dataflow/internal/ContainerFlow.qll

@@ -187,6 +187,24 @@ private predicate qualifierToMethodStep(Expr tracked, MethodAccess sink) {
   tracked = sink.getQualifier()
 }
 
+private predicate taintPreservingQualifierToFunctionalParameter(Method m, Parameter p) {
+  m.getDeclaringType() instanceof IterableType and
+  m.hasName("forEach")
+  or
+  m.(MapMethod).hasName("forEach")
+  or
+  m.getDeclaringType() instanceof IteratorType and
+  m.hasName("forEachRemaining")
+}
+
+private predicate qualifierToFunctionalParameterStep(Expr tracked, Parameter parameter) {
+  exists(MethodAccess ma |
+    taintPreservingQualifierToFunctionalParameter(ma.getMethod(), parameter) and
+    tracked = ma.getQualifier() and
+    ma.getAnArgument().(FunctionalExpr).asMethod().getAParameter() = parameter
+  )
+}
+
 private predicate qualifierToArgumentStep(Expr tracked, Expr sink) {
   exists(MethodAccess ma, CollectionMethod method |
     method = ma.getMethod() and
@@ -406,6 +424,14 @@ predicate containerReturnValueStep(Expr n1, Expr n2) {
   qualifierToMethodStep(n1, n2) or argToMethodStep(n1, n2)
 }
 
+/**
+ * Holds if the step from `n1` to `p1` is extracting a value from a
+ * container. This is restricted to cases where `p1` is the value extracted.
+ */
+predicate containerToParameterStep(Expr n1, Parameter p1) {
+  qualifierToFunctionalParameterStep(n1, p1)
+}
+
 /**
  * Holds if the step from `n1` to `n2` is either extracting a value from a
  * container, inserting a value into a container, or transforming one container

java/ql/src/semmle/code/java/dataflow/internal/TaintTrackingUtil.qll

@@ -48,6 +48,8 @@ predicate localTaintStep(DataFlow::Node src, DataFlow::Node sink) {
 predicate localAdditionalTaintStep(DataFlow::Node src, DataFlow::Node sink) {
   localAdditionalTaintExprStep(src.asExpr(), sink.asExpr())
   or
+  localAdditionalTaintExprToParameterStep(src.asExpr(), sink.asParameter())
+  or
   localAdditionalTaintUpdateStep(src.asExpr(),
     sink.(DataFlow::PostUpdateNode).getPreUpdateNode().asExpr())
   or
@@ -120,6 +122,15 @@ private predicate localAdditionalTaintExprStep(Expr src, Expr sink) {
   formatStep(src, sink)
 }
 
+/**
+ * Holds if taint can flow in one local step from `src` to functional argument `sink`
+ * excluding local data flow steps. That is, `src`, and `sink` are likely to represent
+ * different objects.
+ */
+private predicate localAdditionalTaintExprToParameterStep(Expr src, Parameter sink) {
+  containerToParameterStep(src, sink)
+}
+
 /**
  * Holds if taint can flow in one local step from `src` to `sink` excluding
  * local data flow steps. That is, `src` and `sink` are likely to represent

java/ql/src/semmle/code/java/dispatch/ObjFlow.qll

@@ -129,6 +129,8 @@ private predicate step(Node n1, Node n2) {
   or
   containerStep(n1.asExpr(), n2.asExpr())
   or
+  containerToParameterStep(n1.asExpr(), n2.asParameter())
+  or
   exists(Field v |
     containerStep(n1.asExpr(), v.getAnAccess()) and
     n2.asExpr() = v.getAnAccess()

java/ql/test/library-tests/dataflow/collections/Test.java

@@ -25,5 +25,21 @@ public void run() {
     Iterator<String> it = m.values().iterator();
     String x5 = it.next();
     sink(x5); // Flow
+
+    it.forEachRemaining(x6 -> {
+      sink(x6); // Flow
+    });
+
+    m.forEach((key, value) -> {
+      sink(key); // Flow
+      sink(value); // Flow
+    });
+
+    m.entrySet().forEach(entry -> {
+      String x7 = entry.getKey();
+      sink(x7); // No flow
+      String x8 = entry.getValue();
+      sink(x8); // Flow
+    });
   }
 }

java/ql/test/library-tests/dataflow/collections/flow.expected

@@ -132,3 +132,7 @@
 | Test.java:13:18:13:24 | tainted | Test.java:18:10:18:11 | x3 |
 | Test.java:13:18:13:24 | tainted | Test.java:22:12:22:13 | x4 |
 | Test.java:13:18:13:24 | tainted | Test.java:27:10:27:11 | x5 |
+| Test.java:13:18:13:24 | tainted | Test.java:30:12:30:13 | x6 |
+| Test.java:13:18:13:24 | tainted | Test.java:34:12:34:14 | key |
+| Test.java:13:18:13:24 | tainted | Test.java:35:12:35:16 | value |
+| Test.java:13:18:13:24 | tainted | Test.java:42:12:42:13 | x8 |