Python: Fix getAGuardedNode

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

@@ -130,7 +130,7 @@ module EssaFlow {
   }
 
   predicate useToNextUse(NameNode nodeFrom, NameNode nodeTo) {
-    AdjacentUses::adjacentUseUseSameVar(nodeFrom, nodeTo)
+    AdjacentUses::adjacentUseUse(nodeFrom, nodeTo)
   }
 
   predicate defToFirstUse(EssaVariable var, NameNode nodeTo) {

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

@@ -5,6 +5,7 @@
 private import python
 private import DataFlowPrivate
 import experimental.dataflow.TypeTracker
+private import semmle.python.essa.SsaCompute
 
 /**
  * IPA type for data flow nodes.
@@ -176,10 +177,10 @@ class BarrierGuard extends GuardNode {
 
   /** Gets a node guarded by this guard. */
   final ExprNode getAGuardedNode() {
-    exists(Variable v, NameNode n, boolean testIsTrue |
+    exists(EssaDefinition def, ControlFlowNode node, boolean testIsTrue |
-      n.uses(v) and
+      AdjacentUses::aUse(def, node) and
-      this.checks(n, testIsTrue) and
+      this.checks(node, testIsTrue) and
-      result.asCfgNode().(NameNode).uses(v) and
+      AdjacentUses::aUse(def, result.asCfgNode()) and
       this.controlsNode(result.asCfgNode(), testIsTrue)
     )
   }

python/ql/src/semmle/python/essa/SsaCompute.qll

@@ -443,9 +443,28 @@ private module SsaComputeImpl {
       )
     }
 
+    /**
+     * Holds if `use1` and `use2` form an adjacent use-use-pair of the same
+     * `SsaSourceVariable`, that is, the value read in `use1` can reach `use2`
+     * without passing through any other use or any SSA definition of the variable
+     * except for phi nodes.
+     */
+    cached
+    predicate adjacentUseUse(ControlFlowNode use1, ControlFlowNode use2) {
+      adjacentUseUseSameVar(use1, use2)
+      or
+      exists(SsaSourceVariable v, EssaDefinition def, BasicBlock b1, int i1, BasicBlock b2, int i2 |
+        adjacentVarRefs(v, b1, i1, b2, i2) and
+        variableUse(v, use1, b1, i1) and
+        definesAt(def, v, b2, i2) and
+        firstUse(def, use2) and
+        def instanceof PhiFunction
+      )
+    }
+
     /**
      * Holds if the value defined at `def` can reach `use` without passing through
-     * any other uses, but possibly through phi nodes and uncertain implicit updates.
+     * any other uses, but possibly through phi nodes.
      */
     cached
     predicate firstUse(EssaDefinition def, ControlFlowNode use) {
@@ -482,6 +501,19 @@ private module SsaComputeImpl {
       b = def.(PhiFunction).getBasicBlock() and
       i = -1
     }
+
+    /**
+     * Holds if the value defined at `def` can reach `use`, possibly through phi nodes.
+     */
+    cached
+    predicate aUse(EssaDefinition def, ControlFlowNode use) {
+      firstUse(def, use)
+      or
+      exists(ControlFlowNode firstUse |
+        firstUse(def, firstUse) and
+        adjacentUseUse(firstUse, use)
+      )
+    }
   }
 }
 

python/ql/test/experimental/dataflow/tainttracking/customSanitizer/TestTaint.expected

@@ -2,16 +2,24 @@ test_taint
 | test.py:22 | ok   | test_custom_sanitizer | s |
 | test.py:36 | ok   | test_custom_sanitizer_guard | s |
 | test.py:38 | ok   | test_custom_sanitizer_guard | s |
-| test.py:49 | ok   | test_escape | s2 |
+| test.py:40 | ok   | test_custom_sanitizer_guard | s |
+| test.py:51 | ok   | test_escape | s2 |
 isSanitizer
 | TestTaintTrackingConfiguration | test.py:21:39:21:39 | ControlFlowNode for s |
-| TestTaintTrackingConfiguration | test.py:48:10:48:29 | ControlFlowNode for emulated_escaping() |
+| TestTaintTrackingConfiguration | test.py:50:10:50:29 | ControlFlowNode for emulated_escaping() |
 isSanitizerGuard
 | TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() |
 sanitizerGuardControls
 | TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:36:9:36:26 | ControlFlowNode for ensure_not_tainted | true |
 | TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:36:9:36:29 | ControlFlowNode for ensure_not_tainted() | true |
 | TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:36:28:36:28 | ControlFlowNode for s | true |
-| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:38:9:38:22 | ControlFlowNode for ensure_tainted | false |
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:37:9:37:9 | ControlFlowNode for s | true |
-| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:38:9:38:25 | ControlFlowNode for ensure_tainted() | false |
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:37:13:37:26 | ControlFlowNode for TAINTED_STRING | true |
-| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:38:24:38:24 | ControlFlowNode for s | false |
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:38:9:38:22 | ControlFlowNode for ensure_tainted | true |
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:38:9:38:25 | ControlFlowNode for ensure_tainted() | true |
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:38:24:38:24 | ControlFlowNode for s | true |
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:40:9:40:22 | ControlFlowNode for ensure_tainted | false |
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:40:9:40:25 | ControlFlowNode for ensure_tainted() | false |
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:40:24:40:24 | ControlFlowNode for s | false |
+sanitizerGuardedNode
+| TestTaintTrackingConfiguration | test.py:35:8:35:26 | ControlFlowNode for emulated_is_safe() | test.py:36:28:36:28 | ControlFlowNode for s |

python/ql/test/experimental/dataflow/tainttracking/customSanitizer/TestTaint.ql

@@ -40,3 +40,11 @@ query predicate sanitizerGuardControls(
   conf.isSanitizerGuard(guard) and
   guard.controlsNode(node, testIsTrue)
 }
+
+query predicate sanitizerGuardedNode(
+  TestTaintTrackingConfiguration conf, DataFlow::BarrierGuard guard, DataFlow::ExprNode node
+) {
+  exists(guard.getLocation().getFile().getRelativePath()) and
+  conf.isSanitizerGuard(guard) and
+  node = guard.getAGuardedNode()
+}

python/ql/test/experimental/dataflow/tainttracking/customSanitizer/test.py

@@ -34,6 +34,8 @@ def test_custom_sanitizer_guard():
 
     if emulated_is_safe(s):
         ensure_not_tainted(s)
+        s = TAINTED_STRING
+        ensure_tainted(s)
     else:
         ensure_tainted(s)