Python: Use CallableValue and improve tests

yoff · yoff · commit ce57a28c8f1a · 2020-06-17T19:12:10.000+02:00
diff --git a/python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll b/python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll
@@ -71,18 +71,24 @@ predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
 //--------
 
 /** Represents a callable */
-class DataFlowCallable = FunctionObject; // TODO: consider CallableValue
+class DataFlowCallable = CallableValue;
 
 /** Represents a call to a callable */
 class DataFlowCall extends CallNode {
+  DataFlowCallable callable;
+
+  DataFlowCall() {
+    this = callable.getACall()
+  }
+
   /** Gets the enclosing callable of this call. */
-  abstract DataFlowCallable getEnclosingCallable();
+  DataFlowCallable getEnclosingCallable() { result = callable }
 }
 
 /** A data flow node that represents a call argument. */
 class ArgumentNode extends Node {
   ArgumentNode() {
-    exists( DataFlowCall call, int pos |
+    exists(DataFlowCall call, int pos |
       this.asCfgNode() = call.getArg(pos)
     )
   }
@@ -96,14 +102,9 @@ class ArgumentNode extends Node {
   final DataFlowCall getCall() { this.argumentOf(result, _) }
 }
 
-import semmle.python.pointsto.CallGraph
-
 /** Gets a viable run-time target for the call `call`. */
 DataFlowCallable viableCallable(DataFlowCall call) {
-  exists(FunctionInvocation i |
-    call = i.getCall() and
-    result = i.getFunction()
-  )
+  result = call.getEnclosingCallable()
 }
 
 private newtype TReturnKind = TNormalReturnKind()
diff --git a/python/ql/test/experimental/dataflow/allFlowsConfig.qll b/python/ql/test/experimental/dataflow/allFlowsConfig.qll
@@ -9,6 +9,12 @@ class AllFlowsConfig extends DataFlow::Configuration {
 
   override predicate isSource(DataFlow::Node node) {
     node.asCfgNode().isEntryNode()
+    or
+    node = DataFlow::TEssaNode(_) and
+    not exists(DataFlow::Node pred |
+      pred = DataFlow::TEssaNode(_) and
+      DataFlow::localFlowStep(pred, node)
+    )
   }
 
   override predicate isSink(DataFlow::Node node) {
diff --git a/python/ql/test/experimental/dataflow/local.expected b/python/ql/test/experimental/dataflow/local.expected
@@ -2,5 +2,6 @@
 | test.py:1:5:1:5 | ControlFlowNode for IntegerLiteral | test.py:1:1:1:1 | GSSA Variable a |
 | test.py:2:5:2:5 | ControlFlowNode for a | test.py:2:1:2:1 | GSSA Variable b |
 | test.py:4:1:4:9 | ControlFlowNode for FunctionExpr | test.py:4:5:4:5 | GSSA Variable f |
+| test.py:4:7:4:7 | ControlFlowNode for x | test.py:4:7:4:7 | SSA variable x |
 | test.py:5:7:5:11 | ControlFlowNode for BinaryExpr | test.py:5:3:5:3 | SSA variable y |
 | test.py:8:5:8:8 | ControlFlowNode for f() | test.py:8:1:8:1 | GSSA Variable c |
diff --git a/python/ql/test/experimental/dataflow/sources.expected b/python/ql/test/experimental/dataflow/sources.expected
@@ -1,2 +1,12 @@
 | test.py:0:0:0:0 | Entry node for Module test |
+| test.py:0:0:0:0 | GSSA Variable __name__ |
+| test.py:0:0:0:0 | GSSA Variable __package__ |
+| test.py:0:0:0:0 | GSSA Variable c |
+| test.py:0:0:0:0 | SSA variable $ |
+| test.py:1:1:1:1 | GSSA Variable a |
+| test.py:2:1:2:1 | GSSA Variable b |
 | test.py:4:1:4:9 | Entry node for Function f |
+| test.py:4:5:4:5 | GSSA Variable f |
+| test.py:4:7:4:7 | SSA variable x |
+| test.py:5:3:5:3 | SSA variable y |
+| test.py:8:1:8:1 | GSSA Variable c |
