Merge pull request #12807 from MathiasVP/dataflow-for-keypaths

Swift: Dataflow for keypaths

Author: MathiasVP

swift/ql/lib/codeql/swift/controlflow/CfgNodes.qll

@@ -179,3 +179,25 @@ class ApplyExprCfgNode extends ExprCfgNode {
 class CallExprCfgNode extends ApplyExprCfgNode {
   override CallExpr e;
 }
+
+/** A control-flow node that wraps a key-path application. */
+class KeyPathApplicationExprCfgNode extends ExprCfgNode {
+  override KeyPathApplicationExpr e;
+
+  /**
+   * Gets the control-flow node that wraps the key-path of
+   * this control-flow element.
+   */
+  CfgNode getKeyPath() { result.getAst() = e.getKeyPath() }
+
+  /**
+   * Gets the control-flow node that wraps the base of
+   * this control-flow element.
+   */
+  CfgNode getBase() { result.getAst() = e.getBase() }
+}
+
+/** A control-flow node that wraps a key-path expression. */
+class KeyPathExprCfgNode extends ExprCfgNode {
+  override KeyPathExpr e;
+}

swift/ql/lib/codeql/swift/controlflow/internal/ControlFlowGraphImpl.qll

@@ -58,7 +58,7 @@ module CfgScope {
   }
 
   private class KeyPathScope extends Range_ instanceof KeyPathExpr {
-    AstControlFlowTree tree;
+    KeyPathControlFlowTree tree;
 
     KeyPathScope() { tree.getAst() = this }
 
@@ -76,6 +76,12 @@ module CfgScope {
 
     final override predicate exit(ControlFlowElement last, Completion c) { last(tree, last, c) }
   }
+
+  private class KeyPathControlFlowTree extends StandardPostOrderTree, KeyPathElement {
+    final override ControlFlowElement getChildElement(int i) {
+      result.asAstNode() = expr.getComponent(i)
+    }
+  }
 }
 
 /** Holds if `first` is first executed when entering `scope`. */
@@ -88,6 +94,14 @@ predicate succExit(CfgScope::Range_ scope, ControlFlowElement last, Completion c
   scope.exit(last, c)
 }
 
+private class KeyPathComponentTree extends AstStandardPostOrderTree {
+  override KeyPathComponent ast;
+
+  final override ControlFlowElement getChildElement(int i) {
+    result.asAstNode() = ast.getSubscriptArgument(i).getExpr().getFullyConverted()
+  }
+}
+
 /**
  * Control-flow for statements.
  */

swift/ql/lib/codeql/swift/dataflow/Ssa.qll

@@ -21,7 +21,39 @@ module Ssa {
 
     class ExitBasicBlock = BasicBlocks::ExitBasicBlock;
 
-    class SourceVariable = VarDecl;
+    private newtype TSourceVariable =
+      TNormalSourceVariable(VarDecl v) or
+      TKeyPathSourceVariable(EntryNode entry) { entry.getScope() instanceof KeyPathExpr }
+
+    abstract class SourceVariable extends TSourceVariable {
+      abstract string toString();
+
+      VarDecl asVarDecl() { none() }
+
+      EntryNode asKeyPath() { none() }
+
+      DeclRefExpr getAnAccess() { result.getDecl() = this.asVarDecl() }
+    }
+
+    private class NormalSourceVariable extends SourceVariable, TNormalSourceVariable {
+      VarDecl v;
+
+      NormalSourceVariable() { this = TNormalSourceVariable(v) }
+
+      override string toString() { result = v.toString() }
+
+      override VarDecl asVarDecl() { result = v }
+    }
+
+    private class KeyPathSourceVariable extends SourceVariable, TKeyPathSourceVariable {
+      EntryNode enter;
+
+      KeyPathSourceVariable() { this = TKeyPathSourceVariable(enter) }
+
+      override string toString() { result = enter.toString() }
+
+      override EntryNode asKeyPath() { result = enter }
+    }
 
     predicate variableWrite(BasicBlock bb, int i, SourceVariable v, boolean certain) {
       exists(AssignExpr assign |
@@ -40,17 +72,22 @@ module Ssa {
       // ```
       exists(NamedPattern pattern |
         bb.getNode(i).getNode().asAstNode() = pattern and
-        v = pattern.getVarDecl() and
+        v.asVarDecl() = pattern.getVarDecl() and
         certain = true
       )
       or
-      v instanceof ParamDecl and
-      bb.getNode(i).getNode().asAstNode() = v and
+      exists(ParamDecl p |
+        p = v.asVarDecl() and
+        bb.getNode(i).getNode().asAstNode() = p and
+        certain = true
+      )
+      or
+      bb.getNode(i) = v.asKeyPath() and
       certain = true
       or
       // Mark the subexpression as a write of the local variable declared in the `TapExpr`.
       exists(TapExpr tap |
-        v = tap.getVar() and
+        v.asVarDecl() = tap.getVar() and
         bb.getNode(i).getNode().asAstNode() = tap.getSubExpr() and
         certain = true
       )
@@ -60,7 +97,7 @@ module Ssa {
       exists(DeclRefExpr ref |
         not isLValue(ref) and
         bb.getNode(i).getNode().asAstNode() = ref and
-        v = ref.getDecl() and
+        v.asVarDecl() = ref.getDecl() and
         certain = true
       )
       or
@@ -71,24 +108,26 @@ module Ssa {
       )
       or
       exists(ExitNode exit, AbstractFunctionDecl func |
-        func.getAParam() = v or func.getSelfParam() = v
-      |
+        [func.getAParam(), func.getSelfParam()] = v.asVarDecl() and
         bb.getNode(i) = exit and
-        modifiableParam(v) and
+        modifiableParam(v.asVarDecl()) and
         bb.getScope() = func and
         certain = true
       )
       or
       // Mark the `TapExpr` as a read of the of the local variable.
       exists(TapExpr tap |
-        v = tap.getVar() and
+        v.asVarDecl() = tap.getVar() and
         bb.getNode(i).getNode().asAstNode() = tap and
         certain = true
       )
     }
   }
 
-  private module SsaImpl = SsaImplCommon::Make<SsaInput>;
+  /**
+   * INTERNAL: Do not use.
+   */
+  module SsaImpl = SsaImplCommon::Make<SsaInput>;
 
   cached
   class Definition extends SsaImpl::Definition {
@@ -97,7 +136,7 @@ module Ssa {
 
     cached
     ControlFlowNode getARead() {
-      exists(VarDecl v, SsaInput::BasicBlock bb, int i |
+      exists(SsaInput::SourceVariable v, SsaInput::BasicBlock bb, int i |
         SsaImpl::ssaDefReachesRead(v, this, bb, i) and
         SsaInput::variableRead(bb, i, v, true) and
         result = bb.getNode(i)

swift/ql/lib/codeql/swift/dataflow/internal/DataFlowDispatch.qll

@@ -74,6 +74,7 @@ newtype TDataFlowCall =
   TPropertyGetterCall(PropertyGetterCfgNode getter) or
   TPropertySetterCall(PropertySetterCfgNode setter) or
   TPropertyObserverCall(PropertyObserverCfgNode observer) or
+  TKeyPathCall(KeyPathApplicationExprCfgNode keyPathApplication) or
   TSummaryCall(FlowSummaryImpl::Public::SummarizedCallable c, Node receiver) {
     FlowSummaryImpl::Private::summaryCallbackRange(c, receiver)
   }
@@ -89,6 +90,9 @@ class DataFlowCall extends TDataFlowCall {
   /** Gets the underlying source code call, if any. */
   ApplyExprCfgNode asCall() { none() }
 
+  /** Gets the underlying key-path application node, if any. */
+  KeyPathApplicationExprCfgNode asKeyPath() { none() }
+
   /**
    * Gets the i'th argument of call.class
    * The qualifier is considered to have index `-1`.
@@ -138,6 +142,25 @@ private class NormalCall extends DataFlowCall, TNormalCall {
   override Location getLocation() { result = apply.getLocation() }
 }
 
+private class KeyPathCall extends DataFlowCall, TKeyPathCall {
+  private KeyPathApplicationExprCfgNode apply;
+
+  KeyPathCall() { this = TKeyPathCall(apply) }
+
+  override KeyPathApplicationExprCfgNode asKeyPath() { result = apply }
+
+  override CfgNode getArgument(int i) {
+    i = -1 and
+    result = apply.getBase()
+  }
+
+  override DataFlowCallable getEnclosingCallable() { result = TDataFlowFunc(apply.getScope()) }
+
+  override string toString() { result = apply.toString() }
+
+  override Location getLocation() { result = apply.getLocation() }
+}
+
 class PropertyGetterCall extends DataFlowCall, TPropertyGetterCall {
   private PropertyGetterCfgNode getter;