JS: Augment call graph using type-tracked class instances

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

@@ -642,16 +642,23 @@ class ClassNode extends DataFlow::SourceNode {
    */
   FunctionNode getAStaticMethod() { result = impl.getAStaticMethod() }
 
+  /**
+   * Gets a dataflow node that refers to the superclass of this class.
+   */
+  DataFlow::Node getASuperClassNode() { result = impl.getASuperClassNode() }
+
   /**
    * Gets a direct super class of this class.
    */
   ClassNode getADirectSuperClass() {
-    result.getConstructor().getAstNode() = impl
-          .getASuperClassNode()
-          .analyze()
-          .getAValue()
-          .(AbstractCallable)
-          .getFunction()
+    result.getAClassReference().flowsTo(getASuperClassNode())
+  }
+
+  /**
+   * Gets a direct subclass of this class.
+   */
+  final ClassNode getADirectSubClass() {
+    this = result.getADirectSuperClass()
   }
 
   /**
@@ -662,6 +669,63 @@ class ClassNode extends DataFlow::SourceNode {
     or
     result = getAnInstanceMember().getReceiver()
   }
+
+  /**
+   * Gets the abstract value representing the class itself.
+   */
+  AbstractValue getAbstractClassValue() {
+    result = this.(AnalyzedNode).getAValue()
+  }
+
+  /**
+   * Gets the abstract value representing an instance of this class.
+   */
+  AbstractValue getAbstractInstanceValue() {
+    result = AbstractInstance::of(getAstNode())
+  }
+
+  /**
+   * Gets a dataflow node that refers to this class object.
+   */
+  private DataFlow::SourceNode getAClassReference(DataFlow::TypeTracker t) {
+    t.start() and
+    result.(AnalyzedNode).getAValue() = getAbstractClassValue()
+    or
+    exists(DataFlow::TypeTracker t2 |
+      result = getAClassReference(t2).track(t2, t)
+    )
+  }
+
+  /**
+   * Gets a dataflow node that refers to this class object.
+   */
+  DataFlow::SourceNode getAClassReference() {
+    result = getAClassReference(DataFlow::TypeTracker::end())
+  }
+
+  /**
+   * Gets a dataflow node that refers to an instance of this class.
+   */
+  private DataFlow::SourceNode getAnInstanceReference(DataFlow::TypeTracker t) {
+    result = getAClassReference(t.continue()).getAnInstantiation()
+    or
+    t.start() and
+    result.(AnalyzedNode).getAValue() = getAbstractInstanceValue()
+    or
+    t.start() and
+    result = getAReceiverNode()
+    or
+    exists(DataFlow::TypeTracker t2 |
+      result = getAnInstanceReference(t2).track(t2, t)
+    )
+  }
+
+  /**
+   * Gets a dataflow node that refers to an instance of this class.
+   */
+  DataFlow::SourceNode getAnInstanceReference() {
+    result = getAnInstanceReference(DataFlow::TypeTracker::end())
+  }
 }
 
 module ClassNode {
@@ -803,6 +867,9 @@ module ClassNode {
       kind = MemberKind::method() and
       result = getAPrototypeReference().getAPropertySource(name)
       or
+      kind = MemberKind::method() and
+      result = getConstructor().getReceiver().getAPropertyWrite(name).getRhs().getALocalSource()
+      or
       exists(PropertyAccessor accessor |
         accessor = getAnAccessor(kind) and
         accessor.getName() = name and
@@ -814,6 +881,9 @@ module ClassNode {
       kind = MemberKind::method() and
       result = getAPrototypeReference().getAPropertyWrite().getRhs().getALocalSource()
       or
+      kind = MemberKind::method() and
+      result = getConstructor().getReceiver().getAPropertyWrite().getRhs().getALocalSource()
+      or
       exists(PropertyAccessor accessor |
         accessor = getAnAccessor(kind) and
         result = accessor.getInit().flow()

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

@@ -97,11 +97,33 @@ private module CachedSteps {
     f = cap.getContainer()
   }
 
+  /**
+   * Holds if the method invoked by `invoke` resolved to a member named `name` in `cls`
+   * or one of its super classes.
+   */
+  cached
+  predicate callResolvesToClass(DataFlow::InvokeNode invoke, DataFlow::ClassNode cls, string name) {
+    invoke = cls.getAnInstanceReference().getAMethodCall(name)
+    or
+    exists(DataFlow::ClassNode subclass |
+      callResolvesToClass(invoke, subclass, name) and
+      not exists(subclass.getAnInstanceMember(name)) and
+      cls = subclass.getADirectSuperClass()
+    )
+  }
+
   /**
    * Holds if `invk` may invoke `f`.
    */
   cached
-  predicate calls(DataFlow::InvokeNode invk, Function f) { f = invk.getACallee(0) }
+  predicate calls(DataFlow::InvokeNode invk, Function f) {
+    f = invk.getACallee(0)
+    or
+    exists(DataFlow::ClassNode cls, string name |
+      callResolvesToClass(invk, cls, name) and
+      f = cls.getInstanceMethod(name).getFunction()
+    )
+  }
 
   /**
    * Holds if `invk` may invoke `f` indirectly through the given `callback` argument.