JS: generalize inter procedural IIFE type inference

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

@@ -7,69 +7,6 @@
 import javascript
 import AbstractValuesImpl
 
-/**
- * Flow analysis for immediately-invoked function expressions (IIFEs).
- */
-private class IifeReturnFlow extends DataFlow::AnalyzedValueNode {
-  ImmediatelyInvokedFunctionExpr iife;
-
-  IifeReturnFlow() {
-    astNode = (CallExpr)iife.getInvocation()
-  }
-
-  override AbstractValue getALocalValue() {
-    result = getAReturnValue(iife)
-  }
-}
-
-/**
- * Gets a return value for the immediately-invoked function expression `f`.
- */
-private AbstractValue getAReturnValue(ImmediatelyInvokedFunctionExpr f) {
-  // explicit return value
-  result = f.getAReturnedExpr().analyze().getALocalValue()
-  or
-  // implicit return value
-  (
-    // either because execution of the function may terminate normally
-    mayReturnImplicitly(f)
-    or
-    // or because there is a bare `return;` statement
-    exists (ReturnStmt ret | ret = f.getAReturnStmt() | not exists(ret.getExpr()))
-  ) and
-  result = getDefaultReturnValue(f)
-}
-
-
-/**
- * Holds if the execution of function `f` may complete normally without
- * encountering a `return` or `throw` statement.
- *
- * Note that this is an overapproximation, that is, the predicate may hold
- * of functions that cannot actually complete normally, since it does not
- * account for `finally` blocks and does not check reachability.
- */
-private predicate mayReturnImplicitly(Function f) {
-  exists (ConcreteControlFlowNode final |
-    final.getContainer() = f and
-    final.isAFinalNode() and
-    not final instanceof ReturnStmt and
-    not final instanceof ThrowStmt
-  )
-}
-
-/**
- * Gets the default return value for immediately-invoked function expression `f`,
- * that is, the value that `f` returns if its execution terminates without
- * encountering an explicit `return` statement.
- */
-private AbstractValue getDefaultReturnValue(ImmediatelyInvokedFunctionExpr f) {
-  if f.isGenerator() or f.isAsync() then
-    result = TAbstractOtherObject()
-  else
-    result = TAbstractUndefined()
-}
-
 /**
  * Flow analysis for `this` expressions inside functions.
  */
@@ -190,3 +127,75 @@ private class AnalyzedThisInPropertyFunction extends AnalyzedThisExpr {
   }
 }
 
+/**
+ * Holds if the execution of function `f` may complete normally without
+ * encountering a `return` or `throw` statement.
+ *
+ * Note that this is an overapproximation, that is, the predicate may hold
+ * of functions that cannot actually complete normally, since it does not
+ * account for `finally` blocks and does not check reachability.
+ */
+predicate mayReturnImplicitly(Function f) {
+  exists (ConcreteControlFlowNode final |
+    final.getContainer() = f and
+    final.isAFinalNode() and
+    not final instanceof ReturnStmt and
+    not final instanceof ThrowStmt
+  )
+}
+
+/**
+ * A call with inter-procedural type inference for the return value.
+ */
+abstract class CallWithAnalyzedReturnFlow extends DataFlow::CallNode, DataFlow::AnalyzedValueNode {
+
+  /**
+   * Gets a called function.
+   */
+  abstract Function getAFunction();
+  
+  /**
+   * Gets a return value for this call.
+   */
+  AbstractValue getAReturnValue() {
+    exists (Function f | f = getAFunction() |
+      if f.isGenerator() or f.isAsync() then
+        result = TAbstractOtherObject()
+      else (
+        // explicit return value
+        result = f.getAReturnedExpr().analyze().getALocalValue()
+        or
+        // implicit return value
+        (
+          // either because execution of the function may terminate normally
+          mayReturnImplicitly(f)
+          or
+          // or because there is a bare `return;` statement
+          exists (ReturnStmt ret | ret = f.getAReturnStmt() | not exists(ret.getExpr()))
+        ) and
+        result = TAbstractUndefined()
+      )
+    )
+  }
+
+  override AbstractValue getALocalValue() {
+    result = getAReturnValue()
+  }
+}
+
+/**
+ * Flow analysis for the return value of IIFEs.
+ */
+private class IIFEWithAnalyzedReturnFlow extends CallWithAnalyzedReturnFlow {
+  
+  ImmediatelyInvokedFunctionExpr iife;
+  
+  IIFEWithAnalyzedReturnFlow() {
+    this.asExpr() = iife.getInvocation()
+  }
+  
+  override Function getAFunction() {
+    result = iife
+  }
+  
+}

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

@@ -112,47 +112,43 @@ class AnalyzedVarDef extends VarDef {
 }
 
 /**
- * Flow analysis for simple IIFE parameters.
+ * Flow analysis for simple parameters of selected functions.
  */
-private class AnalyzedIIFEParameter extends AnalyzedVarDef, @vardecl {
-  AnalyzedIIFEParameter() {
-    exists (ImmediatelyInvokedFunctionExpr iife, int parmIdx |
-      this = iife.getParameter(parmIdx) |
-      // we cannot track flow into rest parameters...
-      not this.(Parameter).isRestParameter() and
-      // ...nor flow out of spread arguments
-      exists (int argIdx | argIdx = parmIdx + iife.getArgumentOffset() |
-        not iife.isSpreadArgument([0..argIdx])
-      )
+private class AnalyzedParameter extends AnalyzedVarDef, @vardecl {
+  AnalyzedParameter() {
+    exists (FunctionWithAnalyzedParameters f, int parmIdx |
+      this = f.getParameter(parmIdx) |
+      // we cannot track flow into rest parameters
+      not this.(Parameter).isRestParameter()
     )
   }
 
-  /** Gets the IIFE this is a parameter of. */
-  ImmediatelyInvokedFunctionExpr getIIFE() {
+  /** Gets the function this is a parameter of. */
+  FunctionWithAnalyzedParameters getFunction() {
     this = result.getAParameter()
   }
 
   override DataFlow::AnalyzedNode getRhs() {
-    getIIFE().argumentPassing(this, result.asExpr()) or
+    getFunction().argumentPassing(this, result.asExpr()) or
     result = this.(Parameter).getDefault().analyze()
   }
 
   override AbstractValue getAnRhsValue() {
-    result = AnalyzedVarDef.super.getAnRhsValue() or
-    not getIIFE().argumentPassing(this, _) and result = TAbstractUndefined()
+    result = AnalyzedVarDef.super.getAnRhsValue()
+    or
+    not getFunction().mayReceiveArgument(this) and
+    result = TAbstractUndefined()
   }
 
   override predicate isIncomplete(DataFlow::Incompleteness cause) {
-    exists (ImmediatelyInvokedFunctionExpr iife | iife = getIIFE() |
-      // if the IIFE has a name and that name is referenced, we conservatively
-      // assume that there may be other calls than the direct one
-      exists (iife.getVariable().getAnAccess()) and cause = "call" or
-      // if the IIFE is non-strict and its `arguments` object is accessed, we
-      // also assume that there may be other calls (through `arguments.callee`)
-      not iife.isStrict() and
-      exists (iife.getArgumentsVariable().getAnAccess()) and cause = "call"
+    getFunction().isIncomplete(cause) or
+    (
+      not getFunction().argumentPassing(this, _) and
+      getFunction().mayReceiveArgument(this) and
+      cause = "call"
     )
   }
+  
 }
 
 /**
@@ -650,3 +646,82 @@ private class NamespaceExportVarFlow extends DataFlow::AnalyzedValueNode {
 
   override AbstractValue getALocalValue() { result = TIndefiniteAbstractValue("namespace") }
 }
+
+/**
+ * A function with inter-procedural type inference for its parameters.
+ */
+abstract class FunctionWithAnalyzedParameters extends Function {
+
+  /**
+   * Holds if `p` is a parameter of this function and `arg` is
+   * the corresponding argument.
+   */
+  abstract predicate argumentPassing(SimpleParameter p, Expr arg);
+
+  /**
+   * Holds if `p` is a parameter of this function that may receive a value from an argument.
+   */
+  abstract predicate mayReceiveArgument(Parameter p);
+
+  /**
+   * Holds if flow analysis results for the parameters may be incomplete
+   * due to the given `cause`.
+   */
+  abstract predicate isIncomplete(DataFlow::Incompleteness cause);
+
+}
+
+private abstract class CallWithAnalyzedParameters extends FunctionWithAnalyzedParameters {
+
+  /**
+   * Gets an invocation of this function.
+   */
+  abstract DataFlow::InvokeNode getAnInvocation();
+
+  override predicate argumentPassing(SimpleParameter p, Expr arg) {
+    exists (DataFlow::InvokeNode invk, int argIdx |
+      invk = getAnInvocation() |
+      p = getParameter(argIdx) and not p.isRestParameter() and
+      arg = invk.getArgument(argIdx).asExpr()
+    )
+  }
+
+  override predicate mayReceiveArgument(Parameter p) {
+    exists (DataFlow::InvokeNode invk, int argIdx |
+      invk = getAnInvocation() and
+      p = getParameter(argIdx) |
+      exists (invk.getArgument(argIdx))
+      or
+      invk.asExpr().(InvokeExpr).isSpreadArgument([0..argIdx])
+    )
+  }
+
+}
+
+/**
+ * Flow analysis for simple parameters of IIFEs.
+ */
+private class IIFEWithAnalyzedParameters extends CallWithAnalyzedParameters {
+
+  ImmediatelyInvokedFunctionExpr iife;
+
+  IIFEWithAnalyzedParameters() {
+    this = iife and
+    iife.getInvocationKind() = "direct"
+  }
+
+  override DataFlow::InvokeNode getAnInvocation() {
+    result = iife.getInvocation().flow()
+  }
+
+  override predicate isIncomplete(DataFlow::Incompleteness cause) {
+    // if the IIFE has a name and that name is referenced, we conservatively
+    // assume that there may be other calls than the direct one
+    exists (iife.getVariable().getAnAccess()) and cause = "call" or
+    // if the IIFE is non-strict and its `arguments` object is accessed, we
+    // also assume that there may be other calls (through `arguments.callee`)
+    not iife.isStrict() and
+    exists (iife.getArgumentsVariable().getAnAccess()) and cause = "call"
+  }
+
+}