PS: Also support type tracking of objects constructed with New-Object.

Author: MathiasVP

powershell/ql/consistency-queries/TypeTrackingConsistency.ql

@@ -0,0 +1,8 @@
+import semmle.code.powershell.dataflow.DataFlow
+import semmle.code.powershell.typetracking.internal.TypeTrackingImpl
+
+private module ConsistencyChecksInput implements ConsistencyChecksInputSig {
+  predicate unreachableNodeExclude(DataFlow::Node n) { n instanceof DataFlow::PostUpdateNode }
+}
+
+import ConsistencyChecks<ConsistencyChecksInput>

powershell/ql/lib/powershell.qll

@@ -71,6 +71,7 @@ import semmle.code.powershell.StringConstantExpression
 import semmle.code.powershell.MemberExpr
 import semmle.code.powershell.InvokeMemberExpression
 import semmle.code.powershell.Call
+import semmle.code.powershell.ObjectCreation
 import semmle.code.powershell.SubExpression
 import semmle.code.powershell.ErrorExpr
 import semmle.code.powershell.ConvertExpr

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

@@ -12,6 +12,9 @@ abstract private class AbstractCall extends Ast {
   /** Gets the i'th positional argument to this call. */
   abstract Expr getPositionalArgument(int i);
 
+  /** Holds if an argument with name `name` is provided to this call. */
+  final predicate hasNamedArgument(string name) { exists(this.getNamedArgument(name)) }
+
   /** Gets the argument to this call with the name `name`. */
   abstract Expr getNamedArgument(string name);
 
@@ -25,7 +28,8 @@ abstract private class AbstractCall extends Ast {
   abstract Function getATarget();
 }
 
-private class CmdCall extends AbstractCall instanceof Cmd {
+/** A call to a command. For example, `Write-Host "Hello, world!"`. */
+class CmdCall extends AbstractCall instanceof Cmd {
   final override Expr getCommand() { result = Cmd.super.getCommand() }
 
   final override Expr getPositionalArgument(int i) { result = Cmd.super.getPositionalArgument(i) }
@@ -43,7 +47,8 @@ private class CmdCall extends AbstractCall instanceof Cmd {
   }
 }
 
-private class InvokeMemberCall extends AbstractCall instanceof InvokeMemberExpr {
+/** A call to a method on an object. For example, `$obj.ToString()`. */
+class MethodCall extends AbstractCall instanceof InvokeMemberExpr {
   final override Expr getCommand() { result = super.getMember() }
 
   final override Expr getPositionalArgument(int i) {

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

@@ -49,6 +49,9 @@ class Cmd extends @command, CmdBase {
       )
   }
 
+  /** Holds if this call has an argument named `name`. */
+  predicate hasNamedArgument(string name) { exists(this.getNamedArgument(name)) }
+
   /** Gets the named argument with the given name. */
   Expr getNamedArgument(string name) {
     exists(int i, CmdParameter p |

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

@@ -18,8 +18,30 @@ class InvokeMemberExpr extends @invoke_member_expression, MemberExprBase {
   override predicate isStatic() { this.getQualifier() instanceof TypeNameExpr }
 }
 
+/**
+ * A call to a constructor. For example:
+ *
+ * ```powershell
+ * [System.IO.FileInfo]::new("C:\\file.txt")
+ * ```
+ */
 class ConstructorCall extends InvokeMemberExpr {
-  ConstructorCall() { this.isStatic() and this.getName() = "new" }
-
-  Type getConstructedType() { result = this.getQualifier().(TypeNameExpr).getType() }
+  TypeNameExpr typename;
+
+  ConstructorCall() {
+    this.isStatic() and typename = this.getQualifier() and this.getName() = "new"
+  }
+
+  /**
+   * Gets the type being constructed by this constructor call.
+   *
+   * Note that the type may not exist in the database.
+   *
+   * Use `getConstructedTypeName` to get the name of the type (which will
+   * always exist in the database).
+   */
+  Type getConstructedType() { result = typename.getType() }
+
+  /** Gets the name of the type being constructed by this constructor call. */
+  string getConstructedTypeName() { result = typename.getName() }
 }

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

@@ -0,0 +1,55 @@
+import powershell
+
+abstract private class AbstractObjectCreation extends Call {
+  /**
+   * The type of the object being constructed.
+   * Note that the type may not exist in the database.
+   *
+   * Use `getConstructedTypeName` to get the name of the type (which will
+   * always exist in the database).
+   */
+  abstract Type getConstructedType();
+
+  /** The name of the type of the object being constructed. */
+  abstract string getConstructedTypeName();
+}
+
+/**
+ * An object creation from a call to a constructor. For example:
+ * ```powershell
+ * [System.IO.FileInfo]::new("C:\\file.txt")
+ * ```
+ */
+class NewObjectCreation extends AbstractObjectCreation instanceof ConstructorCall {
+  final override Type getConstructedType() { result = ConstructorCall.super.getConstructedType() }
+
+  final override string getConstructedTypeName() {
+    result = ConstructorCall.super.getConstructedTypeName()
+  }
+}
+
+/**
+ * An object creation from a call to `New-Object`. For example:
+ * ```powershell
+ * New-Object -TypeName System.IO.FileInfo -ArgumentList "C:\\file.txt"
+ * ```
+ */
+class DotNetObjectCreation extends AbstractObjectCreation instanceof Cmd {
+  DotNetObjectCreation() { this.getCommandName() = "New-Object" }
+
+  final override Type getConstructedType() { none() }
+
+  final override string getConstructedTypeName() {
+    // Either it's the named argument `TypeName`
+    result = Cmd.super.getNamedArgument("TypeName").(StringConstExpr).getValue().getValue()
+    or
+    // Or it's the first positional argument if that's the named argument
+    not Cmd.super.hasNamedArgument("TypeName") and
+    exists(StringConstExpr arg | arg = Cmd.super.getPositionalArgument(0) |
+      result = arg.getValue().getValue() and
+      not arg = Cmd.super.getNamedArgument(["ArgumentList", "Property"])
+    )
+  }
+}
+
+final class ObjectCreation = AbstractObjectCreation;

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

@@ -147,6 +147,18 @@ abstract private class AbstractCallCfgNode extends AstCfgNode {
 
 final class CallCfgNode = AbstractCallCfgNode;
 
+class ObjectCreationCfgNode extends CallCfgNode {
+  ObjectCreation objectCreation;
+
+  ObjectCreationCfgNode() { this.getAstNode() = objectCreation }
+
+  ObjectCreation getObjectCreation() { result = objectCreation }
+
+  Type getConstructedType() { result = objectCreation.getConstructedType() }
+
+  string getConstructedTypeName() { result = objectCreation.getConstructedTypeName() }
+}
+
 /** Provides classes for control-flow nodes that wrap AST expressions. */
 module ExprNodes {
   private class VarAccessChildMapping extends ExprChildMapping, VarAccess {
@@ -228,7 +240,7 @@ module ExprNodes {
 
     final override ExprCfgNode getAnArgument() { e.hasCfgChild(e.getAnArgument(), this, result) }
 
-    final override string getName() { none() }
+    final override string getName() { result = e.getName() }
 
     final override ExprCfgNode getCommand() { none() }
   }
@@ -249,6 +261,23 @@ module ExprNodes {
     InvokeMemberCfgNode getInvokeMember() { this = result.getQualifier() }
   }
 
+  class TypeNameChildMapping extends ExprChildMapping, TypeNameExpr {
+    override predicate relevantChild(Ast n) { none() }
+  }
+
+  /** A control-flow node that wraps a `TypeName` expression. */
+  class TypeNameCfgNode extends ExprCfgNode {
+    override string getAPrimaryQlClass() { result = "TypeNameCfgNode" }
+
+    override TypeNameChildMapping e;
+
+    override TypeNameExpr getExpr() { result = super.getExpr() }
+
+    Type getType() { result = this.getExpr().getType() }
+
+    string getTypeName() { result = this.getExpr().getName() }
+  }
+
   class ConditionalChildMapping extends ExprChildMapping, ConditionalExpr {
     override predicate relevantChild(Ast n) { n = this.getCondition() or n = this.getABranch() }
   }
@@ -352,4 +381,19 @@ module StmtNodes {
     /** Gets the RHS of this assignment. */
     final StmtCfgNode getRightHandSide() { s.hasCfgChild(s.getRightHandSide(), this, result) }
   }
+
+  class CmdExprChildMapping extends NonExprChildMapping, CmdExpr {
+    override predicate relevantChild(Ast n) { n = this.getExpr() }
+  }
+
+  /** A control-flow node that wraps a `CmdExpr` expression. */
+  class CmdExprCfgNode extends StmtCfgNode {
+    override string getAPrimaryQlClass() { result = "CmdExprCfgNode" }
+
+    override CmdExprChildMapping s;
+
+    override CmdExpr getStmt() { result = super.getStmt() }
+
+    final ExprCfgNode getExpr() { s.hasCfgChild(s.getExpr(), this, result) }
+  }
 }

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

@@ -122,21 +122,34 @@ class NormalCall extends DataFlowCall, TNormalCall {
   override Location getLocation() { result = c.getLocation() }
 }
 
-/** A call for which we want to compute call targets. */
-private class RelevantCall extends CfgNodes::CallCfgNode { }
+private predicate localFlowStep(Node nodeFrom, Node nodeTo, StepSummary summary) {
+  localFlowStepTypeTracker(nodeFrom, nodeTo) and
+  summary.toString() = "level"
+}
 
 private module TrackInstanceInput implements CallGraphConstruction::InputSig {
-  newtype State = additional MkState(Type m, Boolean exact)
+  private predicate start0(Node start, string typename, boolean exact) {
+    start.(ObjectCreationNode).getObjectCreationNode().getConstructedTypeName() = typename and
+    exact = true
+    or
+    start.asExpr().(CfgNodes::ExprNodes::TypeNameCfgNode).getTypeName() = typename and
+    exact = true
+  }
+
+  newtype State = additional MkState(string typename, Boolean exact) { start0(_, typename, exact) }
 
   predicate start(Node start, State state) {
-    exists(Type tp, boolean exact | state = MkState(tp, exact) |
-      start.asExpr().(CfgNodes::ExprNodes::ConstructorCallCfgNode).getConstructedType() = tp
+    exists(string typename, boolean exact |
+      state = MkState(typename, exact) and
+      start0(start, typename, exact)
     )
   }
 
   pragma[nomagic]
   predicate stepNoCall(Node nodeFrom, Node nodeTo, StepSummary summary) {
     smallStepNoCall(nodeFrom, nodeTo, summary)
+    or
+    localFlowStep(nodeFrom, nodeTo, summary)
   }
 
   predicate stepCall(Node nodeFrom, Node nodeTo, StepSummary summary) {
@@ -155,16 +168,22 @@ private predicate qualifiedCall(CfgNodes::CallCfgNode call, Node receiver, strin
   call.getName() = method
 }
 
-Node trackInstance(Type t, boolean exact) {
+Node trackInstance(string typename, boolean exact) {
   result =
-    CallGraphConstruction::Make<TrackInstanceInput>::track(TrackInstanceInput::MkState(t, exact))
+    CallGraphConstruction::Make<TrackInstanceInput>::track(TrackInstanceInput::MkState(typename,
+        exact))
 }
 
 private CfgScope getTargetInstance(CfgNodes::CallCfgNode call) {
-  exists(Node receiver, string method, Type t |
+  // TODO: Also match argument/parameter types
+  exists(Node receiver, string method, string typename, Type t |
     qualifiedCall(call, receiver, method) and
-    receiver = trackInstance(t, _) and
-    result = t.getMemberFunction(method).getBody()
+    receiver = trackInstance(typename, _) and
+    t.getName() = typename
+  |
+    if method = "new"
+    then result = t.getAConstructor().getBody()
+    else result = t.getMethod(method).getBody()
   )
 }
 

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

@@ -94,6 +94,8 @@ module LocalFlow {
     nodeFrom.asExpr() = nodeTo.asExpr().(CfgNodes::ExprNodes::ConditionalCfgNode).getABranch()
     or
     nodeFrom.asStmt() = nodeTo.asStmt().(CfgNodes::StmtNodes::AssignStmtCfgNode).getRightHandSide()
+    or
+    nodeFrom.asExpr() = nodeTo.asStmt().(CfgNodes::StmtNodes::CmdExprCfgNode).getExpr()
   }
 
   predicate localMustFlowStep(Node nodeFrom, Node nodeTo) {

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

@@ -228,3 +228,20 @@ module BarrierGuard<guardChecksSig/3 guardChecks> {
     none() // TODO
   }
 }
+
+/**
+ * A dataflow node that represents the creation of an object.
+ * 
+ * For example, `[Foo]::new()` or `New-Object Foo`.
+ */
+class ObjectCreationNode extends Node {
+  CfgNodes::ObjectCreationCfgNode objectCreation;
+
+  ObjectCreationNode() {
+    this.asExpr() = objectCreation
+    or
+    this.asStmt() = objectCreation
+  }
+
+  final CfgNodes::ObjectCreationCfgNode getObjectCreationNode() { result = objectCreation }
+}