Python: Implement OutNode

Also, fix test for local flow

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

@@ -26,30 +26,6 @@ abstract class PostUpdateNode extends Node {
   abstract Node getPreUpdateNode();
 }
 
-private newtype TReturnKind = TNormalReturnKind()
-
-/**
- * A return kind. A return kind describes how a value can be returned
- * from a callable. For Python, this is simply a method return.
- */
-class ReturnKind extends TReturnKind {
-  /** Gets a textual representation of this return kind. */
-  string toString() { result = "return" }
-}
-
-/** A data flow node that represents a value returned by a callable. */
-abstract class ReturnNode extends Node {
-  /** Gets the kind of this return node. */
-  abstract ReturnKind getKind();
-}
-
-/** A data flow node that represents the output of a call. */
-abstract class OutNode extends Node {
-  /** Gets the underlying call, where this node is a corresponding output of kind `kind`. */
-  cached
-  abstract DataFlowCall getCall(ReturnKind kind);
-}
-
 class DataFlowExpr = Expr;
 
 //--------
@@ -76,19 +52,26 @@ predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
     nodeTo.asEssaNode() = p.getVariable() and
     nodeFrom.asEssaNode() = p.getShortCircuitInput()
   )
+  // or
+  // exists(EssaNodeDefinition d |
+  //   nodeTo.asEssaNode() = d.getVariable() and
+  //   nodeFrom.asEssaNode().getDefinition().getLocation() = d.(AssignmentDefinition).getValue().getLocation() // TODO: A better way to tie these together
+  // )
   or
-  exists(EssaNodeDefinition d |
-    nodeTo.asEssaNode() = d.getVariable() and
-    nodeFrom.asEssaNode().getDefinition().getLocation() = d.(AssignmentDefinition).getValue().getLocation() // TODO: A better way to tie these together
-  )
+  // As in `taintedAssignment`
+  // `x = f(42)`
+  // nodeTo is any use of `x`
+  // nodeFrom is `f(42)`
+  nodeFrom.asCfgNode() = nodeTo.asEssaNode().getDefinition().(AssignmentDefinition).getValue()
 }
 
+// TODO: Make modules for these headings
 //--------
 // Global flow
 //--------
 
 /** Represents a callable */
-class DataFlowCallable = FunctionObject;
+class DataFlowCallable = FunctionObject; // TODO: consider CallableValue
 
 /** Represents a call to a callable */
 class DataFlowCall extends CallNode {
@@ -112,7 +95,37 @@ import semmle.python.pointsto.CallGraph
 DataFlowCallable viableCallable(DataFlowCall call) {
   exists(FunctionInvocation i |
     call = i.getCall() and
-    result = i.getFunction())
+    result = i.getFunction()
+  )
+}
+
+private newtype TReturnKind = TNormalReturnKind()
+
+/**
+ * A return kind. A return kind describes how a value can be returned
+ * from a callable. For Python, this is simply a method return.
+ */
+class ReturnKind extends TReturnKind {
+  /** Gets a textual representation of this return kind. */
+  string toString() { result = "return" }
+}
+
+/** A data flow node that represents a value returned by a callable. */
+abstract class ReturnNode extends Node {
+  /** Gets the kind of this return node. */
+  abstract ReturnKind getKind();
+}
+
+/** A data flow node that represents the output of a call. */
+class OutNode extends Node {
+  OutNode() { this.asCfgNode() instanceof CallNode}
+
+  /** Gets the underlying call, where this node is a corresponding output of kind `kind`. */
+  cached
+  DataFlowCall getCall(ReturnKind kind) {
+    kind = TNormalReturnKind() and
+    result = this.asCfgNode().(CallNode)
+  }
 }
 
 /**
@@ -121,9 +134,6 @@ DataFlowCallable viableCallable(DataFlowCall call) {
  */
 OutNode getAnOutNode(DataFlowCall call, ReturnKind kind) { call = result.getCall(kind) }
 
-// Extend OutNode here
-// Consider whether to use AST nodes rather than CFG nodes 
-
 //--------
 // Type pruning
 //--------

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

@@ -6,13 +6,17 @@ import python
 private import DataFlowPrivate
 
 /**
- * IPA type for dta flow nodes.
+ * IPA type for data flow nodes.
+ * 
+ * Flow between SSA variables are computed in `Essa.qll`
+ * Flow from SSA variables to control flow nodes is as in
+ * `EssaTaintTracking`.
  */
 newtype TNode =
-  /**
-   * A node corresponding to local flow as computed via SSA.
-   */
-  TEssaNode(EssaVariable var)
+  /** A node corresponding to an SSA variable. */
+  TEssaNode(EssaVariable var) or
+  /** A node corresponding to a control flow node. */
+  TCfgNode(ControlFlowNode node)
 
 /**
  * An element, viewed as a node in a data flow graph. Either an expression
@@ -25,10 +29,19 @@ class Node extends TNode {
    */
   EssaVariable asEssaNode() { this = TEssaNode(result) }
 
+  /**
+   * Get the underlying ControlFlowNode if this is such a node.
+   */
+  ControlFlowNode asCfgNode() { this = TCfgNode(result) }
+
   /**
    * Get a string representation of this data flow node.
    */
-  string toString() { result = this.asEssaNode().toString() }
+  string toString() {
+    result = this.asEssaNode().toString()
+    or
+    result = this.asCfgNode().toString()
+  }
 
   /** Gets the enclosing callable of this node. */
   final DataFlowCallable getEnclosingCallable() {
@@ -53,6 +66,8 @@ class Node extends TNode {
     string filepath, int startline, int startcolumn, int endline, int endcolumn
   ) {
     this.asEssaNode().getDefinition().getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn)
+    or
+    this.asCfgNode().getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn)
   }