diff --git a/README.md b/README.md index f5da2c5..e22b9fd 100644 --- a/README.md +++ b/README.md @@ -64,8 +64,9 @@ PYTHONPATH=src python -m coretrace_python --emit-ir example.py Currently supported inside functions: parameters with defaults and keyword-only or star forms, decorators, assignments to names, attributes, items and unpacked tuples, augmented assignment, list, tuple and dict literals, `and`/`or`, chained comparisons, keyword -arguments, `with`, `assert`, `if`/`elif`/`else`, `while`, `for`, `break`, `continue` and -`raise`. Methods of module-level classes are analysed like functions; other module-level +arguments, `with`, `assert`, `try`/`except`/`else`/`finally`, `await`, `yield`, +`if`/`elif`/`else`, `while`, `for`, `break`, `continue` and `raise`. Blocks inside a `try` +body carry exception edges to the handlers; `finally` is modelled on the normal path only. Methods of module-level classes are analysed like functions; other module-level code is skipped. A function using syntax outside this subset is reported by `--check` as an `unsupported-syntax` note and the other functions are still analysed. Each function is emitted as its control-flow graph: one block per basic block, ending in an explicit diff --git a/src/coretrace_python/abstract/constants.py b/src/coretrace_python/abstract/constants.py index 4ffe733..223bd3f 100644 --- a/src/coretrace_python/abstract/constants.py +++ b/src/coretrace_python/abstract/constants.py @@ -113,6 +113,7 @@ def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> Stat def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) -> Mapping[BlockId, State]: block = self.blocks[block_id] state = self.evaluate(block, incoming) + exits = {target: state for target in block.exception_targets} terminator = block.terminator if isinstance(terminator, Branch): truth = state[terminator.condition].truthiness @@ -121,12 +122,12 @@ def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) - Truth.FALSE: (terminator.else_block,), Truth.UNKNOWN: (terminator.then_block, terminator.else_block), }[truth] - return {target: state for target in targets} + return {**exits, **{target: state for target in targets}} if isinstance(terminator, Jump): - return {terminator.target: state} + return {**exits, terminator.target: state} if isinstance(terminator, ForNext): - return {terminator.body: state, terminator.exit: state} - return {} + return {**exits, terminator.body: state, terminator.exit: state} + return exits # ------------------------------------------------------------------ transfer diff --git a/src/coretrace_python/cfg/builder.py b/src/coretrace_python/cfg/builder.py index fade2b9..274c069 100644 --- a/src/coretrace_python/cfg/builder.py +++ b/src/coretrace_python/cfg/builder.py @@ -35,6 +35,7 @@ def __init__(self, function: nodes.Function) -> None: self.blocks: dict[BlockId, BasicBlock] = {} self.counters: dict[str, int] = {} self.loops: list[tuple[BlockId, BlockId]] = [] + self.handlers: list[tuple[BlockId, ...]] = [] def build(self) -> CFG: entry = BlockId("entry") @@ -50,10 +51,15 @@ def new_id(self, kind: str) -> BlockId: return BlockId(f"{kind}_{self.counters[kind]}") def finish(self, block: _Open, terminator: Terminator) -> None: - self.blocks[block.id] = BasicBlock(block.id, tuple(block.statements), terminator) + self.blocks[block.id] = BasicBlock( + block.id, tuple(block.statements), terminator, self.exception_targets() + ) def header(self, block_id: BlockId, terminator: Terminator) -> None: - self.blocks[block_id] = BasicBlock(block_id, (), terminator) + self.blocks[block_id] = BasicBlock(block_id, (), terminator, self.exception_targets()) + + def exception_targets(self) -> tuple[BlockId, ...]: + return self.handlers[-1] if self.handlers else () # ------------------------------------------------------------------ statements @@ -104,9 +110,49 @@ def statement( return self.loop_statement(node, block, continuation) if isinstance(node, nodes.With): return self.with_statement(node, block) + if isinstance(node, nodes.Try): + return self.try_statement(node, block, continuation) block.statements.append(node) return block + def try_statement( + self, node: nodes.Try, block: _Open, continuation: BlockId | None + ) -> _Open | None: + """Body blocks carry exception edges to every handler; handlers, ``else`` and + ``finally`` join on the normal path (exceptional exits are approximated).""" + + body_id = self.new_id("try") + handler_ids = tuple(self.new_id("handler") for _ in node.handlers) + else_id = self.new_id("else") if node.orelse else None + final_id = self.new_id("finally") if node.finalbody else None + after_id = continuation if continuation is not None and final_id is None else None + if after_id is None: + after_id = self.new_id("after") if final_id is None or continuation is None else None + join = final_id if final_id is not None else after_id + assert join is not None + + self.finish(block, Jump(body_id, node.span)) + self.handlers.append(handler_ids) + try: + # The whole body, including the block that leaves it, may raise into a handler. + end = self.sequence(node.body, _Open(body_id), None, node.span) + if end is not None: + self.finish(end, Jump(else_id if else_id is not None else join, node.span)) + finally: + self.handlers.pop() + if end is not None and else_id is not None: + self.sequence(node.orelse, _Open(else_id), join, node.span) + for handler, handler_id in zip(node.handlers, handler_ids, strict=True): + opened = _Open(handler_id) + opened.statements.append(nodes.EnterHandler(handler, handler.span)) + self.sequence(handler.body, opened, join, node.span) + if final_id is not None: + target = continuation if continuation is not None else after_id + assert target is not None + self.sequence(node.finalbody, _Open(final_id), target, node.span) + return None if continuation is not None else _Open(target) + return None if continuation is not None else _Open(after_id) # type: ignore[arg-type] + def with_statement(self, node: nodes.With, block: _Open) -> _Open | None: """Lay the body out inline; an early exit skips the ``ExitWith`` statements.""" diff --git a/src/coretrace_python/cfg/model.py b/src/coretrace_python/cfg/model.py index 6296f81..573d237 100644 --- a/src/coretrace_python/cfg/model.py +++ b/src/coretrace_python/cfg/model.py @@ -83,6 +83,8 @@ class BasicBlock: id: BlockId statements: tuple[nodes.Statement, ...] terminator: Terminator + exception_targets: tuple[BlockId, ...] = () + """Handler blocks control may reach if a statement of this block raises.""" @dataclass(frozen=True) @@ -105,7 +107,9 @@ def __post_init__(self) -> None: for block_id, block in blocks.items(): if block.id != block_id: raise CFGError(f"block {block.id} is stored under {block_id}") - successors[block_id] = targets(block.terminator) + found = list(targets(block.terminator)) + found.extend(t for t in block.exception_targets if t not in found) + successors[block_id] = tuple(found) for target in successors[block_id]: if target not in blocks: raise CFGError(f"block {block_id} targets unknown block {target}") diff --git a/src/coretrace_python/frontend/ast_adapter.py b/src/coretrace_python/frontend/ast_adapter.py index 66c3577..6a9e08b 100644 --- a/src/coretrace_python/frontend/ast_adapter.py +++ b/src/coretrace_python/frontend/ast_adapter.py @@ -123,6 +123,13 @@ def expression(self, node: ast.expr) -> nodes.Expression: for keyword in node.keywords ) return nodes.Call(self.expression(node.func), arguments, keywords, span) + if isinstance(node, ast.Await): + return nodes.Await(self.expression(node.value), span) + if isinstance(node, ast.Yield): + yielded = self.expression(node.value) if node.value is not None else None + return nodes.Yield(yielded, span) + if isinstance(node, ast.YieldFrom): + self.fail(node, "yield from is not supported yet") if isinstance(node, (ast.ListComp, ast.SetComp, ast.GeneratorExp)): generators = tuple(self.generator(generator) for generator in node.generators) return nodes.Comprehension( @@ -248,6 +255,24 @@ def statement(self, node: ast.stmt) -> nodes.Statement: self.fail(node.cause, "raise from is not supported yet") exception = self.expression(node.exc) if node.exc is not None else None return nodes.Raise(exception, span) + if isinstance(node, ast.Try): + handlers = [] + for handler in node.handlers: + handler_type = self.expression(handler.type) if handler.type is not None else None + handlers.append( + nodes.ExceptHandler( + handler_type, handler.name, self.block(handler.body), self.span(handler) + ) + ) + return nodes.Try( + self.block(node.body), + tuple(handlers), + self.block(node.orelse), + self.block(node.finalbody), + span, + ) + if isinstance(node, getattr(ast, "TryStar", ())): + self.fail(node, "except* groups are not supported yet") if isinstance(node, ast.Global): return nodes.Global(tuple(node.names), span) if isinstance(node, ast.Nonlocal): diff --git a/src/coretrace_python/hir/nodes.py b/src/coretrace_python/hir/nodes.py index 651d15f..9c7911c 100644 --- a/src/coretrace_python/hir/nodes.py +++ b/src/coretrace_python/hir/nodes.py @@ -103,6 +103,18 @@ class Dict: span: SourceSpan +@dataclass(frozen=True, slots=True) +class Await: + value: Expression + span: SourceSpan + + +@dataclass(frozen=True, slots=True) +class Yield: + value: Expression | None + span: SourceSpan + + @dataclass(frozen=True, slots=True) class ComprehensionGenerator: """One ``for target in iterable if condition...`` clause of a comprehension.""" @@ -136,6 +148,8 @@ class Comprehension: | Tuple | List | Dict + | Await + | Yield | Comprehension ) @@ -262,6 +276,33 @@ class ExitWith: span: SourceSpan +@dataclass(frozen=True, slots=True) +class ExceptHandler: + """``except type as name:``; ``type`` is ``None`` for a bare ``except``.""" + + type: Expression | None + name: str | None + body: tuple[Statement, ...] + span: SourceSpan + + +@dataclass(frozen=True, slots=True) +class Try: + body: tuple[Statement, ...] + handlers: tuple[ExceptHandler, ...] + orelse: tuple[Statement, ...] + finalbody: tuple[Statement, ...] + span: SourceSpan + + +@dataclass(frozen=True, slots=True) +class EnterHandler: + """Synthetic statement the CFG builder emits at the start of a handler block.""" + + handler: ExceptHandler + span: SourceSpan + + @dataclass(frozen=True, slots=True) class ImportAlias: name: str @@ -330,6 +371,8 @@ class Class: | With | EnterWith | ExitWith + | Try + | EnterHandler | Import | ImportFrom | Global diff --git a/src/coretrace_python/interprocedural/modulegraph.py b/src/coretrace_python/interprocedural/modulegraph.py index 3eb7a6a..0a49801 100644 --- a/src/coretrace_python/interprocedural/modulegraph.py +++ b/src/coretrace_python/interprocedural/modulegraph.py @@ -99,10 +99,13 @@ def _prefixes(dotted: str) -> list[str]: def _statements(body: Iterable[nodes.Statement]) -> Iterable[nodes.Statement]: for statement in body: yield statement - for attribute in ("body", "orelse"): + for attribute in ("body", "orelse", "finalbody"): nested = getattr(statement, attribute, None) if isinstance(nested, tuple): yield from _statements(nested) + if isinstance(statement, nodes.Try): + for handler in statement.handlers: + yield from _statements(handler.body) __all__ = [ "IGNORED_DIRECTORIES", diff --git a/src/coretrace_python/interprocedural/summaries.py b/src/coretrace_python/interprocedural/summaries.py index 05bc978..af4d1f8 100644 --- a/src/coretrace_python/interprocedural/summaries.py +++ b/src/coretrace_python/interprocedural/summaries.py @@ -185,14 +185,15 @@ def evaluate(self, block: BasicBlock, incoming: Mapping[BlockId, State]) -> Stat def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) -> Mapping[BlockId, State]: block = self.blocks[block_id] state = self.evaluate(block, incoming) + exits = {target: state for target in block.exception_targets} terminator = block.terminator if isinstance(terminator, Branch): - return {terminator.then_block: state, terminator.else_block: state} + return {**exits, terminator.then_block: state, terminator.else_block: state} if isinstance(terminator, Jump): - return {terminator.target: state} + return {**exits, terminator.target: state} if isinstance(terminator, ForNext): - return {terminator.body: state, terminator.exit: state} - return {} + return {**exits, terminator.body: state, terminator.exit: state} + return exits # ------------------------------------------------------------------ transfer diff --git a/src/coretrace_python/ir/lowering.py b/src/coretrace_python/ir/lowering.py index 3714523..a050e28 100644 --- a/src/coretrace_python/ir/lowering.py +++ b/src/coretrace_python/ir/lowering.py @@ -18,6 +18,7 @@ from coretrace_python.hir.visitors import Node, children from coretrace_python.ir.model import ( Assert, + Await, BasicBlock, BinaryOp, BoolOp, @@ -26,6 +27,7 @@ BuildList, BuildTuple, Call, + Catch, Compare, Constant, EffectInstruction, @@ -51,6 +53,7 @@ ValueInstruction, WithEnter, WithExit, + Yield, ) from coretrace_python.semantic import SEMANTIC_ANALYSES from coretrace_python.semantic.scopes import ( @@ -157,6 +160,11 @@ def expression(self, node: nodes.Expression) -> Value: if isinstance(node, nodes.Dict): items = tuple((self.expression(k), self.expression(v)) for k, v in node.items) return self.emit(BuildDict(self.new_value(), node.span, items)) + if isinstance(node, nodes.Await): + return self.emit(Await(self.new_value(), node.span, self.expression(node.value))) + if isinstance(node, nodes.Yield): + value = self.expression(node.value) if node.value is not None else None + return self.emit(Yield(self.new_value(), node.span, value)) if isinstance(node, nodes.Attribute): object_value = self.expression(node.value) return self.emit(GetAttr(self.new_value(), node.span, object_value, node.name)) @@ -212,6 +220,13 @@ def statement(self, node: nodes.Statement) -> None: if isinstance(node, nodes.ExitWith): self.emit_effect(WithExit(None, node.item.span, self.contexts[node.item.span])) return + if isinstance(node, nodes.EnterHandler): + handler = node.handler + caught_type = self.expression(handler.type) if handler.type is not None else None + caught = self.emit(Catch(self.new_value(), handler.span, caught_type)) + if handler.name is not None: + self.store(nodes.Name(handler.name, handler.span), caught) + return if isinstance(node, nodes.ExpressionStatement): self.expression(node.expression) return @@ -276,7 +291,11 @@ def function(self, node: nodes.Function) -> FunctionIR: for statement in cfg_block.statements: self.statement(statement) terminator = self.terminator(cfg_block) - blocks.append(BasicBlock(cfg_block.id, tuple(self.instructions), terminator)) + blocks.append( + BasicBlock( + cfg_block.id, tuple(self.instructions), terminator, cfg_block.exception_targets + ) + ) return FunctionIR( qualified_name(self.scopes, node), parameter_values, self.cfg.entry, tuple(blocks), node.span ) @@ -309,6 +328,8 @@ def names(target: nodes.Target | None) -> None: def walk(node: Node) -> None: if isinstance(node, nodes.Assign | nodes.AugAssign | nodes.For | nodes.WithItem): names(node.target) + elif isinstance(node, nodes.ExceptHandler) and node.name is not None: + assigned.add(node.name) if isinstance(node, nodes.Function | nodes.Class | nodes.Comprehension): return for child in children(node): diff --git a/src/coretrace_python/ir/model.py b/src/coretrace_python/ir/model.py index 85ccbfb..d3f74c0 100644 --- a/src/coretrace_python/ir/model.py +++ b/src/coretrace_python/ir/model.py @@ -153,6 +153,29 @@ class WithEnter(Operands): context: Value +@dataclass(frozen=True) +class Catch(Operands): + """The exception bound by ``except type as name`` at the start of a handler.""" + + result: Value + location: SourceSpan + type: Value | None + + +@dataclass(frozen=True) +class Await(Operands): + result: Value + location: SourceSpan + value: Value + + +@dataclass(frozen=True) +class Yield(Operands): + result: Value + location: SourceSpan + value: Value | None + + @dataclass(frozen=True) class GetAttr(Operands): result: Value @@ -262,6 +285,9 @@ class Assert(Operands): | BuildTuple | BuildDict | WithEnter + | Catch + | Await + | Yield ) EffectInstruction: TypeAlias = StoreLocal | SetAttr | SetItem | WithExit | Assert Instruction: TypeAlias = ValueInstruction | EffectInstruction @@ -316,6 +342,25 @@ class BasicBlock: id: BlockId instructions: tuple[Instruction, ...] terminator: Terminator + exception_targets: tuple[BlockId, ...] = () + + +def successors(block: BasicBlock) -> tuple[BlockId, ...]: + """Terminator targets followed by the exception edges of ``block``.""" + + found = list(_terminator_targets(block.terminator)) + found.extend(t for t in block.exception_targets if t not in found) + return tuple(found) + + +def _terminator_targets(terminator: Terminator) -> tuple[BlockId, ...]: + if isinstance(terminator, Branch): + return (terminator.then_block, terminator.else_block) + if isinstance(terminator, Jump): + return (terminator.target,) + if isinstance(terminator, ForNext): + return (terminator.body, terminator.exit) + return () @dataclass(frozen=True) diff --git a/src/coretrace_python/ir/printer.py b/src/coretrace_python/ir/printer.py index a302ddc..d1833ba 100644 --- a/src/coretrace_python/ir/printer.py +++ b/src/coretrace_python/ir/printer.py @@ -2,6 +2,7 @@ from coretrace_python.ir.model import ( Assert, + Await, BinaryOp, BoolOp, Branch, @@ -9,6 +10,7 @@ BuildList, BuildTuple, Call, + Catch, Compare, Constant, ForNext, @@ -33,6 +35,7 @@ Value, WithEnter, WithExit, + Yield, ) @@ -81,6 +84,14 @@ def _instruction(instruction: Instruction) -> str: return f"{_value(instruction.result)} = with_enter {_value(instruction.context)}" if isinstance(instruction, WithExit): return f"with_exit {_value(instruction.context)}" + if isinstance(instruction, Catch): + suffix = "" if instruction.type is None else f" {_value(instruction.type)}" + return f"{_value(instruction.result)} = catch{suffix}" + if isinstance(instruction, Await): + return f"{_value(instruction.result)} = await {_value(instruction.value)}" + 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, SetAttr): return f"set_attr {_value(instruction.object)}, {instruction.attribute!r}, {_value(instruction.value)}" if isinstance(instruction, SetItem): @@ -142,7 +153,10 @@ def format_module(module: ModuleIR) -> str: parameters = ", ".join(_value(parameter) for parameter in function.parameters) lines = [f"func @{function.name}({parameters}) {{"] for block in function.blocks: - lines.append(f"{block.id}:") + header = str(block.id) + if block.exception_targets: + header += f" [except: {', '.join(str(t) for t in block.exception_targets)}]" + lines.append(f"{header}:") lines.extend(f" {_instruction(instruction)}" for instruction in block.instructions) lines.append(f" {_terminator(block.terminator)}") lines.append("}") diff --git a/src/coretrace_python/ir/ssa.py b/src/coretrace_python/ir/ssa.py index dd0a46c..a449fad 100644 --- a/src/coretrace_python/ir/ssa.py +++ b/src/coretrace_python/ir/ssa.py @@ -33,6 +33,7 @@ Undefined, Value, substitute, + successors, ) @@ -50,7 +51,7 @@ def __init__(self, function: FunctionIR, tree: DominatorTree) -> None: self.terminators: dict[BlockId, Terminator] = {} self.predecessors: dict[BlockId, list[BlockId]] = {b: [] for b in self.blocks} for block in self.blocks.values(): - for successor in _targets(block.terminator): + for successor in successors(block): self.predecessors[successor].append(block.id) self.loop_variables: dict[BlockId, tuple[str, BlockId]] = {} for block in self.blocks.values(): @@ -128,7 +129,7 @@ def push(name: str, value: Value) -> None: terminator = replace(terminator, result=self.new_value()) self.terminators[block_id] = terminator - for successor in _targets(terminator): + for successor in (*_targets(terminator), *self.blocks[block_id].exception_targets): for name, phi in self.phis.get(successor, {}).items(): incoming = (*phi.incoming, (self.current(name), block_id)) self.phis[successor][name] = replace(phi, incoming=incoming) @@ -152,7 +153,11 @@ def build(self) -> FunctionIR: instructions.extend(self.undefined.values()) instructions.extend(self.ordered_phis(block.id)) instructions.extend(self.instructions[block.id]) - blocks.append(BasicBlock(block.id, tuple(instructions), self.terminators[block.id])) + blocks.append( + BasicBlock( + block.id, tuple(instructions), self.terminators[block.id], block.exception_targets + ) + ) return _renumber(replace(self.function, blocks=tuple(blocks))) def ordered_phis(self, block_id: BlockId) -> list[Phi]: @@ -248,6 +253,7 @@ def renamed(value: Value) -> Value: block.id, tuple(substitute(i, renamed, include_result=True) for i in block.instructions), substitute(block.terminator, renamed, include_result=True), + block.exception_targets, ) for block in function.blocks ) diff --git a/src/coretrace_python/semantic/imports.py b/src/coretrace_python/semantic/imports.py index 5d725e7..aababae 100644 --- a/src/coretrace_python/semantic/imports.py +++ b/src/coretrace_python/semantic/imports.py @@ -87,6 +87,12 @@ def body(self, statements: tuple[nodes.Statement, ...], scope_id: ScopeId) -> No self.body(statement.orelse, scope_id) elif isinstance(statement, (nodes.While, nodes.For, nodes.With)): self.body(statement.body, scope_id) + elif isinstance(statement, nodes.Try): + self.body(statement.body, scope_id) + for handler in statement.handlers: + self.body(handler.body, scope_id) + self.body(statement.orelse, scope_id) + self.body(statement.finalbody, scope_id) def base_module(self, statement: nodes.ImportFrom) -> str: if not statement.level: diff --git a/src/coretrace_python/semantic/scopes.py b/src/coretrace_python/semantic/scopes.py index 2d2cb03..da3bc1f 100644 --- a/src/coretrace_python/semantic/scopes.py +++ b/src/coretrace_python/semantic/scopes.py @@ -239,7 +239,17 @@ def statement(self, node: nodes.Statement, scope: _ScopeBuilder) -> None: if item.target is not None: scope.bind(item.target.identifier, BindingKind.LOCAL, item.target.span) self.body(node.body, scope) - elif isinstance(node, nodes.EnterWith | nodes.ExitWith): + elif isinstance(node, nodes.Try): + self.body(node.body, scope) + for handler in node.handlers: + if handler.type is not None: + self.expression(handler.type, scope) + if handler.name is not None: + scope.bind(handler.name, BindingKind.LOCAL, handler.span) + self.body(handler.body, scope) + self.body(node.orelse, scope) + self.body(node.finalbody, scope) + elif isinstance(node, nodes.EnterWith | nodes.ExitWith | nodes.EnterHandler): pass elif isinstance(node, nodes.Return): if node.value is not None: @@ -324,6 +334,11 @@ def expression(self, node: nodes.Expression, scope: _ScopeBuilder) -> None: elif isinstance(node, nodes.BoolOp): for value in node.values: self.expression(value, scope) + elif isinstance(node, nodes.Await): + self.expression(node.value, scope) + elif isinstance(node, nodes.Yield): + if node.value is not None: + self.expression(node.value, scope) elif isinstance(node, nodes.Tuple | nodes.List): for element in node.elements: self.expression(element, scope) diff --git a/src/coretrace_python/taint/engine.py b/src/coretrace_python/taint/engine.py index 9b525c1..4bb6fbe 100644 --- a/src/coretrace_python/taint/engine.py +++ b/src/coretrace_python/taint/engine.py @@ -167,14 +167,15 @@ def evaluate( def flow(self, cfg: CFG, block_id: BlockId, incoming: Mapping[BlockId, State]) -> Mapping[BlockId, State]: block = self.blocks[block_id] state, _ = self.evaluate(block, incoming) + exits = {target: state for target in block.exception_targets} terminator = block.terminator if isinstance(terminator, Branch): - return {terminator.then_block: state, terminator.else_block: state} + return {**exits, terminator.then_block: state, terminator.else_block: state} if isinstance(terminator, Jump): - return {terminator.target: state} + return {**exits, terminator.target: state} if isinstance(terminator, ForNext): - return {terminator.body: state, terminator.exit: state} - return {} + return {**exits, terminator.body: state, terminator.exit: state} + return exits # ------------------------------------------------------------------ transfer diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py new file mode 100644 index 0000000..3dd3a62 --- /dev/null +++ b/tests/test_exceptions.py @@ -0,0 +1,331 @@ +"""Acceptance tests for ``try`` with exception edges, and ``await`` / ``yield``. + +``docs/architecture.md`` §3.2 PyHIR, §5 CFG (exception edges), §6 PyIR. + +A ``try`` body is laid out like any block sequence; every block inside it carries +exception edges to the handler blocks, so a ``raise`` or a call in the body can reach a +handler. ``except E as name`` binds ``name`` to a ``catch`` value; the handlers, the +``else`` body and the ``finally`` body all join after the statement (``finally`` runs on +the normal path only, exceptional exits are approximated). Exception edges carry the +exit state of the raising block, so data produced in the body reaches handlers. + +Expected to remain red until ``Try``, ``Await`` and ``Yield`` reach PyHIR, the CFG and PyIR. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from coretrace_python import engine +from coretrace_python.cfg import CFG, BlockId, Jump, Raise, Return, build_cfg +from coretrace_python.cli import main +from coretrace_python.frontend import build_hir +from coretrace_python.hir import nodes +from coretrace_python.semantic.scopes import BindingKind, analyze_scopes +from coretrace_python.source import SourceManager + +REPO = Path(__file__).resolve().parent.parent +PLUGINS = REPO / "plugins" + + +@pytest.fixture(autouse=True) +def require_try() -> None: + if not hasattr(nodes, "Try") or not hasattr(nodes, "Await"): + pytest.fail("Try, Await and Yield are not in PyHIR yet") + + +def build(source_text: str) -> nodes.Module: + return build_hir(SourceManager().add_source("exc.py", source_text)) + + +def function(source_text: str) -> nodes.Function: + return next(s for s in build(source_text).body if isinstance(s, nodes.Function)) + + +def cfg_for(source_text: str) -> CFG: + return build_cfg(function(source_text)) + + +def emit(source_text: str, tmp_path: Path, capsys, *flags: str) -> str: # type: ignore[no-untyped-def] + path = tmp_path / "exc.py" + path.write_text(source_text, encoding="utf-8") + assert main(["--emit-ir", *flags, str(path)]) == 0, capsys.readouterr().err + return str(capsys.readouterr().out) + + +def b(name: str) -> BlockId: + return BlockId(name) + + +# --------------------------------------------------------------------------- PyHIR + + +def test_try_node_shape() -> None: + f = function( + "def f():\n" + " try:\n" + " risky()\n" + " except ValueError as error:\n" + " handle(error)\n" + " except (TypeError, KeyError):\n" + " pass\n" + " except:\n" + " pass\n" + " else:\n" + " ok()\n" + " finally:\n" + " done()\n" + ) + statement = f.body[0] + assert isinstance(statement, nodes.Try) + assert len(statement.body) == 1 and len(statement.orelse) == 1 and len(statement.finalbody) == 1 + first, second, bare = statement.handlers + assert isinstance(first, nodes.ExceptHandler) + assert isinstance(first.type, nodes.Name) and first.name == "error" + assert isinstance(second.type, nodes.Tuple) and second.name is None + assert bare.type is None and bare.name is None + assert first.span.start_line == 4 + + +def test_await_and_yield_nodes() -> None: + f = function("async def f(x):\n y = await x\n yield y\n yield\n") + assign, first, second = f.body + assert isinstance(assign, nodes.Assign) and isinstance(assign.value, nodes.Await) + assert isinstance(first, nodes.ExpressionStatement) and isinstance(first.expression, nodes.Yield) + assert isinstance(second, nodes.ExpressionStatement) and isinstance(second.expression, nodes.Yield) + assert second.expression.value is None + + +def test_try_star_is_still_a_diagnostic() -> None: + from coretrace_python.frontend import HIRBuildError + + with pytest.raises(HIRBuildError, match="except\\*"): + build("def f():\n try:\n pass\n except* ValueError:\n pass\n") + + +# --------------------------------------------------------------------------- scopes + + +def test_handler_names_are_locals() -> None: + module = build("def f():\n try:\n x = 1\n except Exception as error:\n y = error\n finally:\n z = 2\n") + scopes = analyze_scopes(module) + f = next(s for s in scopes.children(scopes.module_scope.id) if s.name == "f") + + assert f.bindings["error"].kind is BindingKind.LOCAL + assert {"x", "y", "z"} <= set(f.bindings) + + +# --------------------------------------------------------------------------- CFG + + +def test_try_body_blocks_have_exception_edges_to_the_handlers() -> None: + cfg = cfg_for( + "def f():\n" + " try:\n" + " risky()\n" + " except ValueError:\n" + " recover()\n" + " except KeyError:\n" + " other()\n" + " after()\n" + ) + entry = cfg.block(cfg.entry) + + assert isinstance(entry.terminator, Jump) + body = cfg.block(entry.terminator.target) + assert body.id == b("try_1") + assert body.exception_targets == (b("handler_1"), b("handler_2")) + assert cfg.successors(body.id) == (b("after_1"), b("handler_1"), b("handler_2")) + for handler in ("handler_1", "handler_2"): + assert cfg.predecessors(b(handler)) == (b("try_1"),) + assert isinstance(cfg.block(b(handler)).terminator, Jump) + assert cfg.block(b(handler)).terminator.target == b("after_1") # type: ignore[union-attr] + assert cfg.predecessors(b("after_1")) == (b("try_1"), b("handler_1"), b("handler_2")) + assert isinstance(cfg.block(b("after_1")).terminator, Return) + + +def test_raise_inside_try_reaches_the_handler() -> None: + cfg = cfg_for("def f():\n try:\n raise ValueError()\n except ValueError:\n return 1\n return 0\n") + body = cfg.block(b("try_1")) + + assert isinstance(body.terminator, Raise) + assert cfg.successors(body.id) == (b("handler_1"),) + assert b("handler_1") in cfg.reachable() + # The body always raises and the handler returns: the code after is dead. + assert b("after_1") not in cfg.reachable() + assert isinstance(cfg.block(b("handler_1")).terminator, Return) + + +def test_else_and_finally_join_on_the_normal_path() -> None: + cfg = cfg_for( + "def f():\n" + " try:\n" + " x = 1\n" + " except Exception:\n" + " x = 2\n" + " else:\n" + " x = 3\n" + " finally:\n" + " x = 4\n" + " return x\n" + ) + body = cfg.block(b("try_1")) + assert isinstance(body.terminator, Jump) + orelse = cfg.block(body.terminator.target) + assert orelse.id == b("else_1") + assert body.exception_targets == (b("handler_1"),) + assert orelse.exception_targets == () + assert isinstance(orelse.terminator, Jump) and orelse.terminator.target == b("finally_1") + handler = cfg.block(b("handler_1")) + assert isinstance(handler.terminator, Jump) and handler.terminator.target == b("finally_1") + final = cfg.block(b("finally_1")) + assert [type(s).__name__ for s in final.statements] == ["Assign"] + assert isinstance(final.terminator, Jump) and final.terminator.target == b("after_1") + + +def test_control_flow_inside_a_try_body_keeps_its_exception_edges() -> None: + cfg = cfg_for( + "def f(items):\n" + " try:\n" + " for item in items:\n" + " use(item)\n" + " except Exception:\n" + " pass\n" + " return 0\n" + ) + inside = [block for block in cfg.blocks.values() if block.exception_targets] + + assert {block.id for block in inside} >= {b("try_1"), b("loop_1"), b("body_1")} + assert all(block.exception_targets == (b("handler_1"),) for block in inside) + assert cfg.block(b("handler_1")).exception_targets == () + assert cfg.block(b("after_1")).exception_targets == () + + +def test_exception_targets_must_exist() -> None: + from coretrace_python.cfg import BasicBlock, CFGError + + span = function("def f():\n pass\n").span + entry = b("entry") + with pytest.raises(CFGError, match="missing"): + CFG(entry, {entry: BasicBlock(entry, (), Return(None, span), (b("missing"),))}) + + +def test_dominance_sees_exception_edges() -> None: + from coretrace_python.cfg import dominator_tree + + cfg = cfg_for("def f():\n try:\n a()\n except Exception:\n b()\n c()\n") + tree = dominator_tree(cfg) + + assert tree.idom(b("handler_1")) == b("try_1") + assert tree.idom(b("after_1")) == b("try_1") + assert tree.frontier(b("handler_1")) == frozenset({b("after_1")}) + + +# --------------------------------------------------------------------------- PyIR + + +def test_emit_ir_for_try_except(tmp_path, capsys) -> None: # type: ignore[no-untyped-def] + output = emit( + "def f():\n" + " try:\n" + " x = risky()\n" + " except ValueError as error:\n" + " x = error\n" + " return x\n", + tmp_path, + capsys, + ) + assert output == ( + "func @f() {\n" + "entry:\n" + " jump try_1\n" + "try_1 [except: handler_1]:\n" + " %0 = global 'risky'\n" + " %1 = call %0()\n" + ' store_local "x", %1\n' + " jump after_1\n" + "handler_1:\n" + " %2 = symbol @python.builtins.ValueError\n" + " %3 = catch %2\n" + ' store_local "error", %3\n' + ' %4 = load_local "error"\n' + ' store_local "x", %4\n' + " jump after_1\n" + "after_1:\n" + ' %5 = load_local "x"\n' + " return %5\n" + "}\n" + ) + + +def test_exception_edges_carry_the_exit_state_of_the_raising_block(tmp_path, capsys) -> None: # type: ignore[no-untyped-def] + # Approximation: a handler sees the values a body block had when control left it, so + # data produced in the body (``cmd = input()``) is visible in the handler. The older + # value of a name reassigned in the body is not modelled. + output = emit( + "def f(a):\n" + " x = a\n" + " try:\n" + " cmd = risky()\n" + " x = cmd\n" + " except Exception:\n" + " x = cmd\n" + " return x\n", + tmp_path, + capsys, + "--ssa", + ) + assert "phi" not in output.split("handler_1:")[1].split("after_1:")[0] + assert 'after_1:\n %5 = phi "x", [%2, try_1], [%2, handler_1]\n return %5\n' in output + + +def test_emit_ir_for_await_and_yield(tmp_path, capsys) -> None: # type: ignore[no-untyped-def] + output = emit("async def f(x):\n y = await x\n yield y\n yield\n", tmp_path, capsys) + assert " %1 = await %0\n" in output + assert ' %3 = yield %2\n' in output + assert " %4 = yield\n" in output + + +# --------------------------------------------------------------------------- taint + + +def test_taint_reaches_handlers_and_survives_finally() -> None: + findings = engine.check( + SourceManager().add_source( + "exc.py", + "import os\n\n" + "def run():\n" + " cmd = input()\n" + " try:\n" + " risky()\n" + " except Exception:\n" + " os.system(cmd)\n" + " finally:\n" + " os.system(cmd)\n", + ), + [PLUGINS], + ) + assert [(f.rule_id, f.span.start_line) for f in findings] == [ + ("command-injection", 8), + ("command-injection", 10), + ] + + +def test_await_propagates_taint() -> None: + findings = engine.check( + SourceManager().add_source( + "aw.py", "import os\n\nasync def run(reader):\n os.system(await reader.read(input()))\n" + ), + [PLUGINS], + ) + assert [f.rule_id for f in findings] == ["command-injection"] + + +def test_check_no_longer_flags_try_as_unsupported() -> None: + findings = engine.check( + SourceManager().add_source("t.py", "def f():\n try:\n return 1\n except Exception:\n return 2\n"), + [PLUGINS], + ) + assert findings == ()