From 04bcb5487db246230aa08acb7d63b3f139d879ba Mon Sep 17 00:00:00 2001
From: pdobacz <5735525+pdobacz@users.noreply.github.com>
Date: Fri, 14 Jun 2024 17:29:47 +0200
Subject: [PATCH] new(tests): EOF - EIP-3540 container size

Tests https://github.com/ipsilon/eof/pull/125
---
 .../exceptions/evmone_exceptions.py           |   3 +
 .../exceptions/exceptions.py                  |   4 +
 .../eip3540_eof_v1/test_container_size.py     | 114 ++++++++++++++++++
 3 files changed, 121 insertions(+)
 create mode 100644 tests/prague/eip7692_eof_v1/eip3540_eof_v1/test_container_size.py

diff --git a/src/ethereum_test_tools/exceptions/evmone_exceptions.py b/src/ethereum_test_tools/exceptions/evmone_exceptions.py
index f3bd244d1f..112844b06c 100644
--- a/src/ethereum_test_tools/exceptions/evmone_exceptions.py
+++ b/src/ethereum_test_tools/exceptions/evmone_exceptions.py
@@ -73,6 +73,9 @@ class EvmoneExceptionMapper:
         ExceptionMessage(
             EOFException.TOPLEVEL_CONTAINER_TRUNCATED, "err: toplevel_container_truncated"
         ),
+        ExceptionMessage(
+            EOFException.CONTAINER_SIZE_ABOVE_LIMIT, "err: container_size_above_limit"
+        ),
     )
 
     def __init__(self) -> None:
diff --git a/src/ethereum_test_tools/exceptions/exceptions.py b/src/ethereum_test_tools/exceptions/exceptions.py
index 9d19792306..f43ec1dd74 100644
--- a/src/ethereum_test_tools/exceptions/exceptions.py
+++ b/src/ethereum_test_tools/exceptions/exceptions.py
@@ -688,6 +688,10 @@ class EOFException(ExceptionBase):
     """
     Top-level EOF container has data section truncated
     """
+    CONTAINER_SIZE_ABOVE_LIMIT = auto()
+    """
+    EOF container is above size limit
+    """
 
 
 """
diff --git a/tests/prague/eip7692_eof_v1/eip3540_eof_v1/test_container_size.py b/tests/prague/eip7692_eof_v1/eip3540_eof_v1/test_container_size.py
new file mode 100644
index 0000000000..f24dc5ed9a
--- /dev/null
+++ b/tests/prague/eip7692_eof_v1/eip3540_eof_v1/test_container_size.py
@@ -0,0 +1,114 @@
+"""
+EOF validation tests for EIP-3540 container size
+"""
+
+import pytest
+
+from ethereum_test_tools import EOFTestFiller
+from ethereum_test_tools import Opcodes as Op
+from ethereum_test_tools.eof.v1 import Container, EOFException, Section
+from ethereum_test_tools.eof.v1.constants import MAX_INITCODE_SIZE
+
+from .. import EOF_FORK_NAME
+
+REFERENCE_SPEC_GIT_PATH = "EIPS/eip-3540.md"
+REFERENCE_SPEC_VERSION = ""  # todo
+
+pytestmark = pytest.mark.valid_from(EOF_FORK_NAME)
+
+VALID_CONTAINER = Container(sections=[Section.Code(code=Op.STOP)])
+
+
+@pytest.mark.parametrize(
+    "over_limit",
+    [0, 1, 2, 2**16 - MAX_INITCODE_SIZE],
+)
+def test_max_size(
+    eof_test: EOFTestFiller,
+    over_limit: int,
+):
+    """
+    Verify EOF container valid at maximum size, invalid above
+    """
+    minimal_code = bytearray(bytes(VALID_CONTAINER))
+    # Expand the minimal EOF code by more noop code, reaching the desired target container size.
+    code = Container(
+        sections=[
+            Section.Code(
+                code=Op.JUMPDEST * (MAX_INITCODE_SIZE - len(minimal_code) + over_limit) + Op.STOP
+            )
+        ]
+    )
+    assert len(code) == MAX_INITCODE_SIZE + over_limit
+    eof_test(
+        data=bytes(code),
+        expect_exception=None if over_limit == 0 else EOFException.CONTAINER_SIZE_ABOVE_LIMIT,
+    )
+
+
+@pytest.mark.parametrize(
+    "size",
+    [MAX_INITCODE_SIZE + 1, MAX_INITCODE_SIZE * 2],
+)
+def test_above_max_size_raw(
+    eof_test: EOFTestFiller,
+    size: int,
+):
+    """
+    Verify EOF container valid at maximum size, regardless of header contents
+    """
+    code = Op.INVALID * size
+    eof_test(
+        data=bytes(code),
+        expect_exception=EOFException.CONTAINER_SIZE_ABOVE_LIMIT,
+    )
+
+
+@pytest.mark.parametrize(
+    "code",
+    [
+        pytest.param(
+            Container(sections=[Section.Code(code=Op.STOP, custom_size=MAX_INITCODE_SIZE)]),
+            id="1st_code_section",
+        ),
+        pytest.param(
+            Container(
+                sections=[
+                    Section.Code(code=Op.STOP),
+                    Section.Code(code=Op.STOP, custom_size=MAX_INITCODE_SIZE),
+                ]
+            ),
+            id="2nd_code_section",
+        ),
+        pytest.param(
+            Container(
+                sections=[
+                    Section.Code(code=Op.STOP),
+                    Section.Container(container=Op.STOP, custom_size=MAX_INITCODE_SIZE),
+                ]
+            ),
+            id="1st_container_section",
+        ),
+        pytest.param(
+            Container(
+                sections=[
+                    Section.Code(code=Op.STOP),
+                    Section.Container(container=Op.STOP),
+                    Section.Container(container=Op.STOP, custom_size=MAX_INITCODE_SIZE),
+                ]
+            ),
+            id="2nd_container_section",
+        ),
+    ],
+)
+def test_section_after_end_of_container(
+    eof_test: EOFTestFiller,
+    code: Container,
+):
+    """
+    Verify EOF container is invalid if any of sections declares above container size
+    """
+    eof_test(
+        data=bytes(code),
+        expect_exception=EOFException.INVALID_SECTION_BODIES_SIZE,
+    )