Merge pull request #4265 from tausbn/python-add-global-flow-steps

Python: Add `ModuleVariableNode` to keep track of global reads and writes

Author: yoff

python/ql/src/experimental/dataflow/TypeTracker.qll

@@ -47,11 +47,11 @@ class StepSummary extends TStepSummary {
 module StepSummary {
   cached
   predicate step(Node nodeFrom, Node nodeTo, StepSummary summary) {
-    exists(Node mid | EssaFlow::essaFlowStep*(nodeFrom, mid) and smallstep(mid, nodeTo, summary))
+    exists(Node mid | typePreservingStep*(nodeFrom, mid) and smallstep(mid, nodeTo, summary))
   }
 
   predicate smallstep(Node nodeFrom, Node nodeTo, StepSummary summary) {
-    EssaFlow::essaFlowStep(nodeFrom, nodeTo) and
+    typePreservingStep(nodeFrom, nodeTo) and
     summary = LevelStep()
     or
     callStep(nodeFrom, nodeTo) and summary = CallStep()
@@ -68,6 +68,12 @@ module StepSummary {
   }
 }
 
+/** Holds if it's reasonable to expect the data flow step from `nodeFrom` to `nodeTo` to preserve types. */
+private predicate typePreservingStep(Node nodeFrom, Node nodeTo) {
+  EssaFlow::essaFlowStep(nodeFrom, nodeTo) or
+  jumpStep(nodeFrom, nodeTo)
+}
+
 /** Holds if `nodeFrom` steps to `nodeTo` by being passed as a parameter in a call. */
 predicate callStep(ArgumentNode nodeFrom, ParameterNode nodeTo) {
   // TODO: Support special methods?
@@ -111,7 +117,7 @@ predicate returnStep(ReturnNode nodeFrom, Node nodeTo) {
 predicate basicStoreStep(Node nodeFrom, Node nodeTo, string attr) {
   exists(AttributeAssignment a, Node var |
     a.getName() = attr and
-    EssaFlow::essaFlowStep*(nodeTo, var) and
+    simpleLocalFlowStep*(nodeTo, var) and
     var.asVar() = a.getInput() and
     nodeFrom.asCfgNode() = a.getValue()
   )
@@ -276,7 +282,7 @@ class TypeTracker extends TTypeTracker {
       result = this.append(summary)
     )
     or
-    EssaFlow::essaFlowStep(nodeFrom, nodeTo) and
+    typePreservingStep(nodeFrom, nodeTo) and
     result = this
   }
 }

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

@@ -150,10 +150,28 @@ predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
   // If there is ESSA-flow out of a node `node`, we want flow
   // both out of `node` and any post-update node of `node`.
   exists(Node node |
-    not node.(EssaNode).getVar() instanceof GlobalSsaVariable and
-    not nodeTo.(EssaNode).getVar() instanceof GlobalSsaVariable and
     EssaFlow::essaFlowStep(node, nodeTo) and
-    nodeFrom = update(node)
+    nodeFrom = update(node) and
+    (
+      not node instanceof EssaNode or
+      not nodeTo instanceof EssaNode or
+      localEssaStep(node, nodeTo)
+    )
+  )
+}
+
+/**
+ * Holds if there is an Essa flow step from `nodeFrom` to `nodeTo` that does not switch between
+ * local and global SSA variables.
+ */
+private predicate localEssaStep(EssaNode nodeFrom, EssaNode nodeTo) {
+  EssaFlow::essaFlowStep(nodeFrom, nodeTo) and
+  (
+    nodeFrom.getVar() instanceof GlobalSsaVariable and
+    nodeTo.getVar() instanceof GlobalSsaVariable
+    or
+    not nodeFrom.getVar() instanceof GlobalSsaVariable and
+    not nodeTo.getVar() instanceof GlobalSsaVariable
   )
 }
 
@@ -179,7 +197,8 @@ private Node update(Node node) {
  */
 newtype TDataFlowCallable =
   TCallableValue(CallableValue callable) or
-  TClassValue(ClassValue c)
+  TClassValue(ClassValue c) or
+  TModule(Module m)
 
 /** Represents a callable */
 abstract class DataFlowCallable extends TDataFlowCallable {
@@ -233,6 +252,23 @@ class DataFlowClassValue extends DataFlowCallable, TClassValue {
   override string getName() { result = c.getName() }
 }
 
+/** A class representing the scope in which a `ModuleVariableNode` appears. */
+class DataFlowModuleScope extends DataFlowCallable, TModule {
+  Module mod;
+
+  DataFlowModuleScope() { this = TModule(mod) }
+
+  override string toString() { result = mod.toString() }
+
+  override CallNode getACall() { none() }
+
+  override Scope getScope() { result = mod }
+
+  override NameNode getParameter(int n) { none() }
+
+  override string getName() { result = mod.getName() }
+}
+
 newtype TDataFlowCall =
   TCallNode(CallNode call) or
   TSpecialCall(SpecialMethodCallNode special)
@@ -389,21 +425,11 @@ string ppReprType(DataFlowType t) { none() }
  * taken into account.
  */
 predicate jumpStep(Node nodeFrom, Node nodeTo) {
-  // As we have ESSA variables for global variables,
-  // we include ESSA flow steps involving global variables.
-  (
-    nodeFrom.(EssaNode).getVar() instanceof GlobalSsaVariable
-    or
-    nodeTo.(EssaNode).getVar() instanceof GlobalSsaVariable
-  ) and
-  (
-    EssaFlow::essaFlowStep(nodeFrom, nodeTo)
-    or
-    // As jump steps do not respect chronology,
-    // we add jump steps for each def-use pair.
-    nodeFrom.asVar() instanceof GlobalSsaVariable and
-    nodeTo.asCfgNode() = nodeFrom.asVar().getASourceUse()
-  )
+  // Module variable read
+  nodeFrom.(ModuleVariableNode).getARead() = nodeTo
+  or
+  // Module variable write
+  nodeFrom = nodeTo.(ModuleVariableNode).getAWrite()
 }
 
 //--------

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

@@ -24,7 +24,9 @@ newtype TNode =
   /** A node corresponding to a control flow node. */
   TCfgNode(DataFlowCfgNode node) or
   /** A node representing the value of an object after a state change */
-  TPostUpdateNode(PreUpdateNode pre)
+  TPostUpdateNode(PreUpdateNode pre) or
+  /** A node representing a global (module-level) variable in a specific module */
+  TModuleVariableNode(Module m, GlobalVariable v) { v.getScope() = m and v.escapes() }
 
 /**
  * An element, viewed as a node in a data flow graph. Either an SSA variable
@@ -149,6 +151,72 @@ class ParameterNode extends EssaNode {
   override DataFlowCallable getEnclosingCallable() { this.isParameterOf(result, _) }
 }
 
+/**
+ * A data flow node corresponding to a module-level (global) variable that is accessed outside of the module scope.
+ *
+ * Global variables may appear twice in the data flow graph, as both `EssaNode`s and
+ * `ModuleVariableNode`s. The former is used to represent data flow between global variables as it
+ * occurs during module initialization, and the latter is used to represent data flow via global
+ * variable reads and writes during run-time.
+ *
+ * It is possible for data to flow from assignments made at module initialization time to reads made
+ * at run-time, but not vice versa. For example, there will be flow from `SOURCE` to `SINK` in the
+ * following snippet:
+ *
+ * ```python
+ * g = SOURCE
+ *
+ * def foo():
+ *     SINK(g)
+ * ```
+ * but not the other way round:
+ *
+ * ```python
+ * SINK(g)
+ *
+ * def bar()
+ *     global g
+ *     g = SOURCE
+ * ```
+ *
+ * Data flow through `ModuleVariableNode`s is represented as `jumpStep`s, and so any write of a
+ * global variable can flow to any read of the same variable.
+ */
+class ModuleVariableNode extends Node, TModuleVariableNode {
+  Module mod;
+  GlobalVariable var;
+
+  ModuleVariableNode() { this = TModuleVariableNode(mod, var) }
+
+  override Scope getScope() { result = mod }
+
+  override string toString() {
+    result = "ModuleVariableNode for " + var.toString() + " in " + mod.toString()
+  }
+
+  /** Gets the module in which this variable appears. */
+  Module getModule() { result = mod }
+
+  /** Gets the global variable corresponding to this node. */
+  GlobalVariable getVariable() { result = var }
+
+  /** Gets a node that reads this variable. */
+  Node getARead() {
+    result.asCfgNode() = var.getALoad().getAFlowNode() and
+    // Ignore reads that happen when the module is imported. These are only executed once.
+    not result.getScope() = mod
+  }
+
+  /** Gets an `EssaNode` that corresponds to an assignment of this global variable. */
+  EssaNode getAWrite() {
+    result.asVar().getDefinition().(EssaNodeDefinition).definedBy(var, any(DefinitionNode defn))
+  }
+
+  override DataFlowCallable getEnclosingCallable() { result.(DataFlowModuleScope).getScope() = mod }
+
+  override Location getLocation() { result = mod.getLocation() }
+}
+
 /**
  * A node that controls whether other nodes are evaluated.
  */

python/ql/test/experimental/dataflow/basic/globalStep.expected

@@ -1,5 +1,7 @@
 | test.py:1:1:1:21 | ControlFlowNode for FunctionExpr | test.py:1:5:1:17 | GSSA Variable obfuscated_id |
 | test.py:1:1:1:21 | ControlFlowNode for FunctionExpr | test.py:1:5:1:17 | GSSA Variable obfuscated_id |
+| test.py:1:1:1:21 | ControlFlowNode for FunctionExpr | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id |
+| test.py:1:1:1:21 | ControlFlowNode for FunctionExpr | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id |
 | test.py:1:5:1:17 | GSSA Variable obfuscated_id | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id |
 | test.py:1:5:1:17 | GSSA Variable obfuscated_id | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id |
 | test.py:1:19:1:19 | SSA variable x | test.py:2:3:2:3 | SSA variable y |
@@ -70,6 +72,10 @@
 | test.py:6:1:6:1 | GSSA Variable a | test.py:7:19:7:19 | ControlFlowNode for a |
 | test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:6:1:6:1 | GSSA Variable a |
 | test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:6:1:6:1 | GSSA Variable a |
+| test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:7:5:7:20 | GSSA Variable a |
+| test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:7:5:7:20 | GSSA Variable a |
+| test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:7:19:7:19 | ControlFlowNode for a |
+| test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:7:19:7:19 | ControlFlowNode for a |
 | test.py:7:5:7:20 | ControlFlowNode for obfuscated_id() | test.py:7:1:7:1 | GSSA Variable b |
 | test.py:7:5:7:20 | ControlFlowNode for obfuscated_id() | test.py:7:1:7:1 | GSSA Variable b |
 | test.py:7:19:7:19 | ControlFlowNode for a | test.py:1:19:1:19 | SSA variable x |

python/ql/test/experimental/dataflow/basic/local.expected

@@ -3,8 +3,11 @@
 | test.py:0:0:0:0 | GSSA Variable b | test.py:0:0:0:0 | GSSA Variable b |
 | test.py:0:0:0:0 | SSA variable $ | test.py:0:0:0:0 | SSA variable $ |
 | test.py:1:1:1:21 | ControlFlowNode for FunctionExpr | test.py:1:1:1:21 | ControlFlowNode for FunctionExpr |
+| test.py:1:1:1:21 | ControlFlowNode for FunctionExpr | test.py:1:5:1:17 | GSSA Variable obfuscated_id |
+| test.py:1:1:1:21 | ControlFlowNode for FunctionExpr | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id |
 | test.py:1:5:1:17 | ControlFlowNode for obfuscated_id | test.py:1:5:1:17 | ControlFlowNode for obfuscated_id |
 | test.py:1:5:1:17 | GSSA Variable obfuscated_id | test.py:1:5:1:17 | GSSA Variable obfuscated_id |
+| test.py:1:5:1:17 | GSSA Variable obfuscated_id | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id |
 | test.py:1:19:1:19 | ControlFlowNode for x | test.py:1:19:1:19 | ControlFlowNode for x |
 | test.py:1:19:1:19 | SSA variable x | test.py:1:19:1:19 | SSA variable x |
 | test.py:1:19:1:19 | SSA variable x | test.py:2:3:2:3 | SSA variable y |
@@ -31,10 +34,16 @@
 | test.py:4:10:4:10 | ControlFlowNode for z | test.py:4:10:4:10 | ControlFlowNode for z |
 | test.py:6:1:6:1 | ControlFlowNode for a | test.py:6:1:6:1 | ControlFlowNode for a |
 | test.py:6:1:6:1 | GSSA Variable a | test.py:6:1:6:1 | GSSA Variable a |
+| test.py:6:1:6:1 | GSSA Variable a | test.py:7:5:7:20 | GSSA Variable a |
+| test.py:6:1:6:1 | GSSA Variable a | test.py:7:19:7:19 | ControlFlowNode for a |
+| test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:6:1:6:1 | GSSA Variable a |
 | test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral |
+| test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:7:5:7:20 | GSSA Variable a |
+| test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:7:19:7:19 | ControlFlowNode for a |
 | test.py:7:1:7:1 | ControlFlowNode for b | test.py:7:1:7:1 | ControlFlowNode for b |
 | test.py:7:1:7:1 | GSSA Variable b | test.py:7:1:7:1 | GSSA Variable b |
 | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id |
+| test.py:7:5:7:20 | ControlFlowNode for obfuscated_id() | test.py:7:1:7:1 | GSSA Variable b |
 | test.py:7:5:7:20 | ControlFlowNode for obfuscated_id() | test.py:7:5:7:20 | ControlFlowNode for obfuscated_id() |
 | test.py:7:5:7:20 | GSSA Variable a | test.py:7:5:7:20 | GSSA Variable a |
 | test.py:7:19:7:19 | ControlFlowNode for a | test.py:7:19:7:19 | ControlFlowNode for a |

python/ql/test/experimental/dataflow/basic/localStep.expected

@@ -1,5 +1,11 @@
+| test.py:1:1:1:21 | ControlFlowNode for FunctionExpr | test.py:1:5:1:17 | GSSA Variable obfuscated_id |
+| test.py:1:5:1:17 | GSSA Variable obfuscated_id | test.py:7:5:7:17 | ControlFlowNode for obfuscated_id |
 | test.py:1:19:1:19 | SSA variable x | test.py:2:7:2:7 | ControlFlowNode for x |
 | test.py:2:3:2:3 | SSA variable y | test.py:3:7:3:7 | ControlFlowNode for y |
 | test.py:2:7:2:7 | ControlFlowNode for x | test.py:2:3:2:3 | SSA variable y |
 | test.py:3:3:3:3 | SSA variable z | test.py:4:10:4:10 | ControlFlowNode for z |
 | test.py:3:7:3:7 | ControlFlowNode for y | test.py:3:3:3:3 | SSA variable z |
+| test.py:6:1:6:1 | GSSA Variable a | test.py:7:5:7:20 | GSSA Variable a |
+| test.py:6:1:6:1 | GSSA Variable a | test.py:7:19:7:19 | ControlFlowNode for a |
+| test.py:6:5:6:6 | ControlFlowNode for IntegerLiteral | test.py:6:1:6:1 | GSSA Variable a |
+| test.py:7:5:7:20 | ControlFlowNode for obfuscated_id() | test.py:7:1:7:1 | GSSA Variable b |