Swift: flow through writeable keypaths

Author: rdmarsh2

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

@@ -57,6 +57,22 @@ private class KeyPathComponentNodeImpl extends TKeyPathComponentNode, NodeImpl {
   KeyPathComponent getComponent() { result = component }
 }
 
+private class KeyPathComponentPostUpdateNode extends TKeyPathComponentPostUpdateNode, NodeImpl {
+  KeyPathComponent component;
+
+  KeyPathComponentPostUpdateNode() { this = TKeyPathComponentPostUpdateNode(component) }
+
+  override Location getLocationImpl() { result = component.getLocation() }
+
+  override string toStringImpl() { result = "[post] " + 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() }
 
@@ -96,6 +112,9 @@ private module Cached {
     TKeyPathParameterNode(EntryNode entry) { entry.getScope() instanceof KeyPathExpr } or
     TKeyPathReturnNode(ExitNode exit) { exit.getScope() instanceof KeyPathExpr } or
     TKeyPathComponentNode(KeyPathComponent component) or
+    TKeyPathParameterPostUpdateNode(EntryNode entry) { entry.getScope() instanceof KeyPathExpr } or
+    TKeyPathReturnPostUpdateNode(ExitNode exit) { exit.getScope() instanceof KeyPathExpr } or
+    TKeyPathComponentPostUpdateNode(KeyPathComponent component) or
     TExprPostUpdateNode(CfgNode n) {
       // Obviously, the base of setters needs a post-update node
       n = any(PropertySetterCfgNode setter).getBase()
@@ -105,6 +124,8 @@ private module Cached {
       or
       n = any(PropertyObserverCfgNode getter).getBase()
       or
+      n = any(KeyPathApplicationExprCfgNode expr).getBase()
+      or
       // Arguments that are `inout` expressions needs a post-update node,
       // as well as any class-like argument (since a field can be modified).
       // Finally, qualifiers and bases of member reference need post-update nodes to support reverse reads.
@@ -227,6 +248,17 @@ private module Cached {
     nodeTo.(KeyPathComponentNodeImpl).getComponent() =
       nodeFrom.(KeyPathParameterNode).getComponent(0)
     or
+    nodeFrom.(KeyPathComponentPostUpdateNode).getComponent() = 
+      nodeTo.(KeyPathParameterPostUpdateNode).getComponent(0)
+    or
+    // Flow to the result of a keypath assignment
+    // TODO: is there a cleaner way to do this?
+    exists(KeyPathApplicationExpr apply, AssignExpr assign |
+      apply = assign.getDest() and
+      nodeTo.asExpr() = apply and
+      nodeFrom.asExpr() = assign.getSource()
+    )
+    or
     // flow through a flow summary (extension of `SummaryModelCsv`)
     FlowSummaryImpl::Private::Steps::summaryLocalStep(nodeFrom.(FlowSummaryNode).getSummaryNode(),
       nodeTo.(FlowSummaryNode).getSummaryNode(), true)
@@ -383,6 +415,56 @@ class FlowSummaryNode extends NodeImpl, TFlowSummaryNode {
   override string toStringImpl() { result = this.getSummaryNode().toString() }
 }
 
+class KeyPathParameterPostUpdateNode extends NodeImpl, ReturnNode, PostUpdateNodeImpl,
+  TKeyPathParameterPostUpdateNode
+{
+  private EntryNode entry;
+
+  KeyPathParameterPostUpdateNode() { this = TKeyPathParameterPostUpdateNode(entry) }
+
+  override KeyPathParameterNode getPreUpdateNode() {
+    result.getKeyPathExpr() = this.getKeyPathExpr()
+  }
+
+  override Location getLocationImpl() { result = entry.getLocation() }
+
+  override string toStringImpl() { result = "[post] " + entry.toString() }
+
+  override DataFlowCallable getEnclosingCallable() { result.asSourceCallable() = entry.getScope() }
+
+  KeyPathComponent getComponent(int i) { result = entry.getScope().(KeyPathExpr).getComponent(i) }
+
+  KeyPathComponent getAComponent() { result = this.getComponent(_) }
+
+  KeyPathExpr getKeyPathExpr() { result = entry.getScope() }
+
+  override ReturnKind getKind() { result.(ParamReturnKind).getIndex() = -1 }
+}
+
+class KeyPathReturnPostUpdateNode extends NodeImpl, ParameterNodeImpl, PostUpdateNodeImpl,
+  TKeyPathReturnPostUpdateNode
+{
+  private ExitNode exit;
+
+  KeyPathReturnPostUpdateNode() { this = TKeyPathReturnPostUpdateNode(exit) }
+
+  override KeyPathParameterNode getPreUpdateNode() {
+    result.getKeyPathExpr() = this.getKeyPathExpr()
+  }
+
+  override predicate isParameterOf(DataFlowCallable c, ParameterPosition pos) {
+    c.asSourceCallable() = this.getKeyPathExpr() and pos = TPositionalParameter(0) // TODO: new parameter type?
+  }
+
+  override Location getLocationImpl() { result = exit.getLocation() }
+
+  override string toStringImpl() { result = "[post] " + exit.toString() }
+
+  override DataFlowCallable getEnclosingCallable() { result.asSourceCallable() = exit.getScope() }
+
+  KeyPathExpr getKeyPathExpr() { result = exit.getScope() }
+}
+
 /** A data-flow node that represents a call argument. */
 abstract class ArgumentNode extends Node {
   /** Holds if this argument occurs at the given position in the given call. */
@@ -476,6 +558,20 @@ private module ArgumentNodes {
       pos = TThisArgument()
     }
   }
+
+  class KeyPathAssignmentArgumentNode extends ArgumentNode {
+    private KeyPathApplicationExprCfgNode keyPath;
+
+    KeyPathAssignmentArgumentNode() {
+      keyPath = this.getCfgNode() and
+      exists(AssignExpr assign | assign.getDest() = keyPath.getNode().asAstNode())
+    }
+
+    override predicate argumentOf(DataFlowCall call, ArgumentPosition pos) {
+      call.asKeyPath() = keyPath and
+      pos = TPositionalArgument(0) // TODO: new parameter type?
+    }
+  }
 }
 
 import ArgumentNodes
@@ -735,6 +831,24 @@ predicate storeStep(Node node1, ContentSet c, Node node2) {
     c instanceof OptionalSomeContentSet
   )
   or
+  // store through a component in a key-path expression chain
+  exists(KeyPathComponent component |
+    component = node2.(KeyPathComponentPostUpdateNode).getComponent() and
+    (
+      c.isSingleton(any(Content::FieldContent ct | ct.getField() = component.getDeclRef()))
+      or
+      c.isSingleton(any(Content::ArrayContent ac)) and
+      component.isSubscript()
+    )
+  |
+    // the previous node is either the next element in the chain
+    node1.(KeyPathComponentPostUpdateNode).getComponent() = component.getNextComponent()
+    or
+    // or the return node, if this is the last component in the chain
+    not exists(component.getNextComponent()) and
+    node1.(KeyPathReturnPostUpdateNode).getKeyPathExpr() = component.getKeyPathExpr()
+  )
+  or
   FlowSummaryImpl::Private::Steps::summaryStoreStep(node1.(FlowSummaryNode).getSummaryNode(), c,
     node2.(FlowSummaryNode).getSummaryNode())
 }

swift/ql/test/library-tests/dataflow/dataflow/DataFlow.expected

@@ -1,12 +1,14 @@
 edges
 | file://:0:0:0:0 | KeyPathComponent | test.swift:663:13:663:29 | exit #keyPath(...) [some:0] |
 | file://:0:0:0:0 | self [a, x] | file://:0:0:0:0 | .a [x] |
+| file://:0:0:0:0 | self [s, x] | file://:0:0:0:0 | .s [x] |
 | file://:0:0:0:0 | self [str] | file://:0:0:0:0 | .str |
 | file://:0:0:0:0 | self [v2, some:0] | file://:0:0:0:0 | .v2 [some:0] |
 | file://:0:0:0:0 | self [v2] | file://:0:0:0:0 | .v2 |
 | file://:0:0:0:0 | self [v3] | file://:0:0:0:0 | .v3 |
 | file://:0:0:0:0 | self [x, some:0] | file://:0:0:0:0 | .x [some:0] |
 | file://:0:0:0:0 | self [x] | file://:0:0:0:0 | .x |
+| file://:0:0:0:0 | self [x] | file://:0:0:0:0 | .x |
 | file://:0:0:0:0 | value | file://:0:0:0:0 | [post] self [v2] |
 | file://:0:0:0:0 | value | file://:0:0:0:0 | [post] self [v3] |
 | file://:0:0:0:0 | value | file://:0:0:0:0 | [post] self [x] |
@@ -299,6 +301,7 @@ edges
 | test.swift:599:24:599:32 | call to source3() | test.swift:599:13:599:33 | call to MyClass.init(s:) [str] |
 | test.swift:600:13:600:41 | call to MyClass.init(contentsOfFile:) [str] | test.swift:585:9:585:9 | self [str] |
 | test.swift:600:13:600:41 | call to MyClass.init(contentsOfFile:) [str] | test.swift:600:13:600:43 | .str |
+| test.swift:615:7:615:7 | self [x] | file://:0:0:0:0 | self [x] |
 | test.swift:617:8:617:11 | x | test.swift:618:14:618:14 | x |
 | test.swift:618:5:618:5 | [post] self [x] | test.swift:617:3:619:3 | self[return] [x] |
 | test.swift:618:14:618:14 | x | test.swift:618:5:618:5 | [post] self [x] |
@@ -314,6 +317,7 @@ edges
 | test.swift:627:38:627:38 | KeyPathComponent [x] | test.swift:627:36:627:38 | exit #keyPath(...) |
 | test.swift:628:13:628:13 | s [x] | test.swift:627:36:627:38 | enter #keyPath(...) [x] |
 | test.swift:628:13:628:13 | s [x] | test.swift:628:13:628:32 | \\...[...] |
+| test.swift:632:7:632:7 | self [s, x] | file://:0:0:0:0 | self [s, x] |
 | test.swift:634:8:634:11 | s [x] | test.swift:635:14:635:14 | s [x] |
 | test.swift:635:5:635:5 | [post] self [s, x] | test.swift:634:3:636:3 | self[return] [s, x] |
 | test.swift:635:14:635:14 | s [x] | test.swift:635:5:635:5 | [post] self [s, x] |
@@ -438,13 +442,26 @@ edges
 | test.swift:766:29:766:29 | KeyPathComponent [x] | test.swift:766:13:766:29 | exit #keyPath(...) |
 | test.swift:767:15:767:15 | s2 [s, some:0, x] | test.swift:766:13:766:29 | enter #keyPath(...) [s, some:0, x] |
 | test.swift:767:15:767:15 | s2 [s, some:0, x] | test.swift:767:15:767:28 | \\...[...] |
+| test.swift:773:11:773:17 | [post] exit #keyPath(...) | test.swift:773:17:773:17 | [post] KeyPathComponent [x] |
+| test.swift:773:15:773:15 | [post] KeyPathComponent [s, x] | test.swift:773:11:773:17 | [post] enter #keyPath(...) [s, x] |
+| test.swift:773:17:773:17 | [post] KeyPathComponent [x] | test.swift:773:15:773:15 | [post] KeyPathComponent [s, x] |
+| test.swift:774:3:774:3 | [post] s2 [s, x] | test.swift:775:13:775:13 | s2 [s, x] |
+| test.swift:774:3:774:16 | \\...[...] | test.swift:773:11:773:17 | [post] exit #keyPath(...) |
+| test.swift:774:3:774:16 | \\...[...] | test.swift:774:3:774:3 | [post] s2 [s, x] |
+| test.swift:774:20:774:27 | call to source() | test.swift:774:3:774:16 | \\...[...] |
+| test.swift:775:13:775:13 | s2 [s, x] | test.swift:632:7:632:7 | self [s, x] |
+| test.swift:775:13:775:13 | s2 [s, x] | test.swift:775:13:775:16 | .s [x] |
+| test.swift:775:13:775:16 | .s [x] | test.swift:615:7:615:7 | self [x] |
+| test.swift:775:13:775:16 | .s [x] | test.swift:775:13:775:18 | .x |
 nodes
 | file://:0:0:0:0 | .a [x] | semmle.label | .a [x] |
+| file://:0:0:0:0 | .s [x] | semmle.label | .s [x] |
 | file://:0:0:0:0 | .str | semmle.label | .str |
 | file://:0:0:0:0 | .v2 | semmle.label | .v2 |
 | file://:0:0:0:0 | .v2 [some:0] | semmle.label | .v2 [some:0] |
 | file://:0:0:0:0 | .v3 | semmle.label | .v3 |
 | file://:0:0:0:0 | .x | semmle.label | .x |
+| file://:0:0:0:0 | .x | semmle.label | .x |
 | file://:0:0:0:0 | .x [some:0] | semmle.label | .x [some:0] |
 | file://:0:0:0:0 | KeyPathComponent | semmle.label | KeyPathComponent |
 | file://:0:0:0:0 | [post] self [v2, some:0] | semmle.label | [post] self [v2, some:0] |
@@ -453,12 +470,14 @@ nodes
 | file://:0:0:0:0 | [post] self [x, some:0] | semmle.label | [post] self [x, some:0] |
 | file://:0:0:0:0 | [post] self [x] | semmle.label | [post] self [x] |
 | file://:0:0:0:0 | self [a, x] | semmle.label | self [a, x] |
+| file://:0:0:0:0 | self [s, x] | semmle.label | self [s, x] |
 | file://:0:0:0:0 | self [str] | semmle.label | self [str] |
 | file://:0:0:0:0 | self [v2, some:0] | semmle.label | self [v2, some:0] |
 | file://:0:0:0:0 | self [v2] | semmle.label | self [v2] |
 | file://:0:0:0:0 | self [v3] | semmle.label | self [v3] |
 | file://:0:0:0:0 | self [x, some:0] | semmle.label | self [x, some:0] |
 | file://:0:0:0:0 | self [x] | semmle.label | self [x] |
+| file://:0:0:0:0 | self [x] | semmle.label | self [x] |
 | file://:0:0:0:0 | value | semmle.label | value |
 | file://:0:0:0:0 | value | semmle.label | value |
 | file://:0:0:0:0 | value | semmle.label | value |
@@ -769,6 +788,7 @@ nodes
 | test.swift:599:24:599:32 | call to source3() | semmle.label | call to source3() |
 | test.swift:600:13:600:41 | call to MyClass.init(contentsOfFile:) [str] | semmle.label | call to MyClass.init(contentsOfFile:) [str] |
 | test.swift:600:13:600:43 | .str | semmle.label | .str |
+| test.swift:615:7:615:7 | self [x] | semmle.label | self [x] |
 | test.swift:617:3:619:3 | self[return] [x] | semmle.label | self[return] [x] |
 | test.swift:617:8:617:11 | x | semmle.label | x |
 | test.swift:618:5:618:5 | [post] self [x] | semmle.label | [post] self [x] |
@@ -785,6 +805,7 @@ nodes
 | test.swift:627:38:627:38 | KeyPathComponent [x] | semmle.label | KeyPathComponent [x] |
 | test.swift:628:13:628:13 | s [x] | semmle.label | s [x] |
 | test.swift:628:13:628:32 | \\...[...] | semmle.label | \\...[...] |
+| test.swift:632:7:632:7 | self [s, x] | semmle.label | self [s, x] |
 | test.swift:634:3:636:3 | self[return] [s, x] | semmle.label | self[return] [s, x] |
 | test.swift:634:8:634:11 | s [x] | semmle.label | s [x] |
 | test.swift:635:5:635:5 | [post] self [s, x] | semmle.label | [post] self [s, x] |
@@ -915,6 +936,16 @@ nodes
 | test.swift:766:29:766:29 | KeyPathComponent [x] | semmle.label | KeyPathComponent [x] |
 | test.swift:767:15:767:15 | s2 [s, some:0, x] | semmle.label | s2 [s, some:0, x] |
 | test.swift:767:15:767:28 | \\...[...] | semmle.label | \\...[...] |
+| test.swift:773:11:773:17 | [post] enter #keyPath(...) [s, x] | semmle.label | [post] enter #keyPath(...) [s, x] |
+| test.swift:773:11:773:17 | [post] exit #keyPath(...) | semmle.label | [post] exit #keyPath(...) |
+| test.swift:773:15:773:15 | [post] KeyPathComponent [s, x] | semmle.label | [post] KeyPathComponent [s, x] |
+| test.swift:773:17:773:17 | [post] KeyPathComponent [x] | semmle.label | [post] KeyPathComponent [x] |
+| test.swift:774:3:774:3 | [post] s2 [s, x] | semmle.label | [post] s2 [s, x] |
+| test.swift:774:3:774:16 | \\...[...] | semmle.label | \\...[...] |
+| test.swift:774:20:774:27 | call to source() | semmle.label | call to source() |
+| test.swift:775:13:775:13 | s2 [s, x] | semmle.label | s2 [s, x] |
+| test.swift:775:13:775:16 | .s [x] | semmle.label | .s [x] |
+| test.swift:775:13:775:18 | .x | semmle.label | .x |
 subpaths
 | test.swift:75:22:75:22 | x | test.swift:65:16:65:28 | arg1 | test.swift:65:1:70:1 | arg2[return] | test.swift:75:32:75:32 | [post] y |
 | test.swift:114:19:114:19 | arg | test.swift:109:9:109:14 | arg | test.swift:110:12:110:12 | arg | test.swift:114:12:114:22 | call to ... |
@@ -969,6 +1000,9 @@ subpaths
 | test.swift:764:18:764:25 | call to source() | test.swift:617:8:617:11 | x | test.swift:617:3:619:3 | self[return] [x] | test.swift:764:13:764:26 | call to S.init(x:) [x] |
 | test.swift:765:29:765:29 | s [some:0, x] | test.swift:655:8:655:12 | s [some:0, x] | test.swift:655:3:657:3 | self[return] [s, some:0, x] | test.swift:765:14:765:30 | call to S2_Optional.init(s:) [s, some:0, x] |
 | test.swift:767:15:767:15 | s2 [s, some:0, x] | test.swift:766:13:766:29 | enter #keyPath(...) [s, some:0, x] | test.swift:766:13:766:29 | exit #keyPath(...) | test.swift:767:15:767:28 | \\...[...] |
+| test.swift:774:3:774:16 | \\...[...] | test.swift:773:11:773:17 | [post] exit #keyPath(...) | test.swift:773:11:773:17 | [post] enter #keyPath(...) [s, x] | test.swift:774:3:774:3 | [post] s2 [s, x] |
+| test.swift:775:13:775:13 | s2 [s, x] | test.swift:632:7:632:7 | self [s, x] | file://:0:0:0:0 | .s [x] | test.swift:775:13:775:16 | .s [x] |
+| test.swift:775:13:775:16 | .s [x] | test.swift:615:7:615:7 | self [x] | file://:0:0:0:0 | .x | test.swift:775:13:775:18 | .x |
 #select
 | test.swift:7:15:7:15 | t1 | test.swift:6:19:6:26 | call to source() | test.swift:7:15:7:15 | t1 | result |
 | test.swift:9:15:9:15 | t1 | test.swift:6:19:6:26 | call to source() | test.swift:9:15:9:15 | t1 | result |
@@ -1070,3 +1104,4 @@ subpaths
 | test.swift:756:15:756:21 | ...! | test.swift:746:14:746:21 | call to source() | test.swift:756:15:756:21 | ...! | result |
 | test.swift:757:15:757:19 | .v3 | test.swift:747:14:747:21 | call to source() | test.swift:757:15:757:19 | .v3 | result |
 | test.swift:767:15:767:28 | \\...[...] | test.swift:764:18:764:25 | call to source() | test.swift:767:15:767:28 | \\...[...] | result |
+| test.swift:775:13:775:18 | .x | test.swift:774:20:774:27 | call to source() | test.swift:775:13:775:18 | .x | result |