PS: Resolve member function calls using the shared type-tracking library.

Author: MathiasVP

powershell/ql/lib/semmle/code/powershell/Function.qll

@@ -4,6 +4,8 @@ import semmle.code.powershell.controlflow.BasicBlocks
 abstract private class AbstractFunction extends Ast {
   abstract string getName();
 
+  final predicate hasName(string name) { this.getName() = name }
+
   abstract ScriptBlock getBody();
 
   abstract Parameter getFunctionParameter(int i);

powershell/ql/lib/semmle/code/powershell/InvokeMemberExpression.qll

@@ -17,3 +17,9 @@ class InvokeMemberExpr extends @invoke_member_expression, MemberExprBase {
 
   override predicate isStatic() { this.getQualifier() instanceof TypeNameExpr }
 }
+
+class ConstructorCall extends InvokeMemberExpr {
+  ConstructorCall() { this.isStatic() and this.getName() = "new" }
+
+  Type getConstructedType() { result = this.getQualifier().(TypeNameExpr).getType() }
+}

powershell/ql/lib/semmle/code/powershell/Type.qll

@@ -10,4 +10,11 @@ class Type extends @type_definition, Stmt {
   Member getMember(int i) { type_definition_members(this, i, result) }
 
   Member getAMember() { result = this.getMember(_) }
+
+  MemberFunction getMemberFunction(string name) {
+    result = this.getAMember() and
+    result.hasName(name)
+  }
+
+  MemberFunction getAMemberFunction() { result = this.getMemberFunction(_) }
 }

powershell/ql/lib/semmle/code/powershell/TypeExpression.qll

@@ -8,4 +8,7 @@ class TypeNameExpr extends @type_expression, Expr {
   override string toString() { result = this.getName() }
 
   override SourceLocation getLocation() { type_expression_location(this, result) }
+
+  /** Gets the type referred to by this `TypeNameExpr`. */
+  Type getType() { result.getName() = this.getName() }
 }

powershell/ql/lib/semmle/code/powershell/controlflow/CfgNodes.qll

@@ -214,7 +214,7 @@ module ExprNodes {
 
     override InvokeMemberChildMapping e;
 
-    final override InvokeMemberExpr getExpr() { result = super.getExpr() }
+    override InvokeMemberExpr getExpr() { result = super.getExpr() }
 
     final override ExprCfgNode getQualifier() { e.hasCfgChild(e.getQualifier(), this, result) }
 
@@ -231,6 +231,15 @@ module ExprNodes {
     final override ExprCfgNode getCommand() { none() }
   }
 
+    /** A control-flow node that wraps an `ConstructorCall` expression. */
+  class ConstructorCallCfgNode extends InvokeMemberCfgNode {
+    ConstructorCallCfgNode() { super.getExpr() instanceof ConstructorCall }
+
+    final override ConstructorCall getExpr() { result = super.getExpr() }
+
+    Type getConstructedType() { result = this.getExpr().getConstructedType() }
+  }
+
   /** A control-flow node that wraps a qualifier expression. */
   class QualifierCfgNode extends ExprCfgNode {
     QualifierCfgNode() { this = any(InvokeMemberCfgNode invoke).getQualifier() }

powershell/ql/lib/semmle/code/powershell/dataflow/internal/DataFlowDispatch.qll

@@ -1,6 +1,8 @@
 private import powershell
 private import semmle.code.powershell.Cfg
 private import DataFlowPrivate
+private import DataFlowPublic
+private import semmle.code.powershell.typetracking.internal.TypeTrackingImpl
 private import codeql.util.Boolean
 private import codeql.util.Unit
 
@@ -123,11 +125,47 @@ class NormalCall extends DataFlowCall, TNormalCall {
 /** A call for which we want to compute call targets. */
 private class RelevantCall extends CfgNodes::CallCfgNode { }
 
-/** Holds if `call` may resolve to the returned source-code method. */
-private DataFlowCallable viableSourceCallable(DataFlowCall call) {
-  none() // TODO
-  or
-  result = any(AdditionalCallTarget t).viableTarget(call.asCall())
+private module TrackInstanceInput implements CallGraphConstruction::InputSig {
+  newtype State = additional MkState(Type m, Boolean exact)
+
+  predicate start(Node start, State state) {
+    exists(Type tp, boolean exact | state = MkState(tp, exact) |
+      start.asExpr().(CfgNodes::ExprNodes::ConstructorCallCfgNode).getConstructedType() = tp
+    )
+  }
+
+  pragma[nomagic]
+  predicate stepNoCall(Node nodeFrom, Node nodeTo, StepSummary summary) {
+    smallStepNoCall(nodeFrom, nodeTo, summary)
+  }
+
+  predicate stepCall(Node nodeFrom, Node nodeTo, StepSummary summary) {
+    smallStepCall(nodeFrom, nodeTo, summary)
+  }
+
+  class StateProj = Unit;
+
+  Unit stateProj(State state) { exists(state) and exists(result) }
+
+  predicate filter(Node n, Unit u) { none() }
+}
+
+private predicate qualifiedCall(CfgNodes::CallCfgNode call, Node receiver, string method) {
+  call.getQualifier() = receiver.asExpr() and
+  call.getName() = method
+}
+
+private Node trackInstance(Type t, boolean exact) {
+  result =
+    CallGraphConstruction::Make<TrackInstanceInput>::track(TrackInstanceInput::MkState(t, exact))
+}
+
+private CfgScope getTargetInstance(CfgNodes::CallCfgNode call) {
+  exists(Node receiver, string method, Type t |
+    qualifiedCall(call, receiver, method) and
+    receiver = trackInstance(t, _) and
+    result = t.getMemberFunction(method).getBody()
+  )
 }
 
 /**
@@ -154,7 +192,14 @@ private module Cached {
 
   /** Gets a viable run-time target for the call `call`. */
   cached
-  DataFlowCallable viableCallable(DataFlowCall call) { result = viableSourceCallable(call) }
+  DataFlowCallable viableCallable(DataFlowCall call) {
+    result.asCfgScope() = getTargetInstance(call.asCall())
+    or
+    result = any(AdditionalCallTarget t).viableTarget(call.asCall())
+  }
+
+  cached
+  CfgScope getTarget(DataFlowCall call) { result = viableCallable(call).asCfgScope() }
 
   cached
   newtype TArgumentPosition =