diff --git a/python/ql/lib/change-notes/2023-10-17-contextmanager.md b/python/ql/lib/change-notes/2023-10-17-contextmanager.md new file mode 100644 index 00000000000..dabd03a7257 --- /dev/null +++ b/python/ql/lib/change-notes/2023-10-17-contextmanager.md @@ -0,0 +1,4 @@ +--- +category: minorAnalysis +--- +* Added support for functions decorated with `contextlib.contextmanager`. diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll index 78447016913..123fa862c8e 100644 --- a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll +++ b/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`. */ diff --git a/python/ql/test/experimental/dataflow/typetracking/test.py b/python/ql/test/experimental/dataflow/typetracking/test.py index a822ef683eb..a54daa1f046 100644 --- a/python/ql/test/experimental/dataflow/typetracking/test.py +++ b/python/ql/test/experimental/dataflow/typetracking/test.py @@ -195,3 +195,42 @@ class MyClass(object): 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 diff --git a/python/ql/test/experimental/query-tests/Security/CWE-022-TarSlip/TarSlip.expected b/python/ql/test/experimental/query-tests/Security/CWE-022-TarSlip/TarSlip.expected index a699a1bac5f..75e8b40b929 100644 --- a/python/ql/test/experimental/query-tests/Security/CWE-022-TarSlip/TarSlip.expected +++ b/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 |