JS: Summary/steps for iterators and generators

javascript/ql/lib/semmle/javascript/Generators.qll

@@ -11,7 +11,7 @@ private import semmle.javascript.dataflow.internal.PreCallGraphStep
 private module GeneratorDataFlow {
   private import DataFlow::PseudoProperties
 
-  private class ArrayIteration extends PreCallGraphStep {
+  private class ArrayIteration extends LegacyPreCallGraphStep {
     override predicate storeStep(DataFlow::Node pred, DataFlow::SourceNode succ, string prop) {
       exists(DataFlow::FunctionNode f | f.getFunction().isGenerator() |
         prop = iteratorElement() and

javascript/ql/lib/semmle/javascript/internal/flow_summaries/AllFlowSummaries.qll

@@ -1,6 +1,9 @@
 private import AmbiguousCoreMethods
 private import Arrays2
 private import AsyncAwait
+private import ForOfLoops
+private import Generators
+private import Iterators2
 private import Maps2
 private import Promises2
 private import Sets2

javascript/ql/lib/semmle/javascript/internal/flow_summaries/ForOfLoops.qll

@@ -0,0 +1,53 @@
+/**
+ * Contains flow steps to model flow through `for..of` loops.
+ */
+
+private import javascript
+private import semmle.javascript.dataflow.internal.DataFlowNode
+private import semmle.javascript.dataflow.internal.AdditionalFlowInternal
+private import semmle.javascript.dataflow.internal.DataFlowPrivate
+
+class ForOfLoopStep extends AdditionalFlowInternal {
+  override predicate needsSynthesizedNode(AstNode node, string tag, DataFlowCallable container) {
+    // Intermediate nodes to convert (MapKey, MapValue) to a `[key, value]` array.
+    node instanceof ForOfStmt and
+    tag = ["map-key", "map-value"] and
+    container.asSourceCallable() = node.getContainer()
+  }
+
+  override predicate readStep(
+    DataFlow::Node pred, DataFlow::ContentSet contents, DataFlow::Node succ
+  ) {
+    exists(ForOfStmt stmt | pred = stmt.getIterationDomain().flow() |
+      contents =
+        [
+          DataFlow::ContentSet::arrayElement(), DataFlow::ContentSet::setElement(),
+          DataFlow::ContentSet::iteratorElement()
+        ] and
+      succ = DataFlow::lvalueNode(stmt.getLValue())
+      or
+      contents = DataFlow::ContentSet::mapKey() and
+      succ = getSynthesizedNode(stmt, "map-key")
+      or
+      contents = DataFlow::ContentSet::mapValueAll() and
+      succ = getSynthesizedNode(stmt, "map-value")
+      or
+      contents = DataFlow::ContentSet::iteratorError() and
+      succ = stmt.getIterationDomain().getExceptionTarget()
+    )
+  }
+
+  override predicate storeStep(
+    DataFlow::Node pred, DataFlow::ContentSet contents, DataFlow::Node succ
+  ) {
+    exists(ForOfStmt stmt |
+      pred = getSynthesizedNode(stmt, "map-key") and
+      contents.asArrayIndex() = 0
+      or
+      pred = getSynthesizedNode(stmt, "map-value") and
+      contents.asArrayIndex() = 1
+    |
+      succ = DataFlow::lvalueNode(stmt.getLValue())
+    )
+  }
+}

javascript/ql/lib/semmle/javascript/internal/flow_summaries/Generators.qll

@@ -0,0 +1,59 @@
+/**
+ * Contains flow steps to model flow through generator functions.
+ */
+
+private import javascript
+private import semmle.javascript.dataflow.internal.DataFlowNode
+private import semmle.javascript.dataflow.internal.AdditionalFlowInternal
+
+/**
+ * Steps modelling flow out of a generator function:
+ * ```js
+ * function* foo() {
+ *   yield x;  // store 'x' in the return value's IteratorElement
+ *   yield* y; // flow directly to return value, which has expectsContent, so only iterator contents can pass through.
+ *   throw z;  // store 'z' in the return value's IteratorError
+ * }
+ * ```
+ */
+class GeneratorFunctionStep extends AdditionalFlowInternal {
+  override predicate expectsContent(DataFlow::Node node, DataFlow::ContentSet contents) {
+    // Ensure that the return value can only return iterator contents. This is needed for 'yield*'.
+    exists(Function fun |
+      fun.isGenerator() and
+      node = TFunctionReturnNode(fun) and
+      contents = DataFlow::ContentSet::iteratorFilter()
+    )
+  }
+
+  override predicate storeStep(
+    DataFlow::Node pred, DataFlow::ContentSet contents, DataFlow::Node succ
+  ) {
+    // `yield x`. Store into the return value's iterator element.
+    exists(Function fun, YieldExpr yield | fun.isGenerator() |
+      not yield.isDelegating() and
+      yield.getContainer() = fun and
+      pred = yield.getOperand().flow() and
+      contents = DataFlow::ContentSet::iteratorElement() and
+      succ = TFunctionReturnNode(fun)
+    )
+    or
+    exists(Function f | f.isGenerator() |
+      // Store thrown exceptions in the iterator-error
+      pred = TExceptionalFunctionReturnNode(f) and
+      succ = TFunctionReturnNode(f) and
+      contents = DataFlow::ContentSet::iteratorError()
+    )
+  }
+
+  override predicate step(DataFlow::Node pred, DataFlow::Node succ) {
+    // `yield* x`. Flow into the return value, which has expectsContent, so only iterator contents can pass through.
+    exists(Function fun, YieldExpr yield |
+      fun.isGenerator() and
+      yield.getContainer() = fun and
+      yield.isDelegating() and
+      pred = yield.getOperand().flow() and
+      succ = TFunctionReturnNode(fun)
+    )
+  }
+}

javascript/ql/lib/semmle/javascript/internal/flow_summaries/Iterators2.qll

@@ -0,0 +1,29 @@
+/**
+ * Contains flow summaries and steps modelling flow through iterators.
+ */
+
+private import javascript
+private import semmle.javascript.dataflow.internal.DataFlowNode
+private import semmle.javascript.dataflow.FlowSummary
+private import semmle.javascript.dataflow.internal.AdditionalFlowInternal
+private import FlowSummaryUtil
+
+class IteratorNext extends SummarizedCallable {
+  IteratorNext() { this = "Iterator#next" }
+
+  override DataFlow::MethodCallNode getACallSimple() {
+    result.getMethodName() = "next" and
+    result.getNumArgument() = 0
+  }
+
+  override predicate propagatesFlowExt(string input, string output, boolean preservesValue) {
+    preservesValue = true and
+    (
+      input = "Argument[this].IteratorElement" and
+      output = "ReturnValue.Member[value]"
+      or
+      input = "Argument[this].IteratorError" and
+      output = "ReturnValue[exception]"
+    )
+  }
+}