Python: override genEnclosingCallable

yoff · yoff · commit 5a864aab8773 · 2020-06-19T07:29:46.000+02:00
achieved flow out of functions!
diff --git a/python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll b/python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll
@@ -122,11 +122,19 @@ class ReturnKind extends TReturnKind {
 
 /** A data flow node that represents a value returned by a callable. */
 class ReturnNode extends Node {
+  Return ret;
+
   // See `TaintTrackingImplementation::returnFlowStep`
-  ReturnNode() { this.asCfgNode() = any(Return r).getValue().getAFlowNode() }
+  ReturnNode() {
+    this.asCfgNode() = ret.getValue().getAFlowNode()
+  }
 
   /** Gets the kind of this return node. */
   ReturnKind getKind() { result = TNormalReturnKind() }
+
+  override DataFlowCallable getEnclosingCallable() {
+    result.getScope().getAStmt() = ret // TODO: check nested function definitions
+  }
 }
 
 /** A data flow node that represents the output of a call. */
diff --git a/python/ql/src/experimental/dataflow/internal/DataFlowPublic.qll b/python/ql/src/experimental/dataflow/internal/DataFlowPublic.qll
@@ -48,7 +48,7 @@ class Node extends TNode {
   }
 
   /** Gets the enclosing callable of this node. */
-  final DataFlowCallable getEnclosingCallable() {
+  DataFlowCallable getEnclosingCallable() {
     none()
   }
 
@@ -99,13 +99,17 @@ class ParameterNode extends Node {
     this.asEssaNode() instanceof ParameterDefinition
   }
 
-    /**
+  /**
    * Holds if this node is the parameter of callable `c` at the specified
    * (zero-based) position.
    */
   predicate isParameterOf(DataFlowCallable c, int i) {
     this.asEssaNode().(ParameterDefinition).getDefiningNode() = c.getParameter(i)
   }
+
+  override DataFlowCallable getEnclosingCallable() {
+    this.isParameterOf(result, _)
+  }
 }
 
 /**
diff --git a/python/ql/test/experimental/dataflow/callGraph.expected b/python/ql/test/experimental/dataflow/callGraph.expected
@@ -0,0 +1 @@
+| test.py:4:10:4:10 | ControlFlowNode for z | test.py:7:5:7:20 | ControlFlowNode for obfuscated_id() |

Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,7 @@ class Node extends TNode {`
`48`	`48`	`}`
`49`	`49`
`50`	`50`	`/** Gets the enclosing callable of this node. */`
`51`		`- final DataFlowCallable getEnclosingCallable() {`
	`51`	`+ DataFlowCallable getEnclosingCallable() {`
`52`	`52`	`none()`
`53`	`53`	`}`
`54`	`54`
`@@ -99,13 +99,17 @@ class ParameterNode extends Node {`
`99`	`99`	`this.asEssaNode() instanceof ParameterDefinition`
`100`	`100`	`}`
`101`	`101`
`102`		`- /**`
	`102`	`+ /**`
`103`	`103`	* Holds if this node is the parameter of callable `c` at the specified
`104`	`104`	`* (zero-based) position.`
`105`	`105`	`*/`
`106`	`106`	`predicate isParameterOf(DataFlowCallable c, int i) {`
`107`	`107`	`this.asEssaNode().(ParameterDefinition).getDefiningNode() = c.getParameter(i)`
`108`	`108`	`}`
	`109`	`+`
	`110`	`+ override DataFlowCallable getEnclosingCallable() {`
	`111`	`+ this.isParameterOf(result, _)`
	`112`	`+ }`
`109`	`113`	`}`
`110`	`114`
`111`	`115`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+\| test.py:4:10:4:10 \| ControlFlowNode for z \| test.py:7:5:7:20 \| ControlFlowNode for obfuscated_id() \|`