JS: "return" flow through callbacks

javascript/ql/src/semmle/javascript/dataflow/Configuration.qll

@@ -660,7 +660,7 @@ private predicate flowThroughProperty(
  * All of this is done under configuration `cfg`, and `arg` flows along a path
  * summarized by `summary`, while `cb` is only tracked locally.
  */
-private predicate higherOrderCall(
+private predicate summarizedHigherOrderCall(
   DataFlow::Node arg, DataFlow::Node cb, int i, DataFlow::Configuration cfg, PathSummary summary
 ) {
   exists (Function f, DataFlow::InvokeNode outer, DataFlow::InvokeNode inner, int j,
@@ -676,12 +676,52 @@ private predicate higherOrderCall(
     // indirect higher-order call
     exists (DataFlow::Node cbArg, PathSummary newSummary |
       cbParm.flowsTo(cbArg) and
-      higherOrderCall(innerArg, cbArg, i, cfg, newSummary) and
+      summarizedHigherOrderCall(innerArg, cbArg, i, cfg, newSummary) and
       summary = oldSummary.append(PathSummary::call()).append(newSummary)
     )
   )
 }
 
+/**
+ * Holds if `arg` is passed as the `i`th argument to `callback` through a callback invocation.
+ *
+ * This can be a summarized call, that is, `arg` and `callback` flow into a call,
+ * `f(arg, callback)`, which  performs the invocation.
+ *
+ * Alternatively, the callback can flow into a call `f(callback)` which itself provides the `arg`.
+ * That is, `arg` refers to a value defined in `f` or one of its callees.
+ */
+predicate higherOrderCall(
+  DataFlow::Node arg, DataFlow::SourceNode callback, int i, DataFlow::Configuration cfg,
+  PathSummary summary
+) {
+  // Summarized call
+  exists (DataFlow::Node cb |
+    summarizedHigherOrderCall(arg, cb, i, cfg, summary) and
+    callback.flowsTo(cb)
+  )
+  or
+  // Local invocation of a parameter
+  isRelevant(arg, cfg) and
+  exists (DataFlow::InvokeNode invoke |
+    arg = invoke.getArgument(i) and
+    invoke = callback.(DataFlow::ParameterNode).getACall() and
+    summary = PathSummary::call()
+  )
+  or
+  // Forwarding of the callback parameter  (but not the argument).
+  // We use a return summary since flow moves back towards the call site.
+  // This ensures that an argument that is only tainted in some contexts cannot flow
+  // out to every callback.
+  exists(DataFlow::Node cbArg, DataFlow::SourceNode innerCb, PathSummary oldSummary |
+    higherOrderCall(arg, innerCb, i, cfg, oldSummary) and
+    callStep(cbArg, innerCb) and
+    callback.flowsTo(cbArg) and
+    summary = PathSummary::return().append(oldSummary)
+  )
+}
+
+
 /**
  * Holds if `pred` is passed as an argument to a function `f` which also takes a
  * callback parameter `cb` and then invokes `cb`, passing `pred` into parameter `succ`
@@ -693,12 +733,8 @@ private predicate higherOrderCall(
 private predicate flowIntoHigherOrderCall(
   DataFlow::Node pred, DataFlow::Node succ, DataFlow::Configuration cfg, PathSummary summary
 ) {
-  exists(
-    DataFlow::Node fArg, DataFlow::FunctionNode cb,
-    int i, PathSummary oldSummary
-  |
-    higherOrderCall(pred, fArg, i, cfg, oldSummary) and
-    cb = fArg.getALocalSource() and
+  exists(DataFlow::FunctionNode cb, int i, PathSummary oldSummary |
+    higherOrderCall(pred, cb, i, cfg, oldSummary) and
     succ = cb.getParameter(i) and
     summary = oldSummary.append(PathSummary::call())
   )

javascript/ql/src/semmle/javascript/dataflow/internal/FlowSteps.qll

@@ -264,6 +264,12 @@ predicate callback(DataFlow::Node arg, DataFlow::SourceNode cb) {
     callStep(cbArg, cbParm) and
     cb.flowsTo(cbArg)
   )
+  or
+  exists (DataFlow::ParameterNode cbParm, DataFlow::Node cbArg |
+    callback(arg, cbParm) and
+    callStep(cbArg, cbParm) and
+    cb.flowsTo(cbArg)
+  )
 }
 
 /**

javascript/ql/test/library-tests/TaintTracking/BasicTaintTracking.expected

@@ -1,4 +1,14 @@
 | advanced-callgraph.js:2:13:2:20 | source() | advanced-callgraph.js:6:22:6:22 | v |
+| callbacks.js:4:6:4:13 | source() | callbacks.js:34:27:34:27 | x |
+| callbacks.js:4:6:4:13 | source() | callbacks.js:35:27:35:27 | x |
+| callbacks.js:5:6:5:13 | source() | callbacks.js:34:27:34:27 | x |
+| callbacks.js:5:6:5:13 | source() | callbacks.js:35:27:35:27 | x |
+| callbacks.js:25:16:25:23 | source() | callbacks.js:47:26:47:26 | x |
+| callbacks.js:25:16:25:23 | source() | callbacks.js:48:26:48:26 | x |
+| callbacks.js:37:17:37:24 | source() | callbacks.js:37:37:37:37 | x |
+| callbacks.js:44:17:44:24 | source() | callbacks.js:41:10:41:10 | x |
+| callbacks.js:50:18:50:25 | source() | callbacks.js:30:29:30:29 | y |
+| callbacks.js:51:18:51:25 | source() | callbacks.js:30:29:30:29 | y |
 | constructor-calls.js:4:18:4:25 | source() | constructor-calls.js:18:8:18:14 | c.taint |
 | constructor-calls.js:4:18:4:25 | source() | constructor-calls.js:22:8:22:19 | c_safe.taint |
 | constructor-calls.js:10:16:10:23 | source() | constructor-calls.js:26:8:26:14 | d.taint |

javascript/ql/test/library-tests/TaintTracking/callbacks.js

@@ -0,0 +1,52 @@
+import * as dummy from 'dummy'; // treat as module
+
+function provideTaint(cb) {
+  cb(source());
+  cb(source());
+}
+
+function provideTaint2(cb) {
+  provideTaint(cb);
+  provideTaint(cb); // suppress precision gains from functions with unique call site
+}
+
+function forwardTaint(x, cb) {
+  cb(x);
+  cb(x);
+}
+
+function forwardTaint2(x, cb) {
+  forwardTaint(x, cb);
+  forwardTaint(x, cb);
+}
+
+function middleSource(cb) {
+  // The source occurs in-between the callback definition and the callback invocation.
+  forwardTaint(source(), cb);
+}
+
+function middleCallback(x) {
+  // The callback definition occurs in-between the source and the callback invocation.
+  forwardTaint(x, y => sink(y)); // NOT OK
+}
+
+function test() {
+  provideTaint2(x => sink(x)); // NOT OK
+  provideTaint2(x => sink(x)); // NOT OK
+
+  forwardTaint2(source(), x => sink(x)); // NOT OK
+  forwardTaint2("safe", x => sink(x)); // OK
+  
+  function helper1(x) {
+    sink(x); // NOT OK
+    return x;
+  }
+  forwardTaint2(source(), helper1);
+  sink(helper1("safe")); // OK
+
+  middleSource(x => sink(x)); // NOT OK
+  middleSource(x => sink(x)); // NOT OK
+
+  middleCallback(source());
+  middleCallback(source());
+}