JS: Parameterise the module (still only one instantiation)

javascript/ql/lib/semmle/javascript/ApiGraphs.qll

@@ -784,17 +784,6 @@ module API {
       } or
       MkSyntheticCallbackArg(DataFlow::InvokeNode nd)
 
-    private predicate needsDefNode(DataFlow::ClassNode cls) {
-      hasSemantics(cls) and
-      (
-        cls = trackDefNode(_)
-        or
-        cls.getAnInstanceReference() = trackDefNode(_)
-        or
-        needsDefNode(cls.getADirectSubClass())
-      )
-    }
-
     class TDef = MkModuleDef or TNonModuleDef;
 
     class TNonModuleDef = MkModuleExport or MkClassInstance or MkDef or MkSyntheticCallbackArg;
@@ -811,8 +800,26 @@ module API {
       hasSemantics(imp)
     }
 
+    private signature module StageInputSig {
+      /** Holds if `node` should be seen as a use-node root, in addition to module imports (which are the usual roots). */
+      predicate isAdditionalUseRoot(Node node);
+
+      /** Holds if `node` should be seen as a def-node root, in addition to module exports (which are the usual roots). */
+      predicate isAdditionalDefRoot(Node node);
+
+      /**
+       * Holds if `node` is considered "in scope" for this stage, meaning that we allow outgoing labelled edges
+       * to be materialised from here, and continue API graph construction from the successors edges.
+       *
+       * Note that the "additional roots" contributed by the stage inputs may be out of scope but can be tracked to a node in scope.
+       * This predicate should thus not be used to block the tracking of use/def nodes, but only block the creation of new labelled edges.
+       */
+      bindingset[node]
+      predicate inScope(DataFlow::Node node);
+    }
+
     cached
-    private module Stage {
+    private module Stage<StageInputSig S> {
       /**
        * Holds if `rhs` is the right-hand side of a definition of a node that should have an
        * incoming edge from `base` labeled `lbl` in the API graph.
@@ -1005,9 +1012,11 @@ module API {
        */
       cached
       predicate rhs(TApiNode nd, DataFlow::Node rhs) {
+        (S::inScope(rhs) or S::isAdditionalDefRoot(nd)) and
         exists(string m | nd = MkModuleExport(m) | exports(m, rhs))
         or
         rhs(_, _, rhs) and
+        S::inScope(rhs) and
         nd = MkDef(rhs)
       }
 
@@ -1058,7 +1067,8 @@ module API {
           base = MkRoot() and
           exists(EntryPoint e |
             lbl = Label::entryPoint(e) and
-            ref = e.getASource()
+            ref = e.getASource() and
+            S::inScope(ref)
           )
           or
           // property reads
@@ -1230,35 +1240,57 @@ module API {
         )
       }
 
+      private predicate needsDefNode(DataFlow::ClassNode cls) {
+        hasSemantics(cls) and
+        (
+          cls = trackDefNode(_)
+          or
+          cls.getAnInstanceReference() = trackDefNode(_)
+          or
+          needsDefNode(cls.getADirectSubClass())
+          or
+          S::isAdditionalDefRoot(MkClassInstance(cls))
+          or
+          S::isAdditionalUseRoot(MkClassInstance(cls)) // These are also tracked as use-nodes
+        )
+      }
+
       /**
        * Holds if `ref` is a use of node `nd`.
        */
       cached
       predicate use(TApiNode nd, DataFlow::Node ref) {
-        exists(string m, Module mod | nd = MkModuleDef(m) and mod = importableModule(m) |
-          ref = DataFlow::moduleVarNode(mod)
-        )
-        or
-        exists(string m, Module mod | nd = MkModuleExport(m) and mod = importableModule(m) |
-          ref = DataFlow::exportsVarNode(mod)
+        (S::inScope(ref) or S::isAdditionalUseRoot(nd)) and
+        (
+          exists(string m, Module mod | nd = MkModuleDef(m) and mod = importableModule(m) |
+            ref = DataFlow::moduleVarNode(mod)
+          )
           or
-          exists(DataFlow::Node base | use(MkModuleDef(m), base) |
-            ref = trackUseNode(base).getAPropertyRead("exports")
+          exists(string m, Module mod | nd = MkModuleExport(m) and mod = importableModule(m) |
+            ref = DataFlow::exportsVarNode(mod)
+            or
+            exists(DataFlow::Node base | use(MkModuleDef(m), base) |
+              ref = trackUseNode(base).getAPropertyRead("exports")
+            )
+          )
+          or
+          exists(string m |
+            nd = MkModuleImport(m) and
+            ref = DataFlow::moduleImport(m)
           )
         )
         or
-        exists(string m |
-          nd = MkModuleImport(m) and
-          ref = DataFlow::moduleImport(m)
-        )
-        or
-        exists(DataFlow::ClassNode cls | nd = MkClassInstance(cls) |
+        exists(DataFlow::ClassNode cls | nd = MkClassInstance(cls) and needsDefNode(cls) |
           ref = cls.getAReceiverNode()
           or
           ref = cls.(DataFlow::ClassNode).getAPrototypeReference()
         )
         or
         use(_, _, ref) and
+        S::inScope(ref) and
+        nd = MkUse(ref)
+        or
+        S::isAdditionalUseRoot(nd) and
         nd = MkUse(ref)
       }
 
@@ -1498,7 +1530,18 @@ module API {
       }
     }
 
-    import Stage
+    private module Stage1Input implements StageInputSig {
+      pragma[inline]
+      predicate isAdditionalUseRoot(Node node) { none() }
+
+      pragma[inline]
+      predicate isAdditionalDefRoot(Node node) { none() }
+
+      bindingset[node]
+      predicate inScope(DataFlow::Node node) { any() }
+    }
+
+    import Stage<Stage1Input>
 
     /**
      * Holds if there is an edge from `pred` to `succ` in the API graph.