zmake: test final file output size

Zephyr's build allows the image to consume the full flash size
(CONFIG_FLASH_SIZE) since it doesn't assume anything about how
chromium lays out the images (using RO/RW). This means that in
systems with 512Kb of flash space, images taking up more than
256Kb will successfully build (even though the final image will
be larger than 512Kb).

Add a check in the output packers for the final size. This test
ignores by default the .elf extension.

BRANCH=none
BUG=b:190435084
TEST=Added unit tests
TEST=build brya

Signed-off-by: Yuval Peress <peress@chromium.org>
Change-Id: I94f1657d0ff44d79920ae5e8e7f11edf1580de05
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/2948169
Reviewed-by: Denis Brockus <dbrockus@chromium.org>

2 files changed, 102 insertions, 1 deletions

diff --git a/zephyr/zmake/tests/test_packers.py b/zephyr/zmake/tests/test_packers.py
new file mode 100644
index 0000000000..21361a925f
--- /dev/null
+++ b/zephyr/zmake/tests/test_packers.py
@@ -0,0 +1,51 @@
+# Copyright 2021 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import hypothesis
+import hypothesis.strategies as st
+import pathlib
+import pytest
+import tempfile
+import unittest.mock as mock
+
+import zmake.output_packers as packers
+
+# Strategies for use with hypothesis
+absolute_path = st.from_regex(regex=r"\A/[\w/]*\Z")
+
+@hypothesis.given(absolute_path)
+@hypothesis.settings(deadline=60000)
+def test_file_size_unbounded(path):
+    packer = packers.BasePacker(project=None)
+    packer._is_size_bound = mock.Mock(name='_is_size_bound', return_value=False)
+    file = pathlib.Path(path) / 'zephyr.bin'
+    assert packer._check_packed_file_size(file=file, dirs={}) == file
+    packer._is_size_bound.assert_called_once_with(file)
+
+@hypothesis.given(st.binary(min_size=5, max_size=100))
+@hypothesis.settings(deadline=60000)
+def test_file_size_in_bounds(data):
+    packer = packers.BasePacker(project=None)
+    packer._is_size_bound = mock.Mock(name='_is_size_bound', return_value=True)
+    packer._get_max_image_bytes = mock.Mock(name='_get_max_image_bytes',
+                                            return_value=100)
+    with tempfile.TemporaryDirectory() as temp_dir_name:
+        file = pathlib.Path(temp_dir_name) / 'zephyr.bin'
+        with open(file, 'wb') as f:
+            f.write(data)
+        assert packer._check_packed_file_size(file=file, dirs={}) == file
+
+@hypothesis.given(st.binary(min_size=101, max_size=200))
+@hypothesis.settings(deadline=60000)
+def test_file_size_out_of_bounds(data):
+    packer = packers.BasePacker(project=None)
+    packer._is_size_bound = mock.Mock(name='_is_size_bound', return_value=True)
+    packer._get_max_image_bytes = mock.Mock(name='_get_max_image_bytes',
+                                            return_value=100)
+    with tempfile.TemporaryDirectory() as temp_dir_name:
+        file = pathlib.Path(temp_dir_name) / 'zephyr.bin'
+        with open(file, 'wb') as f:
+            f.write(data)
+        with pytest.raises(RuntimeError):
+            packer._check_packed_file_size(file=file, dirs={})
\ No newline at end of file
diff --git a/zephyr/zmake/zmake/output_packers.py b/zephyr/zmake/zmake/output_packers.py
index 586e5c00b3..1e09e9d580 100644
--- a/zephyr/zmake/zmake/output_packers.py
+++ b/zephyr/zmake/zmake/output_packers.py
@@ -45,6 +45,46 @@ class BasePacker:
         """
         raise NotImplementedError('Abstract method not implemented')
 
+    def _get_max_image_bytes(self):
+        """Get the maximum allowed image size (in bytes).
+
+        This value will generally be found in CONFIG_FLASH_SIZE but may vary
+        depending on the specific way things are being packed.
+
+        Returns:
+            The maximum allowed size of the image in bytes.
+        """
+        raise NotImplementedError('Abstract method not implemented')
+
+    def _is_size_bound(self, path):
+        """Check whether the given path should be constrained by size.
+
+        Generally, .elf files will be unconstrained while .bin files will be
+        constrained.
+
+        Args:
+            path: A file's path to test.
+
+        Returns:
+            True if the file size should be checked. False otherwise.
+        """
+        return path.suffix == '.bin'
+
+    def _check_packed_file_size(self, file, dirs):
+        """Check that a packed file passes size constraints.
+
+        Args:
+            file: A file to test.
+            dirs: A map of the arguments to pass to _get_max_image_bytes
+
+        Returns:
+            The file if it passes the test.
+        """
+        if not self._is_size_bound(
+            file) or file.stat().st_size <= self._get_max_image_bytes(**dirs):
+            return file
+        raise RuntimeError('Output file ({}) too large'.format(file))
+
 
 class ElfPacker(BasePacker):
     """Raw proxy for ELF output of a single build."""
@@ -120,10 +160,20 @@ class NpcxPacker(BasePacker):
         if proc.wait(timeout=60):
             raise OSError('Failed to run binman')
 
-        yield work_dir / 'zephyr.bin', 'zephyr.bin'
+        yield self._check_packed_file_size(work_dir / 'zephyr.bin',
+                                           {'ro': ro, 'rw': rw}), 'zephyr.bin'
         yield ro / 'zephyr' / 'zephyr.elf', 'zephyr.ro.elf'
         yield rw / 'zephyr' / 'zephyr.elf', 'zephyr.rw.elf'
 
+    def _get_max_image_bytes(self, ro, rw):
+        ro_size = util.read_kconfig_autoconf_value(
+            ro / 'zephyr' / 'include' / 'generated',
+            'CONFIG_FLASH_SIZE')
+        rw_size = util.read_kconfig_autoconf_value(
+            ro / 'zephyr' / 'include' / 'generated',
+            'CONFIG_FLASH_SIZE')
+        return max(int(ro_size, 0), int(rw_size, 0)) * 1024
+
 
 # A dictionary mapping packer config names to classes.
 packer_registry = {