Merge pull request github#14513 from RasmusWL/yield-modeling

Python: Improve `yield` modeling

Author: RasmusWL

python/ql/lib/change-notes/2023-10-17-contextmanager.md

@@ -0,0 +1,4 @@
+---
+category: minorAnalysis
+---
+* Added support for functions decorated with `contextlib.contextmanager`.

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

@@ -240,6 +240,19 @@ predicate hasPropertyDecorator(Function func) {
   )
 }
 
+/**
+ * Holds if the function `func` has a `contextlib.contextmanager`.
+ */
+predicate hasContextmanagerDecorator(Function func) {
+  exists(ControlFlowNode contextmanager |
+    contextmanager.(NameNode).getId() = "contextmanager" and contextmanager.(NameNode).isGlobal()
+    or
+    contextmanager.(AttrNode).getObject("contextmanager").(NameNode).getId() = "contextlib"
+  |
+    func.getADecorator() = contextmanager.getNode()
+  )
+}
+
 // =============================================================================
 // Callables
 // =============================================================================
@@ -1604,6 +1617,24 @@ class ExtractedReturnNode extends ReturnNode, CfgNode {
   override ReturnKind getKind() { any() }
 }
 
+/**
+ * A data flow node that represents the value yielded by a callable with a
+ * `contextlib.contextmanager` decorator. We treat this as a normal return, which makes
+ * things just work when used in a `with` statement -- technically calling the function
+ * directly will give you a `contextlib._GeneratorContextManager` instance, so it's a
+ * slight workaround solution.
+ *
+ * See https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager
+ */
+class YieldNodeInContextManagerFunction extends ReturnNode, CfgNode {
+  YieldNodeInContextManagerFunction() {
+    hasContextmanagerDecorator(node.getScope()) and
+    node = any(Yield yield).getValue().getAFlowNode()
+  }
+
+  override ReturnKind getKind() { any() }
+}
+
 /** A data-flow node that represents the output of a call. */
 abstract class OutNode extends Node {
   /** Gets the underlying call, where this node is a corresponding output of kind `kind`. */

python/ql/test/experimental/dataflow/typetracking/test.py

@@ -195,3 +195,42 @@ def other_name(self):
     def with_global_modifier(self):
         global some_value
         print(some_value) # $ tracked
+
+# ------------------------------------------------------------------------------
+# yield
+# ------------------------------------------------------------------------------
+
+def yielding_function():
+    x = tracked # $ tracked
+    yield x # $ tracked
+
+def test_yield():
+    for x in yielding_function(): # $ MISSING: tracked
+        print(x) # $ MISSING: tracked
+
+    gen = yielding_function()
+    y = next(gen) # $ MISSING: tracked
+    print(y) # $ MISSING: tracked
+
+
+# see https://docs.python.org/3.11/library/contextlib.html#contextlib.contextmanager
+from contextlib import contextmanager
+import contextlib
+
+@contextmanager
+def managed_resource():
+    x = tracked # $ tracked
+    yield x # $ tracked
+
+def test_context_manager():
+    with managed_resource() as x: # $ tracked
+        print(x) # $ tracked
+
+@contextlib.contextmanager
+def managed_resource2():
+    x = tracked # $ tracked
+    yield x # $ tracked
+
+def test_context_manager2():
+    with managed_resource2() as x: # $ tracked
+        print(x) # $ tracked

python/ql/test/experimental/query-tests/Security/CWE-022-TarSlip/TarSlip.expected

@@ -32,6 +32,12 @@ edges
 | TarSlipImprov.py:141:6:141:29 | ControlFlowNode for Attribute() | TarSlipImprov.py:141:34:141:36 | GSSA Variable tar |
 | TarSlipImprov.py:141:34:141:36 | GSSA Variable tar | TarSlipImprov.py:142:9:142:13 | GSSA Variable entry |
 | TarSlipImprov.py:142:9:142:13 | GSSA Variable entry | TarSlipImprov.py:143:36:143:40 | ControlFlowNode for entry |
+| TarSlipImprov.py:151:14:151:50 | ControlFlowNode for closing() | TarSlipImprov.py:151:55:151:56 | SSA variable tf |
+| TarSlipImprov.py:151:22:151:49 | ControlFlowNode for Attribute() | TarSlipImprov.py:151:14:151:50 | ControlFlowNode for closing() |
+| TarSlipImprov.py:151:55:151:56 | SSA variable tf | TarSlipImprov.py:152:19:152:20 | ControlFlowNode for tf |
+| TarSlipImprov.py:152:19:152:20 | ControlFlowNode for tf | TarSlipImprov.py:157:18:157:40 | ControlFlowNode for py2_tarxz() |
+| TarSlipImprov.py:157:9:157:14 | SSA variable tar_cm | TarSlipImprov.py:162:20:162:23 | SSA variable tarc |
+| TarSlipImprov.py:157:18:157:40 | ControlFlowNode for py2_tarxz() | TarSlipImprov.py:157:9:157:14 | SSA variable tar_cm |
 | TarSlipImprov.py:159:9:159:14 | SSA variable tar_cm | TarSlipImprov.py:162:20:162:23 | SSA variable tarc |
 | TarSlipImprov.py:159:18:159:52 | ControlFlowNode for closing() | TarSlipImprov.py:159:9:159:14 | SSA variable tar_cm |
 | TarSlipImprov.py:159:26:159:51 | ControlFlowNode for Attribute() | TarSlipImprov.py:159:18:159:52 | ControlFlowNode for closing() |
@@ -122,6 +128,12 @@ nodes
 | TarSlipImprov.py:141:34:141:36 | GSSA Variable tar | semmle.label | GSSA Variable tar |
 | TarSlipImprov.py:142:9:142:13 | GSSA Variable entry | semmle.label | GSSA Variable entry |
 | TarSlipImprov.py:143:36:143:40 | ControlFlowNode for entry | semmle.label | ControlFlowNode for entry |
+| TarSlipImprov.py:151:14:151:50 | ControlFlowNode for closing() | semmle.label | ControlFlowNode for closing() |
+| TarSlipImprov.py:151:22:151:49 | ControlFlowNode for Attribute() | semmle.label | ControlFlowNode for Attribute() |
+| TarSlipImprov.py:151:55:151:56 | SSA variable tf | semmle.label | SSA variable tf |
+| TarSlipImprov.py:152:19:152:20 | ControlFlowNode for tf | semmle.label | ControlFlowNode for tf |
+| TarSlipImprov.py:157:9:157:14 | SSA variable tar_cm | semmle.label | SSA variable tar_cm |
+| TarSlipImprov.py:157:18:157:40 | ControlFlowNode for py2_tarxz() | semmle.label | ControlFlowNode for py2_tarxz() |
 | TarSlipImprov.py:159:9:159:14 | SSA variable tar_cm | semmle.label | SSA variable tar_cm |
 | TarSlipImprov.py:159:18:159:52 | ControlFlowNode for closing() | semmle.label | ControlFlowNode for closing() |
 | TarSlipImprov.py:159:26:159:51 | ControlFlowNode for Attribute() | semmle.label | ControlFlowNode for Attribute() |
@@ -199,6 +211,7 @@ subpaths
 | TarSlipImprov.py:130:5:130:7 | ControlFlowNode for tar | TarSlipImprov.py:129:6:129:26 | ControlFlowNode for Attribute() | TarSlipImprov.py:130:5:130:7 | ControlFlowNode for tar | Extraction of tarfile from $@ to a potentially untrusted source $@. | TarSlipImprov.py:129:6:129:26 | ControlFlowNode for Attribute() | ControlFlowNode for Attribute() | TarSlipImprov.py:130:5:130:7 | ControlFlowNode for tar | ControlFlowNode for tar |
 | TarSlipImprov.py:134:1:134:3 | ControlFlowNode for tar | TarSlipImprov.py:133:7:133:39 | ControlFlowNode for Attribute() | TarSlipImprov.py:134:1:134:3 | ControlFlowNode for tar | Extraction of tarfile from $@ to a potentially untrusted source $@. | TarSlipImprov.py:133:7:133:39 | ControlFlowNode for Attribute() | ControlFlowNode for Attribute() | TarSlipImprov.py:134:1:134:3 | ControlFlowNode for tar | ControlFlowNode for tar |
 | TarSlipImprov.py:143:36:143:40 | ControlFlowNode for entry | TarSlipImprov.py:141:6:141:29 | ControlFlowNode for Attribute() | TarSlipImprov.py:143:36:143:40 | ControlFlowNode for entry | Extraction of tarfile from $@ to a potentially untrusted source $@. | TarSlipImprov.py:141:6:141:29 | ControlFlowNode for Attribute() | ControlFlowNode for Attribute() | TarSlipImprov.py:143:36:143:40 | ControlFlowNode for entry | ControlFlowNode for entry |
+| TarSlipImprov.py:169:9:169:12 | ControlFlowNode for tarc | TarSlipImprov.py:151:22:151:49 | ControlFlowNode for Attribute() | TarSlipImprov.py:169:9:169:12 | ControlFlowNode for tarc | Extraction of tarfile from $@ to a potentially untrusted source $@. | TarSlipImprov.py:151:22:151:49 | ControlFlowNode for Attribute() | ControlFlowNode for Attribute() | TarSlipImprov.py:169:9:169:12 | ControlFlowNode for tarc | ControlFlowNode for tarc |
 | TarSlipImprov.py:169:9:169:12 | ControlFlowNode for tarc | TarSlipImprov.py:159:26:159:51 | ControlFlowNode for Attribute() | TarSlipImprov.py:169:9:169:12 | ControlFlowNode for tarc | Extraction of tarfile from $@ to a potentially untrusted source $@. | TarSlipImprov.py:159:26:159:51 | ControlFlowNode for Attribute() | ControlFlowNode for Attribute() | TarSlipImprov.py:169:9:169:12 | ControlFlowNode for tarc | ControlFlowNode for tarc |
 | TarSlipImprov.py:178:36:178:40 | ControlFlowNode for entry | TarSlipImprov.py:176:6:176:31 | ControlFlowNode for Attribute() | TarSlipImprov.py:178:36:178:40 | ControlFlowNode for entry | Extraction of tarfile from $@ to a potentially untrusted source $@. | TarSlipImprov.py:176:6:176:31 | ControlFlowNode for Attribute() | ControlFlowNode for Attribute() | TarSlipImprov.py:178:36:178:40 | ControlFlowNode for entry | ControlFlowNode for entry |
 | TarSlipImprov.py:184:21:184:25 | ControlFlowNode for entry | TarSlipImprov.py:182:6:182:31 | ControlFlowNode for Attribute() | TarSlipImprov.py:184:21:184:25 | ControlFlowNode for entry | Extraction of tarfile from $@ to a potentially untrusted source $@. | TarSlipImprov.py:182:6:182:31 | ControlFlowNode for Attribute() | ControlFlowNode for Attribute() | TarSlipImprov.py:184:21:184:25 | ControlFlowNode for entry | ControlFlowNode for entry |