fix[venom]: port augassign read/write overlap guard

tests/functional/codegen/features/test_assignment.py

@@ -1,5 +1,6 @@
 import pytest
 
+from vyper import compile_code
 from vyper.evm.opcodes import version_check
 from vyper.exceptions import CodegenPanic, ImmutableViolation, InvalidType, TypeMismatch
 
@@ -40,16 +41,16 @@ def augmod(x: int128, y: int128) -> int128:
     print("Passed aug-assignment test")
 
 
-@pytest.mark.parametrize(
-    "source",
-    [
-        """
+# augassigns which can write to the target out from under us;
+# rejected with a (compile-time) panic - GHSA-4w26-8p97-f4jp
+AUGASSIGN_OOB_SOURCES = [
+    """
 @external
 def poc():
     a: DynArray[uint256, 2] = [1, 2]
     a[1] += a.pop()
     """,
-        """
+    """
 a: DynArray[uint256, 2]
 
 def side_effect() -> uint256:
@@ -60,7 +61,7 @@ def poc():
     self.a = [1, 2]
     self.a[1] += self.side_effect()
     """,
-        """
+    """
 a: DynArray[uint256, 2]
 
 def side_effect() -> uint256:
@@ -72,7 +73,7 @@ def poc():
     self.a = [1, 2]
     self.a[1] += self.side_effect()
     """,
-        """
+    """
 a: DynArray[uint256, 2]
 
 interface Foo:
@@ -88,29 +89,19 @@ def poc():
     # panics due to extcall
     self.a[1] += extcall Foo(self).foo()
     """,
-    ],
-)
-@pytest.mark.xfail(strict=True, raises=CodegenPanic)
-def test_augassign_oob(get_contract, tx_failed, source):
-    # xfail here (with panic):
-    c = get_contract(source)
-
-    # not reached until the panic is fixed
-    with tx_failed(c):
-        c.poc()
-
+]
 
-@pytest.mark.parametrize(
-    "source",
-    [
-        """
+# augassigns where the rhs references the lhs but cannot write to it;
+# these compile and run fine
+AUGASSIGN_OVERLAP_OK_SOURCES = [
+    """
 @external
 def entry() -> DynArray[uint256, 2]:
     a: DynArray[uint256, 2] = [1, 1]
     a[1] += a[1]
     return a
     """,
-        """
+    """
 @external
 def entry() -> DynArray[uint256, 2]:
     a: uint256 = 1
@@ -120,7 +111,7 @@ def entry() -> DynArray[uint256, 2]:
     b[0] += b[1] // 2
     return b
     """,
-        """
+    """
 a: DynArray[uint256, 2]
 
 def read() -> uint256:
@@ -132,7 +123,7 @@ def entry() -> DynArray[uint256, 2]:
     self.a[1] += self.read()
     return self.a
     """,
-        """
+    """
 interface Foo:
     def foo() -> uint256: nonpayable
 
@@ -148,7 +139,7 @@ def entry() -> DynArray[uint256, 2]:
     a[1] += extcall Foo(self).foo()
     return a
     """,
-        """
+    """
 interface Foo:
     def foo() -> uint256: nonpayable
 
@@ -168,7 +159,7 @@ def entry() -> DynArray[uint256, 2]:
     a[1] += self.get_foo()
     return a
     """,
-        """
+    """
 a: public(DynArray[uint256, 2])
 
 interface Foo:
@@ -184,13 +175,41 @@ def entry() -> DynArray[uint256, 2]:
     self.a[1] += staticcall Foo(self).foo()
     return self.a
     """,
-    ],
-)
+]
+
+
+@pytest.mark.parametrize("source", AUGASSIGN_OOB_SOURCES)
+@pytest.mark.xfail(strict=True, raises=CodegenPanic)
+def test_augassign_oob(get_contract, tx_failed, source):
+    # xfail here (with panic):
+    c = get_contract(source)
+
+    # not reached until the panic is fixed
+    with tx_failed(c):
+        c.poc()
+
+
+@pytest.mark.parametrize("source", AUGASSIGN_OOB_SOURCES)
+def test_augassign_oob_venom(source):
+    # the venom pipeline must reject the same programs as legacy
+    # (GH issue #5051)
+    with pytest.raises(CodegenPanic):
+        compile_code(source)
+
+
+@pytest.mark.parametrize("source", AUGASSIGN_OVERLAP_OK_SOURCES)
 def test_augassign_rhs_references_lhs2(get_contract, source):
     c = get_contract(source)
     assert c.entry() == [1, 2]
 
 
+@pytest.mark.parametrize("source", AUGASSIGN_OVERLAP_OK_SOURCES)
+def test_augassign_rhs_references_lhs_venom(source):
+    # the venom port of the augassign overlap guard must not reject
+    # programs which legacy accepts (GH issue #5051)
+    assert compile_code(source) is not None
+
+
 @pytest.mark.requires_evm_version("cancun")
 def test_augassign_rhs_references_lhs_transient(get_contract):
     source = """
@@ -213,10 +232,8 @@ def entry() -> DynArray[uint256, 2]:
     assert c.entry() == [2, 3]
 
 
-@pytest.mark.parametrize(
-    "source",
-    [
-        """
+AUGASSIGN_OOB_TRANSIENT_SOURCES = [
+    """
 x: transient(DynArray[uint256, 2])
 
 def write() -> uint256:
@@ -229,7 +246,7 @@ def entry() -> DynArray[uint256, 2]:
     self.x[1] += self.write()
     return self.x
     """,
-        """
+    """
 x: transient(DynArray[uint256, 2])
 
 @external
@@ -239,8 +256,10 @@ def entry() -> DynArray[uint256, 2]:
     self.x[1] += self.x.pop()
     return self.x
     """,
-    ],
-)
+]
+
+
+@pytest.mark.parametrize("source", AUGASSIGN_OOB_TRANSIENT_SOURCES)
 @pytest.mark.xfail(strict=True, raises=CodegenPanic)
 def test_augassign_rhs_references_lhs_transient2(get_contract, tx_failed, source):
     if not version_check(begin="cancun"):
@@ -255,6 +274,15 @@ def test_augassign_rhs_references_lhs_transient2(get_contract, tx_failed, source
         c.entry()
 
 
+@pytest.mark.parametrize("source", AUGASSIGN_OOB_TRANSIENT_SOURCES)
+def test_augassign_oob_transient_venom(source):
+    # the venom pipeline must reject the same programs as legacy
+    # (GH issue #5051). note: the default evm version supports transient
+    # storage, no need for a version check here.
+    with pytest.raises(CodegenPanic):
+        compile_code(source)
+
+
 @pytest.mark.parametrize(
     "typ,in_val,out_val",
     [

vyper/codegen_venom/stmt.py

@@ -15,9 +15,10 @@
 from vyper.codegen_venom.abi import abi_encode_to_buf
 from vyper.codegen_venom.arithmetic import apply_binop
 from vyper.exceptions import CodegenPanic, CompilerPanic, TypeCheckFailure, tag_exceptions
+from vyper.semantics.analysis.utils import get_expr_writes
 from vyper.semantics.data_locations import DataLocation
 from vyper.semantics.types.bytestrings import _BytestringT
-from vyper.semantics.types.function import ContractFunctionT
+from vyper.semantics.types.function import ContractFunctionT, StateMutability
 from vyper.semantics.types.subscriptable import DArrayT, SArrayT, TupleT
 from vyper.semantics.types.user import EventT, StructT
 from vyper.utils import method_id_int
@@ -30,6 +31,67 @@
 from .value import VyperValue
 
 
+# AST-level equivalent of legacy `IRnode.referenced_variables`: all
+# variables read while evaluating `node` (including state accesses of
+# internal functions called within `node`, which the frontend folds
+# into the call node's `_reads`).
+def _referenced_variables(node: vy_ast.VyperNode) -> set:
+    node = node.reduced()
+    ret: set = set()
+    if isinstance(node, vy_ast.ExprNode) and node._expr_info is not None:
+        ret.update(access.variable for access in node._expr_info._reads)
+    for c in node._children:
+        ret |= _referenced_variables(c)
+    return ret
+
+
+# AST-level equivalent of legacy `IRnode.contains_writeable_call`: return
+# True if evaluating `node` can issue a state-modifying call (i.e. one
+# which compiles to `call`/`delegatecall`/`create`/`create2`), including
+# transitively through internal function calls.
+def _contains_writeable_call(node: vy_ast.VyperNode) -> bool:
+    if _emits_writeable_call(node):
+        return True
+    fns = set()
+    for fn_t in _called_internal_functions(node):
+        fns.add(fn_t)
+        fns.update(fn_t.reachable_internal_functions)
+    return any(_emits_writeable_call(fn_t.decl_node) for fn_t in fns)
+
+
+def _called_internal_functions(node: vy_ast.VyperNode) -> set:
+    ret = set()
+    for call in node.get_descendants(vy_ast.Call, include_self=True):
+        func_t = call.func._metadata.get("type")
+        if isinstance(func_t, ContractFunctionT) and (func_t.is_internal or func_t.is_constructor):
+            ret.add(func_t.get_concrete_override())
+    return ret
+
+
+# check if `node` directly contains a state-modifying call (without
+# recursing into internal functions)
+def _emits_writeable_call(node: vy_ast.VyperNode) -> bool:
+    # delayed import due to import cycle
+    from vyper.builtins.functions import RawCall, Send, _CreateBase
+
+    # `extcall` compiles to the `call` opcode (`staticcall`s are emitted
+    # via vy_ast.StaticCall and cannot modify state)
+    if len(node.get_descendants(vy_ast.ExtCall, include_self=True)) > 0:
+        return True
+
+    for call in node.get_descendants(vy_ast.Call, include_self=True):
+        func_t = call.func._metadata.get("type")
+        if isinstance(func_t, RawCall):
+            # raw_call compiles to `staticcall` when `is_static_call=True`
+            if func_t.get_mutability_at_call_site(call) > StateMutability.VIEW:
+                return True
+        elif isinstance(func_t, (Send, _CreateBase)):
+            # `send` compiles to `call`, create builtins to `create`/`create2`
+            return True
+
+    return False
+
+
 class Stmt:
     """Lower Vyper statements to Venom IR."""
 
@@ -290,6 +352,21 @@ def lower_AugAssign(self) -> None:
         if not target_typ._is_prim_word:  # pragma: nocover
             raise TypeCheckFailure("AugAssign only valid for primitive types")
 
+        # oob - GHSA-4w26-8p97-f4jp
+        # ported from legacy codegen (vyper/codegen/stmt.py:parse_AugAssign).
+        # the target pointer is computed (and bounds-checked) before the RHS
+        # is evaluated, so reject the augassign if the RHS could mutate a
+        # complex variable which the target points into (e.g.
+        # `self.a[1] += self.a.pop()`).
+        rhs_writes = set(access.variable for access in get_expr_writes(right_node))
+        for var in _referenced_variables(target):
+            if var.typ._is_prim_word:
+                continue
+            if var in rhs_writes or (
+                var.is_state_variable() and _contains_writeable_call(right_node)
+            ):
+                raise CodegenPanic("unreachable")
+
         # Get target pointer (with location info)
         dst_ptr = self._get_target_ptr(target)