JS: AsyncPackage

asgerf · asgerf · commit a2f7ecec6693 · 2021-03-17T13:26:01.000Z
diff --git a/javascript/ql/src/semmle/javascript/frameworks/AsyncPackage.qll b/javascript/ql/src/semmle/javascript/frameworks/AsyncPackage.qll
@@ -150,11 +150,11 @@ module AsyncPackage {
    *
    * For example: `data -> item` in `async.each(data, (item, cb) => {})`.
    */
-  private class IterationInputTaintStep extends TaintTracking::AdditionalTaintStep, IterationCall {
+  private class IterationInputTaintStep extends TaintTracking::SharedTaintStep {
     override predicate step(DataFlow::Node pred, DataFlow::Node succ) {
-      exists(DataFlow::FunctionNode iteratee |
-        iteratee = getIteratorCallback() and // Require a closure to avoid spurious call/return mismatch.
-        pred = getCollection() and
+      exists(DataFlow::FunctionNode iteratee, IterationCall call |
+        iteratee = call.getIteratorCallback() and // Require a closure to avoid spurious call/return mismatch.
+        pred = call.getCollection() and
         succ = iteratee.getParameter(0)
       )
     }
@@ -166,20 +166,19 @@ module AsyncPackage {
    *
    * For example: `item + taint()` -> result` in `async.map(data, (item, cb) => cb(null, item + taint()), (err, result) => {})`.
    */
-  private class IterationOutputTaintStep extends TaintTracking::AdditionalTaintStep, IterationCall {
-    IterationOutputTaintStep() {
-      name = "concat" or
-      name = "map" or
-      name = "reduce" or
-      name = "reduceRight"
-    }
-
+  private class IterationOutputTaintStep extends TaintTracking::SharedTaintStep {
     override predicate step(DataFlow::Node pred, DataFlow::Node succ) {
-      exists(DataFlow::FunctionNode iteratee, DataFlow::FunctionNode final, int i |
-        iteratee = getIteratorCallback().getALocalSource() and
-        final = getFinalCallback() and // Require a closure to avoid spurious call/return mismatch.
+      exists(DataFlow::FunctionNode iteratee, DataFlow::FunctionNode final, int i, IterationCall call |
+        iteratee = call.getIteratorCallback().getALocalSource() and
+        final = call.getFinalCallback() and // Require a closure to avoid spurious call/return mismatch.
         pred = getLastParameter(iteratee).getACall().getArgument(i) and
-        succ = final.getParameter(i)
+        succ = final.getParameter(i) and
+        exists (string name | name = call.getName() |
+          name = "concat" or
+          name = "map" or
+          name = "reduce" or
+          name = "reduceRight"
+        )
       )
     }
   }
@@ -189,17 +188,13 @@ module AsyncPackage {
    *
    * For example: `data -> result` in `async.sortBy(data, orderingFn, (err, result) => {})`.
    */
-  private class IterationPreserveTaintStep extends TaintTracking::AdditionalTaintStep, IterationCall {
-    IterationPreserveTaintStep() {
-      name = "sortBy"
-      // We don't currently include `filter` and `reject` as they could act as sanitizers.
-    }
-
+  private class IterationPreserveTaintStep extends TaintTracking::SharedTaintStep {
     override predicate step(DataFlow::Node pred, DataFlow::Node succ) {
-      exists(DataFlow::FunctionNode final |
-        final = getFinalCallback() and // Require a closure to avoid spurious call/return mismatch.
-        pred = getCollection() and
-        succ = final.getParameter(1)
+      exists(DataFlow::FunctionNode final, IterationCall call |
+        final = call.getFinalCallback() and // Require a closure to avoid spurious call/return mismatch.
+        pred = call.getCollection() and
+        succ = final.getParameter(1) and
+        call.getName() = "sortBy"
       )
     }
   }
