JS: Precise data-flow for returns from async functions

javascript/ql/src/semmle/javascript/Promises.qll

@@ -455,6 +455,49 @@ private class PromiseTaintStep extends TaintTracking::AdditionalTaintStep {
   }
 }
 
+/**
+ * Defines flow steps for return on async functions.
+ */
+private module AsyncReturnSteps {
+  private predicate valueProp = Promises::valueProp/0;
+
+  private predicate errorProp = Promises::errorProp/0;
+
+  private import semmle.javascript.dataflow.internal.FlowSteps
+
+  /**
+   * A data-flow step for ordinary and exceptional returns from async functions.
+   */
+  private class AsyncReturn extends PreCallGraphStep {
+    override predicate storeStep(DataFlow::Node pred, DataFlow::SourceNode succ, string prop) {
+      exists(DataFlow::FunctionNode f | f.getFunction().isAsync() |
+        // ordinary return
+        prop = valueProp() and
+        pred = f.getAReturn() and
+        succ = f.getReturnNode()
+        or
+        // exceptional return
+        prop = errorProp() and
+        localExceptionStepWithAsyncFlag(pred, succ, true)
+      )
+    }
+  }
+
+  /**
+   * A data-flow step for ordinary return from an async function in a taint configuration.
+   */
+  private class AsyncTaintReturn extends TaintTracking::AdditionalTaintStep, DataFlow::FunctionNode {
+    Function f;
+
+    AsyncTaintReturn() { this.getFunction() = f and f.isAsync() }
+
+    override predicate step(DataFlow::Node pred, DataFlow::Node succ) {
+      returnExpr(f, pred, _) and
+      succ.(DataFlow::FunctionReturnNode).getFunction() = f
+    }
+  }
+}
+
 /**
  * Provides classes for working with the `bluebird` library (http://bluebirdjs.com).
  */

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

@@ -982,8 +982,8 @@ private predicate appendStep(
 private predicate flowThroughCall(
   DataFlow::Node input, DataFlow::Node output, DataFlow::Configuration cfg, PathSummary summary
 ) {
-  exists(Function f, DataFlow::ValueNode ret |
-    ret.asExpr() = f.getAReturnedExpr() and
+  exists(Function f, DataFlow::FunctionReturnNode ret |
+    ret.getFunction() = f and
     (calls(output, f) or callsBound(output, f, _)) and // Do not consider partial calls
     reachableFromInput(f, output, input, ret, cfg, summary) and
     not isBarrierEdge(cfg, ret, output) and
@@ -1000,6 +1000,17 @@ private predicate flowThroughCall(
     not isLabeledBarrierEdge(cfg, ret, output, summary.getEndLabel()) and
     not cfg.isLabeledBarrier(output, summary.getEndLabel())
   )
+  or
+  // exception thrown inside an immediately awaited function call.
+  exists(DataFlow::FunctionNode f, DataFlow::Node invk, DataFlow::Node ret |
+    f.getFunction().isAsync()
+  |
+    (calls(invk, f.getFunction()) or callsBound(invk, f.getFunction(), _)) and
+    reachableFromInput(f.getFunction(), invk, input, ret, cfg, summary) and
+    output = invk.asExpr().getExceptionTarget() and
+    f.getExceptionalReturn() = getThrowTarget(ret) and
+    invk = getAwaitOperand(_)
+  )
 }
 
 /**
@@ -1019,23 +1030,39 @@ private predicate storeStep(
   isAdditionalStoreStep(pred, succ, prop, cfg) and
   summary = PathSummary::level()
   or
-  exists(Function f, DataFlow::Node mid |
+  exists(Function f, DataFlow::Node mid, DataFlow::Node invk |
+    not f.isAsync() and invk = succ
+    or
+    // store in an immediately awaited function call
+    f.isAsync() and
+    invk = getAwaitOperand(succ)
+  |
     // `f` stores its parameter `pred` in property `prop` of a value that flows back to the caller,
     // and `succ` is an invocation of `f`
-    reachableFromInput(f, succ, pred, mid, cfg, summary) and
+    reachableFromInput(f, invk, pred, mid, cfg, summary) and
     (
       returnedPropWrite(f, _, prop, mid)
       or
       exists(DataFlow::SourceNode base | base.flowsToExpr(f.getAReturnedExpr()) |
         isAdditionalStoreStep(mid, base, prop, cfg)
       )
       or
-      succ instanceof DataFlow::NewNode and
+      invk instanceof DataFlow::NewNode and
       receiverPropWrite(f, prop, mid)
     )
   )
 }
 
+/**
+ * Gets a dataflow-node for the operand of the await-expression `await`.
+ */
+private DataFlow::Node getAwaitOperand(DataFlow::Node await) {
+  exists(AwaitExpr awaitExpr |
+    result = awaitExpr.getOperand().getUnderlyingValue().flow() and
+    await.asExpr() = awaitExpr
+  )
+}
+
 /**
  * Holds if `f` may `read` property `prop` of parameter `parm`.
  */
@@ -1128,8 +1155,14 @@ private predicate loadStep(
   isAdditionalLoadStep(pred, succ, prop, cfg) and
   summary = PathSummary::level()
   or
-  exists(Function f, DataFlow::Node read |
-    parameterPropRead(f, succ, pred, prop, read, cfg) and
+  exists(Function f, DataFlow::Node read, DataFlow::Node invk |
+    not f.isAsync() and invk = succ
+    or
+    // load from an immediately awaited function call
+    f.isAsync() and
+    invk = getAwaitOperand(succ)
+  |
+    parameterPropRead(f, invk, pred, prop, read, cfg) and
     reachesReturn(f, read, cfg, summary)
   )
 }
@@ -1624,6 +1657,9 @@ class MidPathNode extends PathNode, MkMidNode {
     // Skip the exceptional return on functions, as this highlights the entire function.
     nd = any(DataFlow::FunctionNode f).getExceptionalReturn()
     or
+    // Skip the special return node for functions, as this highlights the entire function (and the returned expr is the previous node).
+    nd = any(DataFlow::FunctionNode f).getReturnNode()
+    or
     // Skip the synthetic 'this' node, as a ThisExpr will be the next node anyway
     nd = DataFlow::thisNode(_)
     or

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

@@ -922,6 +922,32 @@ module DataFlow {
     override File getFile() { result = function.getFile() }
   }
 
+  /**
+   * A data flow node representing the values returned by a function.
+   */
+  class FunctionReturnNode extends DataFlow::Node, TFunctionReturnNode {
+    Function function;
+
+    FunctionReturnNode() { this = TFunctionReturnNode(function) }
+
+    override string toString() { result = "return of " + function.describe() }
+
+    override predicate hasLocationInfo(
+      string filepath, int startline, int startcolumn, int endline, int endcolumn
+    ) {
+      function.getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn)
+    }
+
+    override BasicBlock getBasicBlock() { result = function.(ExprOrStmt).getBasicBlock() }
+
+    /**
+     * Gets the function corresponding to this return node.
+     */
+    Function getFunction() { result = function }
+
+    override File getFile() { result = function.getFile() }
+  }
+
   /**
    * A data flow node representing the exceptions thrown by the callee of an invocation.
    */
@@ -1265,6 +1291,13 @@ module DataFlow {
     nd = TExceptionalFunctionReturnNode(function)
   }
 
+  /**
+   * INTERNAL: Use `FunctionReturnNode` instead.
+   */
+  predicate functionReturnNode(DataFlow::Node nd, Function function) {
+    nd = TFunctionReturnNode(function)
+  }
+
   /**
    * Gets the data flow node corresponding the given l-value expression, if
    * such a node exists.
@@ -1460,6 +1493,11 @@ module DataFlow {
         localCall(succExpr, f)
       )
     )
+    or
+    // from returned expr to the FunctionReturnNode.
+    exists(Function f | not f.isAsync() |
+      DataFlow::functionReturnNode(succ, f) and pred = valueNode(f.getAReturnedExpr())
+    )
   }
 
   /**

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

@@ -491,6 +491,16 @@ class FunctionNode extends DataFlow::ValueNode, DataFlow::SourceNode {
   DataFlow::ExceptionalFunctionReturnNode getExceptionalReturn() {
     DataFlow::exceptionalFunctionReturnNode(result, astNode)
   }
+
+  /**
+   * Gets the data flow node representing the value returned from this function.
+   *
+   * Note that this differs from `getAReturn()`, in that every function has exactly
+   * one canonical return node, but may have multiple (or zero) returned expressions.
+   * The result of `getAReturn()` is always a predecessor of `getReturnNode()`
+   * in the data-flow graph.
+   */
+  DataFlow::FunctionReturnNode getReturnNode() { DataFlow::functionReturnNode(result, astNode) }
 }
 
 /**

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

@@ -319,6 +319,9 @@ module SourceNode {
       this = DataFlow::destructuredModuleImportNode(_)
       or
       this = DataFlow::globalAccessPathRootPseudoNode()
+      or
+      // Include return nodes because they model the implicit Promise creation in async functions.
+      DataFlow::functionReturnNode(this, _)
     }
   }
 }

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

@@ -243,7 +243,7 @@ private module NodeTracking {
     basicStoreStep(pred, succ, prop) and
     summary = PathSummary::level()
     or
-    exists(Function f, DataFlow::Node mid |
+    exists(Function f, DataFlow::Node mid | not f.isAsync() |
       // `f` stores its parameter `pred` in property `prop` of a value that flows back to the caller,
       // and `succ` is an invocation of `f`
       reachableFromInput(f, succ, pred, mid, summary) and
@@ -266,7 +266,7 @@ private module NodeTracking {
     basicLoadStep(pred, succ, prop) and
     summary = PathSummary::level()
     or
-    exists(Function f, DataFlow::SourceNode parm |
+    exists(Function f, DataFlow::SourceNode parm | not f.isAsync() |
       argumentPassing(succ, pred, f, parm) and
       reachesReturn(f, parm.getAPropertyRead(prop), summary)
     )

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

@@ -27,6 +27,7 @@ newtype TNode =
     exists(decl.getASpecifier().getImportedName())
   } or
   THtmlAttributeNode(HTML::Attribute attr) or
+  TFunctionReturnNode(Function f) or
   TExceptionalFunctionReturnNode(Function f) or
   TExceptionalInvocationReturnNode(InvokeExpr e) or
   TGlobalAccessPathRoot()

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

@@ -61,13 +61,40 @@ predicate localFlowStep(
  * Holds if an exception thrown from `pred` can propagate locally to `succ`.
  */
 predicate localExceptionStep(DataFlow::Node pred, DataFlow::Node succ) {
+  localExceptionStepWithAsyncFlag(pred, succ, false)
+}
+
+/**
+ * Holds if an exception thrown from `pred` can propagate locally to `succ`.
+ *
+ * The `async` flag is true if the step involves wrapping the exception in a rejected Promise.
+ */
+predicate localExceptionStepWithAsyncFlag(DataFlow::Node pred, DataFlow::Node succ, boolean async) {
+  exists(DataFlow::Node target | target = getThrowTarget(pred) |
+    async = false and
+    succ = target and
+    not succ = any(DataFlow::FunctionNode f | f.getFunction().isAsync()).getExceptionalReturn()
+    or
+    async = true and
+    exists(DataFlow::FunctionNode f | f.getExceptionalReturn() = target |
+      succ = f.getReturnNode() // returns a rejected promise - therefore using the ordinary return node.
+    )
+  )
+}
+
+/**
+ * Gets the dataflow-node that an exception thrown at `thrower` will flow to.
+ *
+ * The predicate that all functions are not async.
+ */
+DataFlow::Node getThrowTarget(DataFlow::Node thrower) {
   exists(Expr expr |
     expr = any(ThrowStmt throw).getExpr() and
-    pred = expr.flow()
+    thrower = expr.flow()
     or
-    DataFlow::exceptionalInvocationReturnNode(pred, expr)
+    DataFlow::exceptionalInvocationReturnNode(thrower, expr)
   |
-    succ = expr.getExceptionTarget()
+    result = expr.getExceptionTarget()
   )
 }
 
@@ -213,14 +240,14 @@ private module CachedSteps {
 
   /**
    * Holds if there is a flow step from `pred` to `succ` through:
-   * - returning a value from a function call, or
+   * - returning a value from a function call (from the special `FunctionReturnNode`), or
    * - throwing an exception out of a function call, or
    * - the receiver flowing out of a constructor call.
    */
   cached
   predicate returnStep(DataFlow::Node pred, DataFlow::Node succ) {
     exists(Function f | calls(succ, f) or callsBound(succ, f, _) |
-      returnExpr(f, pred, _)
+      DataFlow::functionReturnNode(pred, f)
       or
       succ instanceof DataFlow::NewNode and
       DataFlow::thisNode(pred, f)

javascript/ql/test/library-tests/CallGraphs/FullTest/tests.expected

@@ -67,6 +67,7 @@ test_getAFunctionValue
 | es2015.js:2:14:4:3 | () {\\n   ... ");\\n  } | es2015.js:2:14:4:3 | () {\\n   ... ");\\n  } |
 | es2015.js:6:5:6:16 | ExampleClass | es2015.js:2:14:4:3 | () {\\n   ... ");\\n  } |
 | es2015.js:8:2:12:1 | functio ... \\n  };\\n} | es2015.js:8:2:12:1 | functio ... \\n  };\\n} |
+| es2015.js:8:2:12:1 | return of anonymous function | es2015.js:9:10:11:3 | () => { ... ();\\n  } |
 | es2015.js:9:10:11:3 | () => { ... ();\\n  } | es2015.js:9:10:11:3 | () => { ... ();\\n  } |
 | es2015.js:10:5:10:20 | arguments.callee | es2015.js:8:2:12:1 | functio ... \\n  };\\n} |
 | es2015.js:10:5:10:22 | arguments.callee() | es2015.js:9:10:11:3 | () => { ... ();\\n  } |

javascript/ql/test/library-tests/DataFlow/flowStep.expected

@@ -17,10 +17,12 @@
 | sources.js:1:6:1:6 | x | sources.js:1:11:1:11 | x |
 | sources.js:1:6:1:11 | x => x | sources.js:1:5:1:12 | (x => x) |
 | sources.js:1:11:1:11 | x | sources.js:1:1:1:12 | new (x => x) |
+| sources.js:1:11:1:11 | x | sources.js:1:6:1:11 | return of anonymous function |
 | sources.js:3:2:5:1 | functio ... x+19;\\n} | sources.js:3:1:5:2 | (functi ... +19;\\n}) |
 | sources.js:3:11:3:11 | x | sources.js:3:11:3:11 | x |
 | sources.js:3:11:3:11 | x | sources.js:4:10:4:10 | x |
 | sources.js:4:10:4:13 | x+19 | sources.js:3:1:5:6 | (functi ... \\n})(23) |
+| sources.js:4:10:4:13 | x+19 | sources.js:3:2:5:1 | return of anonymous function |
 | sources.js:5:4:5:5 | 23 | sources.js:3:11:3:11 | x |
 | sources.js:9:14:9:18 | array | sources.js:9:14:9:18 | array |
 | sources.js:9:14:9:18 | array | sources.js:10:19:10:23 | array |
@@ -89,7 +91,9 @@
 | tst.js:16:13:16:13 | a | tst.js:16:13:16:13 | a |
 | tst.js:16:13:16:13 | a | tst.js:18:12:18:12 | a |
 | tst.js:18:12:18:12 | a | tst.js:16:1:20:9 | (functi ... ("arg") |
+| tst.js:18:12:18:12 | a | tst.js:16:2:20:1 | return of function f |
 | tst.js:19:10:19:11 | "" | tst.js:16:1:20:9 | (functi ... ("arg") |
+| tst.js:19:10:19:11 | "" | tst.js:16:2:20:1 | return of function f |
 | tst.js:20:4:20:8 | "arg" | tst.js:16:13:16:13 | a |
 | tst.js:22:5:22:25 | readFileSync | tst.js:23:1:23:12 | readFileSync |
 | tst.js:22:7:22:18 | readFileSync | tst.js:22:5:22:25 | readFileSync |
@@ -102,11 +106,13 @@
 | tst.js:28:2:28:1 | x | tst.js:29:3:29:3 | x |
 | tst.js:28:2:29:3 | () =>\\n  x | tst.js:28:1:30:1 | (() =>\\n ... ables\\n) |
 | tst.js:29:3:29:3 | x | tst.js:28:1:30:3 | (() =>\\n ... les\\n)() |
+| tst.js:29:3:29:3 | x | tst.js:28:2:29:3 | return of anonymous function |
 | tst.js:32:1:32:0 | x | tst.js:33:10:33:10 | x |
 | tst.js:32:1:34:1 | functio ... ables\\n} | tst.js:32:10:32:10 | g |
 | tst.js:32:10:32:10 | g | tst.js:35:1:35:1 | g |
 | tst.js:32:10:32:10 | g | tst.js:60:1:60:1 | g |
 | tst.js:32:10:32:10 | g | tst.js:62:4:62:4 | g |
+| tst.js:33:10:33:10 | x | tst.js:32:1:34:1 | return of function g |
 | tst.js:37:5:42:1 | o | tst.js:43:1:43:1 | o |
 | tst.js:37:5:42:1 | o | tst.js:44:1:44:1 | o |
 | tst.js:37:5:42:1 | o | tst.js:61:3:61:3 | o |
@@ -142,6 +148,7 @@
 | tst.js:90:15:90:15 | o | tst.js:90:4:90:11 | { r: z } |
 | tst.js:90:15:90:15 | o | tst.js:90:4:90:15 | { r: z } = o |
 | tst.js:91:10:91:18 | x + y + z | tst.js:87:1:96:2 | (functi ... r: 0\\n}) |
+| tst.js:91:10:91:18 | x + y + z | tst.js:87:2:92:1 | return of anonymous function |
 | tst.js:92:4:96:1 | {\\n  p:  ...  r: 0\\n} | tst.js:87:11:87:24 | { p: x, ...o } |
 | tst.js:98:2:103:1 | functio ...  + z;\\n} | tst.js:98:1:103:2 | (functi ... + z;\\n}) |
 | tst.js:98:11:98:24 | rest | tst.js:99:15:99:18 | rest |
@@ -156,6 +163,7 @@
 | tst.js:101:13:101:16 | rest | tst.js:101:3:101:9 | [ , z ] |
 | tst.js:101:13:101:16 | rest | tst.js:101:3:101:16 | [ , z ] = rest |
 | tst.js:102:10:102:18 | x + y + z | tst.js:98:1:103:17 | (functi ... 3, 0 ]) |
+| tst.js:102:10:102:18 | x + y + z | tst.js:98:2:103:1 | return of anonymous function |
 | tst.js:103:4:103:16 | [ 19, 23, 0 ] | tst.js:98:11:98:24 | [ x, ...rest ] |
 | tst.js:105:1:105:1 | x | tst.js:105:1:105:6 | x ?? y |
 | tst.js:105:6:105:6 | y | tst.js:105:1:105:6 | x ?? y |