Java: Switch array steps and one containerstep.

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

@@ -1,6 +1,8 @@
 import java
 import semmle.code.java.Collections
 import semmle.code.java.Maps
+private import semmle.code.java.dataflow.SSA
+private import DataFlowUtil
 
 private class EntryType extends RefType {
   EntryType() {
@@ -426,3 +428,44 @@ predicate containerStep(Expr n1, Expr n2) {
   containerReturnValueStep(n1, n2) or
   containerUpdateStep(n1, n2)
 }
+
+predicate arrayStoreStep(Node node1, Node node2) {
+  exists(Argument arg |
+    node1.asExpr() = arg and
+    arg.isVararg() and
+    node2.(ImplicitVarargsArray).getCall() = arg.getCall()
+  )
+  or
+  node2.asExpr().(ArrayInit).getAnInit() = node1.asExpr()
+  or
+  exists(Assignment assign | assign.getSource() = node1.asExpr() |
+    node2.(PostUpdateNode).getPreUpdateNode().asExpr() = assign.getDest().(ArrayAccess).getArray()
+  )
+}
+
+private predicate enhancedForStmtStep(Node node1, Node node2, Type containerType) {
+  exists(EnhancedForStmt for, Expr e, SsaExplicitUpdate v |
+    for.getExpr() = e and
+    node1.asExpr() = e and
+    containerType = e.getType() and
+    v.getDefiningExpr() = for.getVariable() and
+    v.getAFirstUse() = node2.asExpr()
+  )
+}
+
+predicate arrayReadStep(Node node1, Node node2, Type elemType) {
+  exists(ArrayAccess aa |
+    aa.getArray() = node1.asExpr() and
+    aa.getType() = elemType and
+    node2.asExpr() = aa
+  )
+  or
+  exists(Array arr |
+    enhancedForStmtStep(node1, node2, arr) and
+    arr.getComponentType() = elemType
+  )
+}
+
+predicate collectionReadStep(Node node1, Node node2) {
+  enhancedForStmtStep(node1, node2, any(Type t | not t instanceof Array))
+}

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

@@ -4,6 +4,7 @@ private import DataFlowImplCommon
 private import DataFlowDispatch
 private import semmle.code.java.controlflow.Guards
 private import semmle.code.java.dataflow.SSA
+private import ContainerFlow
 private import FlowSummaryImpl as FlowSummaryImpl
 import DataFlowNodes::Private
 
@@ -137,13 +138,15 @@ class MapValueContent extends Content, TMapValueContent {
  * Thus, `node2` references an object with a field `f` that contains the
  * value of `node1`.
  */
-predicate storeStep(Node node1, Content f, PostUpdateNode node2) {
+predicate storeStep(Node node1, Content f, Node node2) {
   exists(FieldAccess fa |
     instanceFieldAssign(node1.asExpr(), fa) and
-    node2.getPreUpdateNode() = getFieldQualifier(fa) and
+    node2.(PostUpdateNode).getPreUpdateNode() = getFieldQualifier(fa) and
     f.(FieldContent).getField() = fa.getField()
   )
   or
+  f instanceof ArrayContent and arrayStoreStep(node1, node2)
+  or
   FlowSummaryImpl::Private::Steps::summaryStoreStep(node1, f, node2)
 }
 
@@ -171,6 +174,10 @@ predicate readStep(Node node1, Content f, Node node2) {
     node2.asExpr() = get
   )
   or
+  f instanceof ArrayContent and arrayReadStep(node1, node2, _)
+  or
+  f instanceof CollectionContent and collectionReadStep(node1, node2)
+  or
   FlowSummaryImpl::Private::Steps::summaryReadStep(node1, f, node2)
 }
 

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

@@ -144,6 +144,8 @@ predicate simpleLocalFlowStep(Node node1, Node node2) {
   or
   node2.asExpr().(AssignExpr).getSource() = node1.asExpr()
   or
+  node2.asExpr().(ArrayCreationExpr).getInit() = node1.asExpr()
+  or
   exists(MethodAccess ma, ValuePreservingMethod m, int argNo |
     ma.getCallee().getSourceDeclaration() = m and m.returnsValue(argNo)
   |

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

@@ -30,6 +30,12 @@ DataFlowType getContentType(Content c) {
   or
   c instanceof ArrayContent and
   result instanceof TypeObject
+  or
+  c instanceof MapKeyContent and
+  result instanceof TypeObject
+  or
+  c instanceof MapValueContent and
+  result instanceof TypeObject
 }
 
 /** Gets the return type of kind `rk` for callable `c`. */

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

@@ -60,12 +60,6 @@ private module Cached {
     localAdditionalTaintUpdateStep(src.asExpr(),
       sink.(DataFlow::PostUpdateNode).getPreUpdateNode().asExpr())
     or
-    exists(Argument arg |
-      src.asExpr() = arg and
-      arg.isVararg() and
-      sink.(DataFlow::ImplicitVarargsArray).getCall() = arg.getCall()
-    )
-    or
     FlowSummaryImpl::Private::Steps::summaryLocalStep(src, sink, false)
   }
 
@@ -103,20 +97,8 @@ private predicate localAdditionalTaintExprStep(Expr src, Expr sink) {
   or
   sink.(AssignAddExpr).getSource() = src and sink.getType() instanceof TypeString
   or
-  sink.(ArrayCreationExpr).getInit() = src
-  or
-  sink.(ArrayInit).getAnInit() = src
-  or
-  sink.(ArrayAccess).getArray() = src
-  or
   sink.(LogicExpr).getAnOperand() = src
   or
-  exists(EnhancedForStmt for, SsaExplicitUpdate v |
-    for.getExpr() = src and
-    v.getDefiningExpr() = for.getVariable() and
-    v.getAFirstUse() = sink
-  )
-  or
   containerReturnValueStep(src, sink)
   or
   constructorStep(src, sink)
@@ -141,10 +123,6 @@ private predicate localAdditionalTaintExprStep(Expr src, Expr sink) {
  * This is restricted to cases where the step updates the value of `sink`.
  */
 private predicate localAdditionalTaintUpdateStep(Expr src, Expr sink) {
-  exists(Assignment assign | assign.getSource() = src |
-    sink = assign.getDest().(ArrayAccess).getArray()
-  )
-  or
   containerUpdateStep(src, sink)
   or
   qualifierToArgumentStep(src, sink)