Java: Add additional taint steps through collections.

java/ql/src/semmle/code/java/dataflow/TaintTracking.qll

@@ -12,6 +12,7 @@ private import DefUse
 private import semmle.code.java.security.SecurityTests
 private import semmle.code.java.security.Validation
 private import semmle.code.java.frameworks.android.Intent
+private import semmle.code.java.Maps
 
 module TaintTracking {
   /**
@@ -209,6 +210,12 @@ module TaintTracking {
       sink = assign.getDest().(ArrayAccess).getArray()
     )
     or
+    exists(EnhancedForStmt for, SsaExplicitUpdate v |
+      for.getExpr() = src and
+      v.getDefiningExpr() = for.getVariable() and
+      v.getAFirstUse() = sink
+    )
+    or
     constructorStep(src, sink)
     or
     qualifierToMethodStep(src, sink)
@@ -418,8 +425,49 @@ module TaintTracking {
     or
     m instanceof IntentGetExtraMethod
     or
-    m instanceof CollectionMethod and
-    m.hasName("toArray")
+    m
+        .getDeclaringType()
+        .getSourceDeclaration()
+        .getASourceSupertype*()
+        .hasQualifiedName("java.util", "Map<>$Entry") and
+    m.hasName("getValue")
+    or
+    m
+        .getDeclaringType()
+        .getSourceDeclaration()
+        .getASourceSupertype*()
+        .hasQualifiedName("java.lang", "Iterable") and
+    m.hasName("iterator")
+    or
+    m
+        .getDeclaringType()
+        .getSourceDeclaration()
+        .getASourceSupertype*()
+        .hasQualifiedName("java.util", "Iterator") and
+    m.hasName("next")
+    or
+    m.getDeclaringType().getSourceDeclaration().hasQualifiedName("java.util", "Enumeration") and
+    m.hasName("nextElement")
+    or
+    m.(MapMethod).hasName("entrySet")
+    or
+    m.(MapMethod).hasName("get")
+    or
+    m.(MapMethod).hasName("remove")
+    or
+    m.(MapMethod).hasName("values")
+    or
+    m.(CollectionMethod).hasName("toArray")
+    or
+    m.(CollectionMethod).hasName("get")
+    or
+    m.(CollectionMethod).hasName("remove") and m.getParameterType(0).(PrimitiveType).hasName("int")
+    or
+    m.(CollectionMethod).hasName("subList")
+    or
+    m.(CollectionMethod).hasName("firstElement")
+    or
+    m.(CollectionMethod).hasName("lastElement")
     or
     m.getDeclaringType().hasQualifiedName("java.nio", "ByteBuffer") and
     m.hasName("get")
@@ -596,6 +644,18 @@ module TaintTracking {
     method.getDeclaringType().hasQualifiedName("java.io", "ByteArrayOutputStream") and
     method.hasName("write") and
     arg = 0
+    or
+    method.(MapMethod).hasName("put") and arg = 1
+    or
+    method.(MapMethod).hasName("putAll") and arg = 0
+    or
+    method.(CollectionMethod).hasName("add") and arg = method.getNumberOfParameters() - 1
+    or
+    method.(CollectionMethod).hasName("addAll") and arg = method.getNumberOfParameters() - 1
+    or
+    method.(CollectionMethod).hasName("addElement") and arg = 0
+    or
+    method.(CollectionMethod).hasName("set") and arg = 1
   }
 
   /** A comparison or equality test with a constant. */

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

@@ -0,0 +1,29 @@
+import java.util.*;
+
+public class Test {
+  static String tainted;
+
+  void sink(Object o) { }
+
+  public void run() {
+    HashMap<String, String> m = new HashMap<>();
+    String x1 = m.get("key");
+    sink(x1); // No flow
+
+    m.put("key", tainted);
+    String x2 = m.get("key");
+    sink(x2); // Flow
+
+    String x3 = m.values().toArray(new String[1])[0];
+    sink(x3); // Flow
+
+    for(Map.Entry<String, String> e : m.entrySet()) {
+      String x4 = e.getValue();
+      sink(x4); // Flow
+    }
+
+    Iterator<String> it = m.values().iterator();
+    String x5 = it.next();
+    sink(x5); // Flow
+  }
+}

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

@@ -0,0 +1,4 @@
+| Test.java:13:18:13:24 | tainted | Test.java:15:10:15:11 | x2 |
+| 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 |

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

@@ -0,0 +1,21 @@
+import java
+import semmle.code.java.dataflow.TaintTracking
+
+class Conf extends TaintTracking::Configuration {
+  Conf() { this = "conf" }
+
+  override predicate isSource(DataFlow::Node src) {
+    src.asExpr().(VarAccess).getVariable().hasName("tainted")
+  }
+
+  override predicate isSink(DataFlow::Node sink) {
+    exists(MethodAccess ma |
+      sink.asExpr() = ma.getAnArgument() and
+      ma.getMethod().hasName("sink")
+    )
+  }
+}
+
+from Conf c, DataFlow::Node src, DataFlow::Node sink
+where c.hasFlow(src, sink)
+select src, sink