C++: output iterator flow through operator= models

cpp/ql/src/semmle/code/cpp/models/implementations/Iterator.qll

@@ -78,6 +78,20 @@ private FunctionInput getIteratorArgumentInput(Operator op, int index) {
   )
 }
 
+private FunctionOutput getIteratorArgumentOutput(Operator op, int index) {
+  exists(Type t |
+    t =
+      op
+          .getACallToThisFunction()
+          .getArgument(index)
+          .getExplicitlyConverted()
+          .getType()
+          .stripTopLevelSpecifiers()
+  |
+    result.isParameterDeref(index) // TODO: does this work with an rvalue reference?
+  )
+}
+
 /**
  * A non-member prefix `operator*` function for an iterator type.
  */
@@ -92,6 +106,9 @@ class IteratorPointerDereferenceOperator extends Operator, TaintFunction, Iterat
   override predicate hasTaintFlow(FunctionInput input, FunctionOutput output) {
     input = iteratorInput and
     output.isReturnValue()
+    or
+    input.isReturnValueDeref() and
+    output = getIteratorArgumentOutput(this, 0)
   }
 }
 
@@ -180,6 +197,9 @@ class IteratorPointerDereferenceMemberOperator extends MemberFunction, TaintFunc
   override predicate hasTaintFlow(FunctionInput input, FunctionOutput output) {
     input.isQualifierObject() and
     output.isReturnValue()
+    or
+    input.isReturnValueDeref() and
+    output.isQualifierObject()
   }
 }
 
@@ -274,6 +294,27 @@ class IteratorArrayMemberOperator extends MemberFunction, TaintFunction, Iterato
   }
 }
 
+/**
+ * An `operator=` member function of an iterator class that is not a copy or move assignment
+ * operator.
+ * 
+ * The `hasTaintFlow` override provides flow through output iterators that return themselves with
+ * `operator*` and use their own `operator=` to assign to the container.
+ */
+class IteratorAssignmentMemberOperator extends MemberFunction, TaintFunction {
+  IteratorAssignmentMemberOperator() {
+    this.hasName("operator=") and
+    this.getDeclaringType() instanceof Iterator and
+    not this instanceof CopyAssignmentOperator and
+    not this instanceof MoveAssignmentOperator
+  }
+
+  override predicate hasTaintFlow(FunctionInput input, FunctionOutput output) {
+    input.isParameter(0) and
+    output.isQualifierObject()
+  }
+}
+
 /**
  * A `begin` or `end` member function, or a related member function, that
  * returns an iterator.

cpp/ql/test/library-tests/dataflow/taint-tests/localTaint.expected

@@ -7135,8 +7135,13 @@
 | vector.cpp:427:13:427:30 | call to back_inserter | vector.cpp:428:4:428:5 | it |  |
 | vector.cpp:427:32:427:34 | ref arg out | vector.cpp:429:8:429:10 | out |  |
 | vector.cpp:427:32:427:34 | ref arg out | vector.cpp:430:2:430:2 | out |  |
+| vector.cpp:428:3:428:3 | ref arg call to operator* | vector.cpp:428:6:428:6 | ref arg call to operator++ | TAINT |
+| vector.cpp:428:3:428:3 | ref arg call to operator* | vector.cpp:429:8:429:10 | out |  |
+| vector.cpp:428:3:428:3 | ref arg call to operator* | vector.cpp:430:2:430:2 | out |  |
 | vector.cpp:428:4:428:5 | it | vector.cpp:428:6:428:6 | call to operator++ |  |
 | vector.cpp:428:6:428:6 | call to operator++ | vector.cpp:428:3:428:3 | call to operator* | TAINT |
+| vector.cpp:428:6:428:6 | ref arg call to operator++ | vector.cpp:428:4:428:5 | ref arg it |  |
+| vector.cpp:428:11:428:36 | call to basic_string | vector.cpp:428:3:428:3 | ref arg call to operator* | TAINT |
 | vector.cpp:428:23:428:35 | source_string | vector.cpp:428:11:428:36 | call to basic_string | TAINT |
 | vector.cpp:429:8:429:10 | ref arg out | vector.cpp:430:2:430:2 | out |  |
 | vector.cpp:433:20:433:22 | call to vector | vector.cpp:434:32:434:34 | out |  |
@@ -7145,6 +7150,11 @@
 | vector.cpp:434:13:434:30 | call to back_inserter | vector.cpp:435:4:435:5 | it |  |
 | vector.cpp:434:32:434:34 | ref arg out | vector.cpp:436:8:436:10 | out |  |
 | vector.cpp:434:32:434:34 | ref arg out | vector.cpp:437:2:437:2 | out |  |
+| vector.cpp:435:3:435:3 | ref arg call to operator* | vector.cpp:435:6:435:6 | ref arg call to operator++ | TAINT |
+| vector.cpp:435:3:435:3 | ref arg call to operator* | vector.cpp:436:8:436:10 | out |  |
+| vector.cpp:435:3:435:3 | ref arg call to operator* | vector.cpp:437:2:437:2 | out |  |
 | vector.cpp:435:4:435:5 | it | vector.cpp:435:6:435:6 | call to operator++ |  |
 | vector.cpp:435:6:435:6 | call to operator++ | vector.cpp:435:3:435:3 | call to operator* | TAINT |
+| vector.cpp:435:6:435:6 | ref arg call to operator++ | vector.cpp:435:4:435:5 | ref arg it |  |
+| vector.cpp:435:11:435:16 | call to source | vector.cpp:435:3:435:3 | ref arg call to operator* | TAINT |
 | vector.cpp:436:8:436:10 | ref arg out | vector.cpp:437:2:437:2 | out |  |

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

@@ -657,3 +657,5 @@
 | vector.cpp:409:7:409:9 | v13 | vector.cpp:408:11:408:16 | call to source |
 | vector.cpp:414:7:414:9 | v14 | vector.cpp:413:11:413:16 | call to source |
 | vector.cpp:422:8:422:10 | out | vector.cpp:417:33:417:45 | source_string |
+| vector.cpp:429:8:429:10 | out | vector.cpp:417:33:417:45 | source_string |
+| vector.cpp:436:8:436:10 | out | vector.cpp:435:11:435:16 | call to source |

cpp/ql/test/library-tests/dataflow/taint-tests/test_diff.expected

@@ -383,3 +383,5 @@
 | vector.cpp:409:7:409:9 | vector.cpp:408:11:408:16 | AST only |
 | vector.cpp:414:7:414:9 | vector.cpp:413:11:413:16 | AST only |
 | vector.cpp:422:8:422:10 | vector.cpp:417:33:417:45 | AST only |
+| vector.cpp:429:8:429:10 | vector.cpp:417:33:417:45 | AST only |
+| vector.cpp:436:8:436:10 | vector.cpp:435:11:435:16 | AST only |