Swift: Add dataflow through key-path expressios by modeling them as lambdas that perform a sequence of read steps.

Author: MathiasVP

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

@@ -179,3 +179,15 @@ class ApplyExprCfgNode extends ExprCfgNode {
 class CallExprCfgNode extends ApplyExprCfgNode {
   override CallExpr e;
 }
+
+class KeyPathApplicationExprCfgNode extends ExprCfgNode {
+  override KeyPathApplicationExpr e;
+
+  CfgNode getKeyPath() { result.getAst() = e.getKeyPath() }
+
+  CfgNode getBase() { result.getAst() = e.getBase() }
+}
+
+class KeyPathExprCfgNode extends ExprCfgNode {
+  override KeyPathExpr e;
+}

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,17 +108,16 @@ 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
       )
@@ -97,7 +133,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;
 

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

@@ -40,6 +40,22 @@ private class ExprNodeImpl extends ExprNode, NodeImpl {
   override DataFlowCallable getEnclosingCallable() { result = TDataFlowFunc(n.getScope()) }
 }
 
+private class KeyPathComponentNodeImpl extends TKeyPathComponentNode, NodeImpl {
+  KeyPathComponent component;
+
+  KeyPathComponentNodeImpl() { this = TKeyPathComponentNode(component) }
+
+  override Location getLocationImpl() { result = component.getLocation() }
+
+  override string toStringImpl() { result = component.toString() }
+
+  override DataFlowCallable getEnclosingCallable() {
+    result.asSourceCallable() = component.getKeyPathExpr()
+  }
+
+  KeyPathComponent getComponent() { result = component }
+}
+
 private class PatternNodeImpl extends PatternNode, NodeImpl {
   override Location getLocationImpl() { result = pattern.getLocation() }
 
@@ -78,6 +94,9 @@ private module Cached {
       FlowSummaryImpl::Private::summaryNodeRange(c, state)
     } or
     TSourceParameterNode(ParamDecl param) or
+    TKeyPathParameterNode(EntryNode entry) { entry.getScope() instanceof KeyPathExpr } or
+    TKeyPathReturnNode(ExitNode entry) { entry.getScope() instanceof KeyPathExpr } or
+    TKeyPathComponentNode(KeyPathComponent component) or
     TSummaryParameterNode(FlowSummary::SummarizedCallable c, ParameterPosition pos) {
       FlowSummaryImpl::Private::summaryParameterNodeRange(c, pos)
     } or
@@ -105,7 +124,7 @@ private module Cached {
     (
       nodeTo instanceof InoutReturnNode
       implies
-      nodeTo.(InoutReturnNode).getParameter() = def.getSourceVariable()
+      nodeTo.(InoutReturnNode).getParameter() = def.getSourceVariable().asVarDecl()
     )
   }
 
@@ -135,7 +154,7 @@ private module Cached {
       (
         nodeTo instanceof InoutReturnNode
         implies
-        nodeTo.(InoutReturnNode).getParameter() = def.getSourceVariable()
+        nodeTo.(InoutReturnNode).getParameter() = def.getSourceVariable().asVarDecl()
       )
       or
       // use-use flow
@@ -198,6 +217,11 @@ private module Cached {
         nodeFrom.asPattern().(TypedPattern).getSubPattern()
       ]
     or
+    // Flow from the unique parameter of a key path expression to
+    // the first component in the chain.
+    nodeTo.(KeyPathComponentNodeImpl).getComponent() =
+      nodeFrom.(KeyPathParameterNode).getComponent(0)
+    or
     // flow through a flow summary (extension of `SummaryModelCsv`)
     FlowSummaryImpl::Private::Steps::summaryLocalStep(nodeFrom, nodeTo, true)
   }
@@ -311,6 +335,28 @@ private module ParameterNodes {
 
     override DataFlowCallable getEnclosingCallable() { this.isParameterOf(result, _) }
   }
+
+  class KeyPathParameterNode extends ParameterNodeImpl, TKeyPathParameterNode {
+    private EntryNode entry;
+
+    KeyPathParameterNode() { this = TKeyPathParameterNode(entry) }
+
+    override predicate isParameterOf(DataFlowCallable c, ParameterPosition pos) {
+      c.asSourceCallable() = entry.getScope() and pos = TThisParameter()
+    }
+
+    override Location getLocationImpl() { result = entry.getLocation() }
+
+    override string toStringImpl() { result = entry.toString() }
+
+    override DataFlowCallable getEnclosingCallable() { this.isParameterOf(result, _) }
+
+    KeyPathComponent getComponent(int i) { result = entry.getScope().(KeyPathExpr).getComponent(i) }
+
+    KeyPathComponent getAComponent() { result = this.getComponent(_) }
+
+    KeyPathExpr getKeyPathExpr() { result = entry.getScope() }
+  }
 }
 
 import ParameterNodes
@@ -412,6 +458,17 @@ private module ArgumentNodes {
       FlowSummaryImpl::Private::summaryArgumentNode(call, this, pos)
     }
   }
+
+  class KeyPathArgumentNode extends ExprNode, ArgumentNode {
+    private KeyPathApplicationExprCfgNode keyPath;
+
+    KeyPathArgumentNode() { keyPath.getBase() = this.getCfgNode() }
+
+    override predicate argumentOf(DataFlowCall call, ArgumentPosition pos) {
+      call.asKeyPath() = keyPath and
+      pos = TThisArgument()
+    }
+  }
 }
 
 import ArgumentNodes
@@ -474,6 +531,24 @@ private module ReturnNodes {
 
     override ReturnKind getKind() { result = rk }
   }
+
+  class KeyPathReturnNodeImpl extends ReturnNode, TKeyPathReturnNode, NodeImpl {
+    ExitNode exit;
+
+    KeyPathReturnNodeImpl() { this = TKeyPathReturnNode(exit) }
+
+    override ReturnKind getKind() { result instanceof NormalReturnKind }
+
+    override ControlFlowNode getCfgNode() { result = exit }
+
+    override DataFlowCallable getEnclosingCallable() { result.asSourceCallable() = exit.getScope() }
+
+    override Location getLocationImpl() { result = exit.getLocation() }
+
+    override string toStringImpl() { result = exit.toString() }
+
+    KeyPathExpr getKeyPathExpr() { result = exit.getScope() }
+  }
 }
 
 import ReturnNodes
@@ -495,6 +570,16 @@ private module OutNodes {
     }
   }
 
+  class KeyPathOutNode extends OutNode, ExprNodeImpl {
+    KeyPathApplicationExprCfgNode keyPath;
+
+    KeyPathOutNode() { keyPath = this.getCfgNode() }
+
+    override DataFlowCall getCall(ReturnKind kind) {
+      result.asKeyPath() = keyPath and kind instanceof NormalReturnKind
+    }
+  }
+
   class SummaryOutNode extends OutNode, SummaryNode {
     SummaryOutNode() { FlowSummaryImpl::Private::summaryOutNode(_, this, _) }
 
@@ -658,6 +743,20 @@ predicate readStep(Node node1, ContentSet c, Node node2) {
     node2.asPattern() = pat.getSubPattern() and
     c instanceof OptionalSomeContentSet
   )
+  or
+  // read of a component in a key-path expression chain
+  exists(KeyPathComponent component, FieldDecl f |
+    component = node1.(KeyPathComponentNodeImpl).getComponent() and
+    f = component.getDeclRef() and
+    c.isSingleton(any(Content::FieldContent ct | ct.getField() = f))
+  |
+    // the next node is either the next element in the chain
+    node2.(KeyPathComponentNodeImpl).getComponent() = component.getNextComponent()
+    or
+    // or the return node, if this is the last component in the chain
+    not exists(component.getNextComponent()) and
+    node2.(KeyPathReturnNodeImpl).getKeyPathExpr() = component.getKeyPathExpr()
+  )
 }
 
 /**
@@ -786,6 +885,9 @@ predicate lambdaCall(DataFlowCall call, LambdaCallKind kind, Node receiver) {
   or
   kind = TLambdaCallKind() and
   receiver = call.(SummaryCall).getReceiver()
+  or
+  kind = TLambdaCallKind() and
+  receiver.asExpr() = call.asKeyPath().getExpr().(KeyPathApplicationExpr).getKeyPath()
 }
 
 /** Extra data-flow steps needed for lambda flow analysis. */

swift/ql/lib/codeql/swift/elements/KeyPathComponent.qll

@@ -1,4 +1,5 @@
 private import codeql.swift.generated.KeyPathComponent
+private import swift
 
 class KeyPathComponent extends Generated::KeyPathComponent {
   /**
@@ -35,4 +36,23 @@ class KeyPathComponent extends Generated::KeyPathComponent {
    * Tuple indexing like `.1`.
    */
   predicate isTupleIndexing() { getKind() = 9 }
+
+  /** Gets the underlying key-path expression which this is a component of. */
+  KeyPathExpr getKeyPathExpr() { result.getAComponent() = this }
+
+  /** Holds if this component is the i'th component of the underling key-path expression. */
+  predicate hasIndex(int i) { any(KeyPathExpr e).getComponent(i) = this }
+
+  /** Gets the next component of the underlying key-path expression. */
+  KeyPathComponent getNextComponent() {
+    exists(int i, KeyPathExpr e |
+      hasKeyPathExprAndIndex(e, i, this) and
+      hasKeyPathExprAndIndex(e, i + 1, result)
+    )
+  }
+}
+
+pragma[nomagic]
+private predicate hasKeyPathExprAndIndex(KeyPathExpr e, int i, KeyPathComponent c) {
+  e.getComponent(i) = c
 }