Merge pull request github#17653 from yoff/python/typetracking-through…

…-comprehensions
Kwstubbs · Oct 8, 2024 · 1f1b1b7 · 1f1b1b7
2 parents 3c1a19c + 6d486f9
commit 1f1b1b7
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/python/ql/lib/change-notes/2024-10-03-typetracking-through-comprehensions.md b/python/ql/lib/change-notes/2024-10-03-typetracking-through-comprehensions.md
@@ -0,0 +1,4 @@
+---
+category: minorAnalysis
+---
+* Type tracking, and hence the API graph, is now able to correctly trace trough comprehensions.
diff --git a/python/ql/lib/semmle/python/ApiGraphs.qll b/python/ql/lib/semmle/python/ApiGraphs.qll
@@ -843,6 +843,13 @@ module API {
         ref = pred.getSubscript(_) and
         ref.asCfgNode().isLoad()
         or
+        // Subscript via comprehension
+        lbl = Label::subscript() and
+        exists(PY::Comp comp |
+          pred.asExpr() = comp.getIterable() and
+          ref.asExpr() = comp.getNthInnerLoop(0).getTarget()
+        )
+        or
         // Subclassing a node
         lbl = Label::subclass() and
         exists(PY::ClassExpr clsExpr, DataFlow::Node superclass | pred.flowsTo(superclass) |

diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll b/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll
@@ -304,7 +304,7 @@ module TypeTrackingInput implements Shared::TypeTrackingInput {
       var.hasDefiningNode(def)
     |
       nodeTo.(DataFlowPublic::ScopeEntryDefinitionNode).getDefinition() = e and
-      nodeFrom.asCfgNode() = def.getValue() and
+      nodeFrom.asCfgNode() = def and
       var.getScope().getScope*() = nodeFrom.getScope()
     )
   }

diff --git a/python/ql/test/library-tests/frameworks/stdlib/http_server.py b/python/ql/test/library-tests/frameworks/stdlib/http_server.py
@@ -30,7 +30,7 @@ def test_cgi_FieldStorage_taint():
         form['key'][0].value, # $ tainted
         form['key'][0].file, # $ tainted
         form['key'][0].filename, # $ tainted
-        [field.value for field in form['key']], # $ MISSING: tainted
+        [field.value for field in form['key']], # $ tainted
 
         # `form.getvalue('key')` will be a list, if multiple fields named "key" are provided
         form.getvalue('key'), # $ tainted
@@ -40,7 +40,7 @@ def test_cgi_FieldStorage_taint():
 
         form.getlist('key'), # $ tainted
         form.getlist('key')[0], # $ tainted
-        [field.value for field in form.getlist('key')], # $ MISSING: tainted
+        [field.value for field in form.getlist('key')], # $ tainted
     )