Python: subscript on local source nodes

yoff · yoff · commit 6114d71d3dc2 · 2022-09-26T13:39:59.000+02:00
and adjust comment on awaited
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/LocalSources.qll b/python/ql/lib/semmle/python/dataflow/new/internal/LocalSources.qll
@@ -101,6 +101,11 @@ class LocalSourceNode extends Node {
    */
   Node getAnAwaited() { Cached::await(this, result) }
 
+  /**
+   * Gets a subscript of this node.
+   */
+  CfgNode getASubscript() { Cached::subscript(this, result) }
+
   /**
    * Gets a call to the method `methodName` on this node.
    *
@@ -233,7 +238,7 @@ private module Cached {
   }
 
   /**
-   * Holds if `node` flows to the awaited value of `awaited`.
+   * Holds if `node` flows to a value that, when awaited, results in `awaited`.
    */
   cached
   predicate await(LocalSourceNode node, Node awaited) {
@@ -242,4 +247,15 @@ private module Cached {
       awaited = awaited(awaitedValue)
     )
   }
+
+  /**
+   * Holds if `node` flows to a sequence of which `subscript` is a subscript.
+   */
+  cached
+  predicate subscript(LocalSourceNode node, CfgNode subscript) {
+    exists(CfgNode seq, SubscriptNode subscriptNode | subscriptNode = subscript.getNode() |
+      node.flowsTo(seq) and
+      seq.getNode() = subscriptNode.getObject()
+    )
+  }
 }
