Java: Improve taint step modeling to use postupdate nodes.

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

@@ -161,10 +161,27 @@ private predicate argToQualifierStep(Expr tracked, Expr sink) {
 /**
  * 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
- * to another.
+ * to another. This is restricted to cases where `n2` is the returned value of
+ * a call.
  */
-predicate containerStep(Expr n1, Expr n2) {
-  qualifierToMethodStep(n1, n2) or
+predicate containerReturnValueStep(Expr n1, Expr n2) { qualifierToMethodStep(n1, n2) }
+
+/**
+ * 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
+ * to another. This is restricted to cases where the value of `n2` is being modified.
+ */
+predicate containerUpdateStep(Expr n1, Expr n2) {
   qualifierToArgumentStep(n1, n2) or
   argToQualifierStep(n1, n2)
 }
+
+/**
+ * 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
+ * to another.
+ */
+predicate containerStep(Expr n1, Expr n2) {
+  containerReturnValueStep(n1, n2) or
+  containerUpdateStep(n1, n2)
+}

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

@@ -26,6 +26,8 @@ private newtype TNode =
     e instanceof Argument and not e.getType() instanceof ImmutableType
     or
     exists(FieldAccess fa | fa.getField() instanceof InstanceField and e = fa.getQualifier())
+    or
+    exists(ArrayAccess aa | e = aa.getArray())
   } or
   TImplicitExprPostUpdate(InstanceAccessExt ia) {
     implicitInstanceArgument(_, ia)

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

@@ -41,6 +41,9 @@ predicate localTaintStep(DataFlow::Node src, DataFlow::Node sink) {
 predicate localAdditionalTaintStep(DataFlow::Node src, DataFlow::Node sink) {
   localAdditionalTaintExprStep(src.asExpr(), sink.asExpr())
   or
+  localAdditionalTaintUpdateStep(src.asExpr(),
+    sink.(DataFlow::PostUpdateNode).getPreUpdateNode().asExpr())
+  or
   exists(Argument arg |
     src.asExpr() = arg and
     arg.isVararg() and
@@ -105,30 +108,20 @@ private predicate localAdditionalTaintExprStep(Expr src, Expr sink) {
   or
   sink.(LogicExpr).getAnOperand() = src
   or
-  exists(Assignment assign | assign.getSource() = src |
-    sink = assign.getDest().(ArrayAccess).getArray()
-  )
-  or
   exists(EnhancedForStmt for, SsaExplicitUpdate v |
     for.getExpr() = src and
     v.getDefiningExpr() = for.getVariable() and
     v.getAFirstUse() = sink
   )
   or
-  containerStep(src, sink)
+  containerReturnValueStep(src, sink)
   or
   constructorStep(src, sink)
   or
   qualifierToMethodStep(src, sink)
   or
-  qualifierToArgumentStep(src, sink)
-  or
   argToMethodStep(src, sink)
   or
-  argToArgStep(src, sink)
-  or
-  argToQualifierStep(src, sink)
-  or
   comparisonStep(src, sink)
   or
   stringBuilderStep(src, sink)
@@ -136,6 +129,26 @@ private predicate localAdditionalTaintExprStep(Expr src, Expr sink) {
   serializationStep(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
+ * different objects.
+ * 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)
+  or
+  argToArgStep(src, sink)
+  or
+  argToQualifierStep(src, sink)
+}
+
 private class BulkData extends RefType {
   BulkData() {
     this.(Array).getElementType().(PrimitiveType).getName().regexpMatch("byte|char")
@@ -242,7 +255,7 @@ private predicate constructorStep(Expr tracked, ConstructorCall sink) {
 }
 
 /** Access to a method that passes taint from qualifier to argument. */
-private predicate qualifierToArgumentStep(Expr tracked, RValue sink) {
+private predicate qualifierToArgumentStep(Expr tracked, Expr sink) {
   exists(MethodAccess ma, int arg |
     taintPreservingQualifierToArgument(ma.getMethod(), arg) and
     tracked = ma.getQualifier() and
@@ -458,7 +471,7 @@ private predicate taintPreservingArgumentToMethod(Method method, int arg) {
  * Holds if `tracked` and `sink` are arguments to a method that transfers taint
  * between arguments.
  */
-private predicate argToArgStep(Expr tracked, RValue sink) {
+private predicate argToArgStep(Expr tracked, Expr sink) {
   exists(MethodAccess ma, Method method, int input, int output |
     taintPreservingArgToArg(method, input, output) and
     ma.getMethod() = method and

java/ql/test/library-tests/dataflow/taint/A.java

@@ -24,4 +24,52 @@ public class A {
     bInput.read(b2);
     sink(b2);
   }
+
+  void streamWrite(ByteArrayOutputStream baos, byte[] data) {
+    baos.write(data);
+  }
+
+  void test3(ByteArrayOutputStream baos) {
+    streamWrite(baos, taint());
+    sink(baos.toByteArray());
+  }
+
+  static class BaosHolder {
+    ByteArrayOutputStream baos = new ByteArrayOutputStream();
+  }
+
+  void streamWriteHolder(BaosHolder bh, byte[] data) {
+    bh.baos.write(data);
+  }
+
+  void test4(BaosHolder bh) {
+    streamWriteHolder(bh, taint());
+    sink(bh.baos.toByteArray());
+  }
+
+  static class DataHolder {
+    byte[] data = new byte[10];
+  }
+
+  void test5_a(DataHolder dh) {
+    ByteArrayInputStream bais = new ByteArrayInputStream(taint());
+    bais.read(dh.data);
+    test5_b(dh);
+  }
+
+  void test5_b(DataHolder dh) {
+    sink(dh.data);
+  }
+
+  void arrayWrite(byte[] from, byte[] to) {
+    for (int i = 0; i < 10; i++) {
+      to[i] = from[i];
+    }
+  }
+
+  void test6() {
+    byte[] b = new byte[10];
+    arrayWrite(taint(), b);
+    sink(b);
+  }
 }

java/ql/test/library-tests/dataflow/taint/test.expected

@@ -1,5 +1,9 @@
 | A.java:10:19:10:25 | taint(...) | A.java:15:10:15:11 | b2 |
 | A.java:20:19:20:25 | taint(...) | A.java:25:10:25:11 | b2 |
+| A.java:33:23:33:29 | taint(...) | A.java:34:10:34:27 | toByteArray(...) |
+| A.java:46:27:46:33 | taint(...) | A.java:47:10:47:30 | toByteArray(...) |
+| A.java:55:58:55:64 | taint(...) | A.java:61:10:61:16 | dh.data |
+| A.java:72:16:72:22 | taint(...) | A.java:73:10:73:10 | b |
 | B.java:15:21:15:27 | taint(...) | B.java:18:10:18:16 | aaaargs |
 | B.java:15:21:15:27 | taint(...) | B.java:21:10:21:10 | s |
 | B.java:15:21:15:27 | taint(...) | B.java:24:10:24:15 | concat |