Python: module for import time flow

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll

@@ -381,6 +381,86 @@ module EssaFlow {
 //--------
 // Local flow
 //--------
+/** A module for transforming step relations. */
+module StepRelationTransformations {
+  /**
+   * Holds if there is a step from `nodeFrom` to `nodeTo` in
+   * the step relation to be transformed.
+   *
+   * This is the input relation to the transformations.
+   */
+  signature predicate stepSig(Node nodeFrom, Node nodeTo);
+
+  /**
+   * A module to separate import-time from run-time.
+   *
+   * We really have two local flow relations, one for module initialisation time (or _import time_) and one for runtime.
+   * Consider a read from a global variable `x = foo`. At import time there should be a local flow step from `foo` to `x`,
+   * while at runtime there should be a jump step from the module variable corresponding to `foo` to `x`.
+   *
+   * Similarly, for a write `foo = y`, at import time, there is a local flow step from `y` to `foo` while at runtime there
+   * is a jump step from `y` to the module variable corresponding to `foo`.
+   *
+   * We need a way of distinguishing if we are looking at import time or runtime. We have the following helpful facts:
+   * - All top-level executable statements are import time (and import time only)
+   * - All non-top-level code may be executed at runtime (but could also be executed at import time)
+   *
+   * We could write an analysis to determine which functions are called at import time, but until we have that, we will go
+   * with the heuristic that global variables act according to import time rules at top-level program points and according
+   * to runtime rules everywhere else. This will forego some import time local flow but otherwise be consistent.
+   */
+  module PhaseDependentFlow<stepSig/2 rawStep> {
+    /**
+     * Holds if `node` is found at the top level of a module.
+     */
+    pragma[inline]
+    private predicate isTopLevel(Node node) { node.getScope() instanceof Module }
+
+    /** Holds if a step can be taken from `nodeFrom` to `nodeTo` at import time. */
+    predicate importTimeStep(Node nodeFrom, Node nodeTo) {
+      // As a proxy for whether statements can be executed at import time,
+      // we check if they appear at the top level.
+      // This will miss statements inside functions called from the top level.
+      isTopLevel(nodeFrom) and
+      isTopLevel(nodeTo) and
+      rawStep(nodeFrom, nodeTo)
+    }
+
+    /** Holds if a step can be taken from `nodeFrom` to `nodeTo` at runtime. */
+    predicate runtimeStep(Node nodeFrom, Node nodeTo) {
+      // Anything not at the top level can be executed at runtime.
+      not isTopLevel(nodeFrom) and
+      not isTopLevel(nodeTo) and
+      rawStep(nodeFrom, nodeTo)
+    }
+
+    /**
+     * Holds if a step can be taken from `nodeFrom` to `nodeTo`.
+     */
+    predicate step(Node nodeFrom, Node nodeTo) {
+      importTimeStep(nodeFrom, nodeTo) or
+      runtimeStep(nodeFrom, nodeTo)
+    }
+  }
+
+  /**
+   * A module to add steps from post-update nodes.
+   * Whenever there is a step from `x` to `y`,
+   * we add a step from `[post] x` to `y`.
+   */
+  module IncludePostUpdateFlow<stepSig/2 rawStep> {
+    predicate step(Node nodeFrom, Node nodeTo) {
+      // We either have a raw step from `nodeFrom`...
+      rawStep(nodeFrom, nodeTo)
+      or
+      // ...or we have a raw step from a pre-update node of `nodeFrom`
+      rawStep(nodeFrom.(PostUpdateNode).getPreUpdateNode(), nodeTo)
+    }
+  }
+}
+
+import StepRelationTransformations
+
 /**
  * This is the local flow predicate that is used as a building block in global
  * data flow.
@@ -401,69 +481,16 @@ predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
  * or at runtime when callables in the module are called.
  */
 predicate simpleLocalFlowStepForTypetracking(Node nodeFrom, Node nodeTo) {
-  // If there is local flow out of a node `node`, we want flow
-  // both out of `node` and any post-update node of `node`.
-  exists(Node node |
-    nodeFrom = update(node) and
-    (
-      importTimeLocalFlowStep(node, nodeTo) or
-      runtimeLocalFlowStep(node, nodeTo)
-    )
-  )
-}
-
-/**
- * Holds if `node` is found at the top level of a module.
- */
-pragma[inline]
-predicate isTopLevel(Node node) { node.getScope() instanceof Module }
-
-/** Holds if there is local flow from `nodeFrom` to `nodeTo` at import time. */
-predicate importTimeLocalFlowStep(Node nodeFrom, Node nodeTo) {
-  // As a proxy for whether statements can be executed at import time,
-  // we check if they appear at the top level.
-  // This will miss statements inside functions called from the top level.
-  isTopLevel(nodeFrom) and
-  isTopLevel(nodeTo) and
-  EssaFlow::essaFlowStep(nodeFrom, nodeTo)
-}
-
-/** Holds if there is local flow from `nodeFrom` to `nodeTo` at runtime. */
-predicate runtimeLocalFlowStep(Node nodeFrom, Node nodeTo) {
-  // Anything not at the top level can be executed at runtime.
-  not isTopLevel(nodeFrom) and
-  not isTopLevel(nodeTo) and
-  EssaFlow::essaFlowStep(nodeFrom, nodeTo)
+  IncludePostUpdateFlow<PhaseDependentFlow<EssaFlow::essaFlowStep/2>::step/2>::step(nodeFrom, nodeTo)
 }
 
-predicate summaryFlowSteps(Node nodeFrom, Node nodeTo) {
-  // If there is local flow out of a node `node`, we want flow
-  // both out of `node` and any post-update node of `node`.
-  exists(Node node |
-    nodeFrom = update(node) and
-    (
-      importTimeSummaryFlowStep(node, nodeTo) or
-      runtimeSummaryFlowStep(node, nodeTo)
-    )
-  )
-}
-
-predicate importTimeSummaryFlowStep(Node nodeFrom, Node nodeTo) {
-  // As a proxy for whether statements can be executed at import time,
-  // we check if they appear at the top level.
-  // This will miss statements inside functions called from the top level.
-  isTopLevel(nodeFrom) and
-  isTopLevel(nodeTo) and
+private predicate summaryLocalStep(Node nodeFrom, Node nodeTo) {
   FlowSummaryImpl::Private::Steps::summaryLocalStep(nodeFrom.(FlowSummaryNode).getSummaryNode(),
     nodeTo.(FlowSummaryNode).getSummaryNode(), true)
 }
 
-predicate runtimeSummaryFlowStep(Node nodeFrom, Node nodeTo) {
-  // Anything not at the top level can be executed at runtime.
-  not isTopLevel(nodeFrom) and
-  not isTopLevel(nodeTo) and
-  FlowSummaryImpl::Private::Steps::summaryLocalStep(nodeFrom.(FlowSummaryNode).getSummaryNode(),
-    nodeTo.(FlowSummaryNode).getSummaryNode(), true)
+predicate summaryFlowSteps(Node nodeFrom, Node nodeTo) {
+  IncludePostUpdateFlow<PhaseDependentFlow<summaryLocalStep/2>::step/2>::step(nodeFrom, nodeTo)
 }
 
 /** `ModuleVariable`s are accessed via jump steps at runtime. */
@@ -487,15 +514,6 @@ predicate runtimeJumpStep(Node nodeFrom, Node nodeTo) {
   )
 }
 
-/**
- * Holds if `result` is either `node`, or the post-update node for `node`.
- */
-private Node update(Node node) {
-  result = node
-  or
-  result.(PostUpdateNode).getPreUpdateNode() = node
-}
-
 //--------
 // Type pruning
 //--------

python/ql/test/experimental/dataflow/module-initialization/localFlow.ql

@@ -12,7 +12,7 @@ module ImportTimeLocalFlowTest implements FlowTestSig {
     // results are displayed next to `nodeTo`, so we need a line to write on
     nodeTo.getLocation().getStartLine() > 0 and
     nodeTo.asVar() instanceof GlobalSsaVariable and
-    DP::importTimeLocalFlowStep(nodeFrom, nodeTo)
+    DP::PhaseDependentFlow<DP::EssaFlow::essaFlowStep/2>::importTimeStep(nodeFrom, nodeTo)
   }
 }