diff --git a/docs/usage.md b/docs/usage.md index 45367e6..cb7d371 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -366,8 +366,10 @@ boolean operators, chained comparisons, keyword and starred arguments, `with`, ` comprehensions, lambdas, nested functions and classes, `global`, `nonlocal`, `del`, assignment expressions (`(y := f(x))`, laid out as an assignment before the statement), `yield from`, starred assignment targets (`first, *rest = items`) and `match` with every -pattern kind. `finally` is modelled on the normal path only and a -`nonlocal` write stays inside the nested function. Methods of module-level classes are +pattern kind. `finally` is modelled on the normal path only. A +`nonlocal` write made by a nested function flows back to the enclosing variable when the +enclosing function calls that nested function directly by name; a nested function passed +around and called elsewhere keeps its writes to itself. Methods of module-level classes are analysed like functions. ## Looking at the intermediate representation diff --git a/src/coretrace_python/cache.py b/src/coretrace_python/cache.py index 17c3833..51b2344 100644 --- a/src/coretrace_python/cache.py +++ b/src/coretrace_python/cache.py @@ -29,6 +29,7 @@ KnownFunction, ModuleGraph, Mutation, + NonlocalWrite, SummaryIndex, Target, UnknownTarget, @@ -36,7 +37,7 @@ from coretrace_python.semantic.symbols import SymbolId from coretrace_python.source import SourceId, SourceSpan -CACHE_FORMAT = 2 +CACHE_FORMAT = 3 @dataclass(frozen=True) @@ -247,6 +248,10 @@ def _encode_summary(summary: FunctionSummary) -> dict[str, Any]: for m in summary.mutations ], "side_effects": sorted(summary.side_effects), + "nonlocal_writes": [ + {"name": w.name, "dependencies": sorted(w.dependencies), "externals": sorted(str(s) for s in w.externals)} + for w in summary.nonlocal_writes + ], } @@ -268,6 +273,12 @@ def _decode_summary(data: Mapping[str, Any]) -> FunctionSummary: for m in data["mutations"] ), frozenset(_string(name) for name in data["side_effects"]), + tuple( + NonlocalWrite( + _string(w["name"]), _indices(w["dependencies"]), frozenset(SymbolId(_string(s)) for s in w["externals"]) + ) + for w in data["nonlocal_writes"] + ), ) diff --git a/src/coretrace_python/interprocedural/__init__.py b/src/coretrace_python/interprocedural/__init__.py index f53b497..d9ef053 100644 --- a/src/coretrace_python/interprocedural/__init__.py +++ b/src/coretrace_python/interprocedural/__init__.py @@ -20,6 +20,7 @@ ExternalCall, FunctionSummary, Mutation, + NonlocalWrite, ProjectSummaries, SummaryAnalysis, SummaryIndex, @@ -36,6 +37,7 @@ "KnownFunction", "ModuleGraph", "Mutation", + "NonlocalWrite", "ProjectSummaries", "SummaryAnalysis", "SummaryIndex", diff --git a/src/coretrace_python/interprocedural/summaries.py b/src/coretrace_python/interprocedural/summaries.py index f70ef84..f2d556c 100644 --- a/src/coretrace_python/interprocedural/summaries.py +++ b/src/coretrace_python/interprocedural/summaries.py @@ -50,10 +50,12 @@ Instruction, Jump, MakeFunction, + NonlocalResult, Phi, Return, SetAttr, SetItem, + SetNonlocal, Symbol, Undefined, Value, @@ -102,6 +104,16 @@ class Mutation: externals: frozenset[SymbolId] +@dataclass(frozen=True) +class NonlocalWrite: + """What the function stores into its ``nonlocal`` ``name``: data of its parameters + and results of external symbols, written back to the caller's variable.""" + + name: str + dependencies: Dependencies + externals: frozenset[SymbolId] + + @dataclass(frozen=True) class FunctionSummary: name: str @@ -112,6 +124,7 @@ class FunctionSummary: return_externals: frozenset[SymbolId] = frozenset() mutations: tuple[Mutation, ...] = () side_effects: frozenset[str] = frozenset() + nonlocal_writes: tuple[NonlocalWrite, ...] = () class SummaryIndex: @@ -184,6 +197,7 @@ def __init__( self.returns: Dep = EMPTY self.stores: dict[HeapLocation, Dep] = {} self.mutated_globals: set[str] = set() + self.nonlocal_stores: dict[str, Dep] = {} # ------------------------------------------------------------------ heap @@ -224,6 +238,9 @@ def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> Stat self.store(state, instruction.object, ATTRIBUTES, state.get(instruction.value, EMPTY)) elif isinstance(instruction, SetItem): self.store(state, instruction.object, ELEMENTS, state.get(instruction.value, EMPTY)) + elif isinstance(instruction, SetNonlocal): + previous = self.nonlocal_stores.get(instruction.name, EMPTY) + self.nonlocal_stores[instruction.name] = previous | state.get(instruction.value, EMPTY) if instruction.result is None: continue if isinstance(instruction, Phi): @@ -233,6 +250,8 @@ def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> Stat deps |= incoming[predecessor].get(value, EMPTY) elif isinstance(instruction, Call): deps = self.call(instruction, state) + elif isinstance(instruction, NonlocalResult): + deps = self.nonlocal_result(instruction, state) else: deps = self.instruction(instruction, state) if isinstance(instruction, GetAttr | GetItem | GetIter): @@ -324,6 +343,34 @@ def call(self, call: Call, state: dict[Key, Dep]) -> Dep: assert isinstance(target, UnknownTarget) return everything | state.get(call.callee, EMPTY) + def nonlocal_result(self, instruction: NonlocalResult, state: dict[Key, Dep]) -> Dep: + """What the nested callee left in its nonlocal ``name``, in the caller's terms.""" + + call = self.defs.get(instruction.call) + if not isinstance(call, Call): + return EMPTY + target = self.graph.target_at(self.name, call.location) + if not isinstance(target, KnownFunction) or target.name not in self.table: + return EMPTY + callee = self.table[target.name] + arguments = tuple(self.deep(a, state) for a in call.arguments) + made = self.defs.get(call.callee) + if isinstance(made, MakeFunction): + arguments = (*arguments, *(self.deep(v, state) for v in made.captured)) + bound = self.receiver(call, target.name) + arguments = (*(self.deep(r, state) for r in bound), *arguments) + keywords = EMPTY + for value in (*call.starred, *(v for _, v in call.keywords)): + keywords |= self.deep(value, state) + deps = EMPTY + for write in callee.nonlocal_writes: + if write.name != instruction.name: + continue + for index in write.dependencies: + deps |= arguments[index] if index < len(arguments) else keywords + deps |= Dep(externals=write.externals) + return deps + def receiver(self, call: Call, name: str) -> tuple[Value, ...]: """The object a method call runs on: the receiver of ``obj.method(...)`` or the new object of ``Class(...)``, the method's implicit first parameter.""" @@ -407,6 +454,7 @@ def summarize( problem.returns = EMPTY problem.stores = {} problem.mutated_globals = set() + problem.nonlocal_stores = {} for block in function.blocks: if solution.reached(block.id): problem.evaluate(block, solution.incoming(block.id)) @@ -426,6 +474,11 @@ def summarize( return_externals=problem.returns.externals, mutations=tuple(mutations), side_effects=frozenset(problem.mutated_globals), + nonlocal_writes=tuple( + NonlocalWrite(name, deps.parameters, deps.externals) + for name, deps in sorted(problem.nonlocal_stores.items()) + if deps.parameters or deps.externals + ), ) diff --git a/src/coretrace_python/ir/lowering.py b/src/coretrace_python/ir/lowering.py index dcfee6d..0d532ab 100644 --- a/src/coretrace_python/ir/lowering.py +++ b/src/coretrace_python/ir/lowering.py @@ -50,11 +50,13 @@ MakeClass, MakeFunction, ModuleIR, + NonlocalResult, Raise, Return, SetAttr, SetGlobal, SetItem, + SetNonlocal, StoreLocal, Symbol, Terminator, @@ -166,7 +168,9 @@ def expression(self, node: nodes.Expression) -> Value: callee = self.expression(node.callee) arguments, starred = self.spread(node.arguments) keywords = tuple((k.name, self.expression(k.value)) for k in node.keywords) - return self.emit(Call(self.new_value(), node.span, callee, arguments, keywords, starred)) + result = self.emit(Call(self.new_value(), node.span, callee, arguments, keywords, starred)) + self.write_back_nonlocals(node, result) + return result if isinstance(node, nodes.BoolOp): values = tuple(self.expression(value) for value in node.values) return self.emit(BoolOp(self.new_value(), node.span, node.operator, values)) @@ -221,6 +225,20 @@ def spread(self, elements: tuple[nodes.Expression, ...]) -> tuple[tuple[Value, . plain.append(self.expression(element)) return tuple(plain), tuple(unpacked) + def write_back_nonlocals(self, node: nodes.Call, call: Value) -> None: + """After a direct call to a nested function defined here, store what it left in + each of its ``nonlocal`` names back into this function's variable.""" + + if not isinstance(node.callee, nodes.Name): + return + nested = self.nested_defs.get(node.callee.identifier) + if nested is None: + return + bindings = self.scopes.scope_for(nested).bindings + for name in sorted(n for n, b in bindings.items() if b.kind is BindingKind.NONLOCAL): + written = self.emit(NonlocalResult(self.new_value(), node.span, call, name)) + self.store(nodes.Name(name, node.span), written) + # ------------------------------------------------------------------ statements def store(self, target: nodes.Target, value: Value) -> None: @@ -232,6 +250,8 @@ def store(self, target: nodes.Target, value: Value) -> None: if resolution.kind is not ResolutionKind.LOCAL and target.identifier not in self.nonlocals: self.fail(target, "assignment to a free variable that is not declared nonlocal") self.emit_effect(StoreLocal(None, target.span, target.identifier, value)) + if target.identifier in self.nonlocals: + self.emit_effect(SetNonlocal(None, target.span, target.identifier, value)) elif isinstance(target, nodes.Attribute): object_value = self.expression(target.value) self.emit_effect(SetAttr(None, target.span, object_value, target.name, value)) @@ -300,6 +320,7 @@ def statement(self, node: nodes.Statement) -> None: captured = tuple(self.captured_value(name, node.span) for name in captured_names(node, self.scopes)) made = self.emit(MakeFunction(self.new_value(), node.span, node.name, captured)) self.store(nodes.Name(node.name, node.span), made) + self.nested_defs[node.name] = node return if isinstance(node, nodes.AugAssign): current = self.expression(node.target) @@ -405,6 +426,9 @@ def function(self, node: nodes.Function) -> FunctionIR: # assigned in this body, so it lives in a local slot initialised from the # captured value; the others are never reassigned. bindings = self.scopes.scope_for(node).bindings + # Nested definitions bound in this body, so a direct call to one can write its + # ``nonlocal`` names back into this function's variables. + self.nested_defs: dict[str, nodes.Function] = {} self.nonlocals = frozenset( name for name in captured if name in bindings and bindings[name].kind is BindingKind.NONLOCAL ) @@ -429,7 +453,11 @@ def function(self, node: nodes.Function) -> FunctionIR: ) ) return FunctionIR( - qualified_name(self.scopes, node), parameter_values, self.cfg.entry, tuple(blocks), node.span + qualified_name(self.scopes, node), + parameter_values, + self.cfg.entry, + tuple(blocks), + node.span, ) diff --git a/src/coretrace_python/ir/model.py b/src/coretrace_python/ir/model.py index 69c66c3..7dddd02 100644 --- a/src/coretrace_python/ir/model.py +++ b/src/coretrace_python/ir/model.py @@ -234,6 +234,17 @@ class Yield(Operands): value: Value | None +@dataclass(frozen=True) +class NonlocalResult(Operands): + """The value a nested function left in its ``nonlocal`` ``name`` when ``call`` + returned; the caller stores it back into its own variable.""" + + result: Value + location: SourceSpan + call: Value + name: str + + @dataclass(frozen=True) class GetAttr(Operands): result: Value @@ -340,6 +351,17 @@ class DelAttr(Operands): attribute: str +@dataclass(frozen=True) +class SetNonlocal(Operands): + """A store to a name this function declares ``nonlocal``, kept through SSA so the + summary can tell the enclosing function what was written back.""" + + result: None + location: SourceSpan + name: str + value: Value + + @dataclass(frozen=True) class SetGlobal(Operands): """Assignment to a name declared ``global``.""" @@ -389,9 +411,11 @@ class Import(Operands): | Catch | Await | Yield + | NonlocalResult ) EffectInstruction: TypeAlias = ( StoreLocal | SetAttr | SetItem | WithExit | Assert | Import | SetGlobal | DelItem | DelAttr + | SetNonlocal ) Instruction: TypeAlias = ValueInstruction | EffectInstruction diff --git a/src/coretrace_python/ir/printer.py b/src/coretrace_python/ir/printer.py index f1e6031..ccb2793 100644 --- a/src/coretrace_python/ir/printer.py +++ b/src/coretrace_python/ir/printer.py @@ -30,12 +30,14 @@ MakeClass, MakeFunction, ModuleIR, + NonlocalResult, Phi, Raise, Return, SetAttr, SetGlobal, SetItem, + SetNonlocal, StoreLocal, Symbol, Terminator, @@ -107,6 +109,8 @@ def _instruction(instruction: Instruction) -> str: return f'{_value(instruction.result)} = make_class "{instruction.name}"' if isinstance(instruction, SetGlobal): return f'set_global "{instruction.name}", {_value(instruction.value)}' + if isinstance(instruction, SetNonlocal): + return f'set_nonlocal "{instruction.name}", {_value(instruction.value)}' if isinstance(instruction, BuildString): return f"{_value(instruction.result)} = build_string {', '.join(_value(v) for v in instruction.parts)}".rstrip() if isinstance(instruction, BuildSlice): @@ -124,6 +128,8 @@ def _instruction(instruction: Instruction) -> str: if isinstance(instruction, Yield): suffix = "" if instruction.value is None else f" {_value(instruction.value)}" return f"{_value(instruction.result)} = yield{suffix}" + if isinstance(instruction, NonlocalResult): + return f"{_value(instruction.result)} = nonlocal_result {_value(instruction.call)}, {instruction.name!r}" if isinstance(instruction, SetAttr): return f"set_attr {_value(instruction.object)}, {instruction.attribute!r}, {_value(instruction.value)}" if isinstance(instruction, SetItem): diff --git a/src/coretrace_python/taint/engine.py b/src/coretrace_python/taint/engine.py index d77b873..b32ad22 100644 --- a/src/coretrace_python/taint/engine.py +++ b/src/coretrace_python/taint/engine.py @@ -57,6 +57,7 @@ Instruction, Jump, MakeFunction, + NonlocalResult, Phi, SetAttr, SetItem, @@ -227,6 +228,8 @@ def evaluate( taint = taint.join(incoming[predecessor].get(value, Taint.none())) elif isinstance(instruction, Call): taint = self.call(instruction, state, flows) + elif isinstance(instruction, NonlocalResult): + taint = self.nonlocal_result(instruction, state) else: taint = self.instruction(instruction, state) if isinstance(instruction, GetAttr | GetItem | GetIter): @@ -317,6 +320,36 @@ def call(self, call: Call, state: dict[Key, Taint], flows: list[TaintFlow]) -> T return self.external(returned, everything, call, state, flows) return everything.join(state.get(call.callee, Taint.none())) + def nonlocal_result(self, instruction: NonlocalResult, state: dict[Key, Taint]) -> Taint: + """The taint a nested callee left in its nonlocal ``name``, in the caller's terms.""" + + call = self.defs.get(instruction.call) + if not isinstance(call, Call): + return Taint.none() + target = self.graph.target_at(self.name, call.location) + if not isinstance(target, KnownFunction): + return Taint.none() + summary = self.summaries.summary(target.name) + arguments = ( + *(self.deep(r, state) for r in self.receiver(call, target.name)), + *(self.deep(a, state) for a in call.arguments), + *(self.deep(v, state) for v in self.captured(call)), + ) + keywords = Taint.none() + for value in (*call.starred, *(v for _, v in call.keywords)): + keywords = keywords.join(self.deep(value, state)) + result = Taint.none() + for write in summary.nonlocal_writes: + if write.name != instruction.name: + continue + for index in sorted(write.dependencies): + result = result.join(arguments[index] if index < len(arguments) else keywords) + for symbol in sorted(write.externals, key=str): + source = self.models.source(symbol) + if source is not None: + result = result.join(Taint(source.kinds, frozenset({source}))) + return result + def receiver(self, call: Call, name: str) -> tuple[Value, ...]: """The object a method call runs on, its implicit first parameter: the receiver of ``obj.method(...)``, or the new object of ``Class(...)``.""" diff --git a/tests/test_nonlocal_writeback.py b/tests/test_nonlocal_writeback.py new file mode 100644 index 0000000..32ccc1b --- /dev/null +++ b/tests/test_nonlocal_writeback.py @@ -0,0 +1,103 @@ +"""Acceptance tests for ``nonlocal`` writes flowing back to the enclosing function (issue #71). + +A nested function that declares ``nonlocal name`` and assigns it rebinds the enclosing +function's variable when it is called. The lowering of the enclosing function emits, after +a call to such a nested function it defined, a ``nonlocal_result`` per written name and +stores it into the variable; the nested function's summary records what it stores into +each nonlocal, and the dependence and taint engines read that back at the call. + +Expected to remain red until ``ir.model.NonlocalResult`` and ``FunctionSummary.nonlocal_writes`` exist. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from coretrace_python import engine +from coretrace_python.findings import Finding +from coretrace_python.frontend import build_hir +from coretrace_python.interprocedural import FunctionSummary +from coretrace_python.ir import model +from coretrace_python.ir.lowering import lower_module +from coretrace_python.ir.printer import format_module +from coretrace_python.source import SourceManager + +MISSING = None if hasattr(model, "NonlocalResult") else "NonlocalResult is missing" + + +@pytest.fixture(autouse=True) +def require_writeback() -> None: + if MISSING is not None: + pytest.fail(f"nonlocal write-back is not implemented yet: {MISSING}") + + +REPO = Path(__file__).resolve().parent.parent +PLUGINS = REPO / "src" / "coretrace_python" / "bundled" + + +def check(text: str) -> tuple[Finding, ...]: + return engine.check(SourceManager().add_source("m.py", text), [PLUGINS]) + + +def rules(findings: tuple[Finding, ...]) -> list[tuple[str, int, str | None]]: + return sorted((f.rule_id, f.span.start_line, f.function) for f in findings) + + +def printed(text: str) -> str: + return format_module(lower_module(build_hir(SourceManager().add_source("m.py", text)))) + + +OUTSIDE = "import os\n\ndef outer():\n cmd = 'ls'\n def set_cmd():\n nonlocal cmd\n cmd = input()\n set_cmd()\n os.system(cmd)\n" + + +def test_the_call_rebinds_the_written_nonlocal_in_the_enclosing_function() -> None: + outer = printed(OUTSIDE).split("func @outer.set_cmd")[0] + + assert "nonlocal_result" in outer + assert outer.index("= call") < outer.index("nonlocal_result") < outer.index('store_local "cmd"', outer.index("= call")) + + +def test_a_nonlocal_write_flows_back_to_the_caller() -> None: + assert rules(check(OUTSIDE)) == [("command-injection", 9, "outer")] + + +def test_a_nonlocal_write_from_a_parameter_flows_back() -> None: + text = ( + "import os\n\ndef outer():\n cmd = 'ls'\n def set_cmd(value):\n nonlocal cmd\n" + " cmd = value\n set_cmd(input())\n os.system(cmd)\n" + ) + assert rules(check(text)) == [("command-injection", 9, "outer")] + + +def test_a_nonlocal_written_with_a_constant_does_not_taint_the_caller() -> None: + text = ( + "import os\n\ndef outer():\n cmd = input()\n def reset():\n nonlocal cmd\n" + " cmd = 'ls'\n reset()\n os.system(cmd)\n" + ) + assert check(text) == () + + +def test_unwritten_captures_and_other_locals_are_untouched() -> None: + text = ( + "import os\n\ndef outer():\n cmd = 'ls'\n other = 'pwd'\n def set_cmd():\n nonlocal cmd\n" + " cmd = input()\n os.system(other)\n set_cmd()\n os.system(other)\n" + ) + assert check(text) == () + + +def test_summaries_record_nonlocal_writes_and_survive_the_cache(tmp_path: Path) -> None: + from coretrace_python.cache import ProjectCache + + (tmp_path / "m.py").write_text(OUTSIDE, encoding="utf-8") + cache = ProjectCache(tmp_path / "cache") + first = engine.analyze_project(tmp_path, [PLUGINS], cache=cache) + second = engine.analyze_project(tmp_path, [PLUGINS], cache=cache) + + assert second.reused == ("m",) + assert rules(second.findings) == rules(first.findings) == [("command-injection", 9, "outer")] + entries = [json.loads(p.read_text(encoding="utf-8")) for p in (tmp_path / "cache").glob("*.json")] + assert any("nonlocal_writes" in json.dumps(entry) for entry in entries) + assert "nonlocal_writes" in FunctionSummary.__dataclass_fields__ diff --git a/tests/test_remaining_limits.py b/tests/test_remaining_limits.py index 9a657f5..de6c303 100644 --- a/tests/test_remaining_limits.py +++ b/tests/test_remaining_limits.py @@ -116,9 +116,9 @@ def test_nonlocal_assignments_lower_as_locals_of_the_nested_function() -> None: def test_a_nonlocal_write_is_visible_inside_the_writer_but_not_yet_outside() -> None: inside = "import os\n\ndef outer():\n cmd = 'ls'\n def set_cmd():\n nonlocal cmd\n cmd = input()\n os.system(cmd)\n set_cmd()\n" assert rules(check(inside)) == [("command-injection", 8, "set_cmd")] - # Documented limit: the write does not flow back to the enclosing function. + # The write flows back to the enclosing function when it calls the writer directly. outside = "import os\n\ndef outer():\n cmd = 'ls'\n def set_cmd():\n nonlocal cmd\n cmd = input()\n set_cmd()\n os.system(cmd)\n" - assert check(outside) == () + assert rules(check(outside)) == [("command-injection", 9, "outer")] # --------------------------------------------------------------------------- classes inside functions