Python: Add post-update nodes

python/ql/src/experimental/dataflow/internal/DataFlowImplConsistency.qll

@@ -123,8 +123,18 @@ module Consistency {
     n.getEnclosingCallable() != call.getEnclosingCallable()
   }
 
+  // This predicate helps the compiler forget that in some languages
+  // it is impossible for a result of `getPreUpdateNode` to be an
+  // instance of `PostUpdateNode`.
+  private Node getPre(PostUpdateNode n) {
+    result = n.getPreUpdateNode()
+    or
+    none()
+  }
+
   query predicate postIsNotPre(PostUpdateNode n, string msg) {
-    n.getPreUpdateNode() = n and msg = "PostUpdateNode should not equal its pre-update node."
+    getPre(n) = n and
+    msg = "PostUpdateNode should not equal its pre-update node."
   }
 
   query predicate postHasUniquePre(PostUpdateNode n, string msg) {
@@ -152,12 +162,6 @@ module Consistency {
     msg = "Origin of readStep is missing a PostUpdateNode."
   }
 
-  query predicate storeIsPostUpdate(Node n, string msg) {
-    storeStep(_, _, n) and
-    not n instanceof PostUpdateNode and
-    msg = "Store targets should be PostUpdateNodes."
-  }
-
   query predicate argHasPostUpdate(ArgumentNode n, string msg) {
     not hasPost(n) and
     not isImmutableOrUnobservable(n) and

python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll

@@ -15,6 +15,32 @@ class DataFlowCfgNode extends ControlFlowNode {
   DataFlowCfgNode() { isExpressionNode(this) }
 }
 
+/** A data flow node which should have an associated post-update node. */
+abstract class PreUpdateNode extends Node { }
+
+/** An argument might have its value changed as a result of a call. */
+class ArgumentPreUpdateNode extends PreUpdateNode, ArgumentNode { }
+
+/** An object might have its value changed after a store. */
+class StorePreUpdateNode extends PreUpdateNode, CfgNode {
+  StorePreUpdateNode() {
+    exists(Attribute a |
+      node = a.getObject().getAFlowNode() and
+      a.getCtx() instanceof Store
+    )
+  }
+}
+
+/** A node marking the state change of an object after a read */
+class ReadPreUpdateNode extends PreUpdateNode, CfgNode {
+  ReadPreUpdateNode() {
+    exists(Attribute a |
+      node = a.getObject().getAFlowNode() and
+      a.getCtx() instanceof Load
+    )
+  }
+}
+
 /**
  * A node associated with an object after an operation that might have
  * changed its state.
@@ -24,12 +50,21 @@ class DataFlowCfgNode extends ControlFlowNode {
  * an update to the field.
  *
  * Nodes corresponding to AST elements, for example `ExprNode`, usually refer
- * to the value before the update with the exception of `ObjectCreation`,
- * which represents the value after the constructor has run.
+ * to the value before the update.
  */
-abstract class PostUpdateNode extends Node {
+class PostUpdateNode extends Node, TPostUpdateNode {
+  PreUpdateNode pre;
+
+  PostUpdateNode() { this = TPostUpdateNode(pre) }
+
   /** Gets the node before the state update. */
-  abstract Node getPreUpdateNode();
+  Node getPreUpdateNode() { result = pre }
+
+  override string toString() { result = "[post] " + pre.toString() }
+
+  override Scope getScope() { result = pre.getScope() }
+
+  override Location getLocation() { result = pre.getLocation() }
 }
 
 class DataFlowExpr = Expr;
@@ -98,7 +133,17 @@ module EssaFlow {
 predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
   not nodeFrom.(EssaNode).getVar() instanceof GlobalSsaVariable and
   not nodeTo.(EssaNode).getVar() instanceof GlobalSsaVariable and
-  EssaFlow::essaFlowStep(nodeFrom, nodeTo)
+  EssaFlow::essaFlowStep(update(nodeFrom), nodeTo)
+}
+
+private Node update(Node node) {
+  exists(PostUpdateNode pun |
+    node = pun.getPreUpdateNode() and
+    result = pun
+  )
+  or
+  not exists(PostUpdateNode pun | node = pun.getPreUpdateNode()) and
+  result = node
 }
 
 // TODO: Make modules for these headings

python/ql/src/experimental/dataflow/internal/DataFlowPublic.qll

@@ -20,7 +20,9 @@ newtype TNode =
   /** A node corresponding to an SSA variable. */
   TEssaNode(EssaVariable var) or
   /** A node corresponding to a control flow node. */
-  TCfgNode(DataFlowCfgNode node)
+  TCfgNode(DataFlowCfgNode node) or
+  /** A node representing the value of an object after a state change */
+  TPostUpdateNode(PreUpdateNode pre)
 
 /**
  * An element, viewed as a node in a data flow graph. Either an SSA variable