[PATCH 2/3] test: fit: cover the kernel_noload header-size and lying-header paths

Reshape and extend the kernel_noload decompression pytests to match the
new bootm behaviour that reads the uncompressed size from the compressor
header:

  - Rename test_fit_kernel_noload_decomp_overflow to
    test_fit_kernel_noload_decomp_gzip_lying_hdr. Its setup (a 4 MiB
    payload of zeros gzipped) used to force the failure via the 8x
    heuristic starving the buffer; now that bootm reads gzip ISIZE, the
    honest trailer sizes the buffer correctly, so overwrite ISIZE with a
    tiny value instead and verify the resulting decompression is still
    stopped at the buffer boundary. This is the direct test of the
    CONFIG_SYS_BOOTM_LEN cap on the attacker-controlled header value.

  - Add test_fit_kernel_noload_decomp_gzip_hdr_sized: a 6 MiB gzipped
    payload whose compression ratio is past the 8x heuristic
    decompresses cleanly because ISIZE is consulted.

  - Add test_fit_kernel_noload_decomp_lz4_hdr_sized: the same, for lz4
    with --content-size so the frame's FLG bit is set.

  - Add test_fit_kernel_noload_decomp_zstd_hdr_sized: the same, for
    zstd whose default encoder embeds Frame_Content_Size in a
    single-segment frame.

  - Rename the pre-existing test_fit_kernel_noload_decomp_boundary to
    test_fit_kernel_noload_decomp_gzip_boundary so every noload_decomp
    test carries the compressor in its name.

  - Parametrise NOLOAD_ITS on compression so lz4, zstd, and future
    formats can share the template.

The lying-header case is covered for gzip only because the
CONFIG_SYS_BOOTM_LEN cap and buffer allocation live in a single
format-agnostic branch of bootm_load_os(): every parser feeds the same
code path, so one test is enough to exercise the security invariant
end-to-end. Per-parser correctness is covered by the hdr_sized tests
above.

The lzma branch of the helper is exercised separately in
test/lib/compression.c because standard Ubuntu ships xz-utils' lzma
shim which always writes the header size as "unknown".

Signed-off-by: Aristo Chen <[email protected]>
---
 test/py/tests/test_fit.py | 182 +++++++++++++++++++++++++++++++++-----
 1 file changed, 161 insertions(+), 21 deletions(-)

diff --git a/test/py/tests/test_fit.py b/test/py/tests/test_fit.py
index 76adb98e2c5..0e1175fbea6 100755
--- a/test/py/tests/test_fit.py
+++ b/test/py/tests/test_fit.py
@@ -118,8 +118,9 @@ host save hostfs 0 %(loadables2_addr)x %(loadables2_out)s %(loadables2_size)x
 '''
 
 # A minimal ITS for a compressed 'kernel_noload' kernel. bootm allocates a
-# per-image decompression buffer for this image type, sized as a multiple of
-# the compressed length; see the test_fit_kernel_noload_decomp_* tests.
+# per-image decompression buffer for this image type, sized either from the
+# compressor header or as a multiple of the compressed length; see the
+# test_fit_kernel_noload_decomp_* tests.
 NOLOAD_ITS = '''
 /dts-v1/;
 
@@ -133,7 +134,7 @@ NOLOAD_ITS = '''
                         type = "kernel_noload";
                         arch = "sandbox";
                         os = "linux";
-                        compression = "gzip";
+                        compression = "%(compression)s";
                         load = <0>;
                         entry = <0>;
                 };
@@ -511,14 +512,13 @@ class TestFitImage:
             + output)
 
     @pytest.mark.buildconfigspec('gzip')
-    def test_fit_kernel_noload_decomp_overflow(self, ubman, fsetup):
-        """Test that an over-large compressed kernel_noload image is rejected
+    def test_fit_kernel_noload_decomp_gzip_lying_hdr(self, ubman, fsetup):
+        """A tampered gzip ISIZE cannot shrink the buffer past the payload
 
-        For a compressed 'kernel_noload' kernel, bootm_load_os() allocates a
-        decompression buffer of ALIGN(image_len * 8, SZ_1M) and must bound the
-        decompressor by that buffer. A kernel that decompresses to far more
-        than eight times its compressed size must therefore fail with a
-        decompression error instead of overflowing the buffer.
+        bootm_load_os() sizes the kernel_noload decompression buffer from the
+        compressor header (gzip ISIZE). That value is attacker-controlled;
+        rewriting ISIZE to understate the real size must not let decompression
+        overflow the resulting buffer.
         """
         sz_1m = 1 << 20
 
@@ -527,23 +527,24 @@ class TestFitImage:
         # per-image kernel_noload buffer rather than by that global limit.
         bootm_len = int(ubman.config.buildconfig['config_sys_bootm_len'], 0)
 
-        # 4MB of zeros compresses to a few KB, so the decompression buffer
-        # (ALIGN(image_len * 8, SZ_1M), i.e. 1MB here) ends up far smaller
-        # than the uncompressed image.
         decomp_size = 4 * sz_1m
+        assert decomp_size <= bootm_len, (
+            'Test setup error: uncompressed size (%#x) must be <= '
+            'CONFIG_SYS_BOOTM_LEN (%#x)' % (decomp_size, bootm_len))
         kernel = fit_util.make_fname(ubman, 'test-noload-kernel.bin')
         with open(kernel, 'wb') as fd:
             fd.write(b'\0' * decomp_size)
         kernel_gz = self.make_compressed(ubman, kernel)
 
-        image_len = self.filesize(kernel_gz)
-        req_size = (image_len * 8 + sz_1m - 1) // sz_1m * sz_1m
-        assert req_size < decomp_size <= bootm_len, (
-            'Test setup error: need decomp buffer (%#x) < image (%#x) <= '
-            'CONFIG_SYS_BOOTM_LEN (%#x)' % (req_size, decomp_size, bootm_len))
+        # Rewrite gzip ISIZE (the last 4 bytes) to claim a tiny image, so
+        # bootm allocates ALIGN(<lie>, SZ_1M) = 1 MiB and the real 4 MiB
+        # decompression has to overrun that buffer.
+        with open(kernel_gz, 'r+b') as fd:
+            fd.seek(-4, os.SEEK_END)
+            fd.write((256).to_bytes(4, 'little'))
 
         fit = fit_util.make_fit(ubman, fsetup['mkimage'], NOLOAD_ITS,
-                                {'kernel': kernel_gz})
+                                {'kernel': kernel_gz, 'compression': 'gzip'})
         fit_addr = fsetup['fit_addr']
 
         ubman.run_command_list([
@@ -563,7 +564,146 @@ class TestFitImage:
             ubman.restart_uboot()
 
     @pytest.mark.buildconfigspec('gzip')
-    def test_fit_kernel_noload_decomp_boundary(self, ubman, fsetup):
+    def test_fit_kernel_noload_decomp_gzip_hdr_sized(self, ubman, fsetup):
+        """A well-compressed kernel_noload image fits when ISIZE is honest
+
+        bootm_load_os() reads gzip ISIZE to size the decompression buffer.
+        For a well-compressed image whose ratio exceeds the 8x fallback
+        heuristic (e.g. 6 MiB of zeros gzipping to a few KiB), an ISIZE-sized
+        buffer is the only way the decompression fits.
+        """
+        sz_1m = 1 << 20
+        bootm_len = int(ubman.config.buildconfig['config_sys_bootm_len'], 0)
+
+        # Stay under CONFIG_SYS_BOOTM_LEN so the ISIZE hint isn't rejected as
+        # bogus; still large enough that image_len * 8 falls well short.
+        decomp_size = 6 * sz_1m
+        assert decomp_size <= bootm_len, (
+            'Test setup error: decomp_size (%#x) must be <= '
+            'CONFIG_SYS_BOOTM_LEN (%#x)' % (decomp_size, bootm_len))
+        kernel = fit_util.make_fname(ubman, 'test-noload-kernel-hdrsized.bin')
+        with open(kernel, 'wb') as fd:
+            fd.write(b'\0' * decomp_size)
+        kernel_gz = self.make_compressed(ubman, kernel)
+
+        image_len = self.filesize(kernel_gz)
+        heuristic_bound = (image_len * 8 + sz_1m - 1) // sz_1m * sz_1m
+        assert heuristic_bound < decomp_size, (
+            'Test setup error: 8x heuristic bound (%#x) must be < uncompressed '
+            'size (%#x); if this fires, the compressor got less effective and '
+            'the test needs a bigger payload' % (heuristic_bound, decomp_size))
+
+        fit = fit_util.make_fit(ubman, fsetup['mkimage'], NOLOAD_ITS,
+                                {'kernel': kernel_gz, 'compression': 'gzip'},
+                                basename='test-noload-hdrsized.fit')
+        fit_addr = fsetup['fit_addr']
+
+        # Decompression must succeed: bootm read ISIZE and allocated a big
+        # enough buffer despite the ratio being past the fallback heuristic.
+        output = ubman.run_command_list([
+            'host load hostfs 0 %x %s' % (fit_addr, fit),
+            'bootm start %x' % fit_addr,
+            'bootm loados',
+        ])
+        text = '\n'.join(output)
+        assert 'Image too large' not in text, (
+            'bootm rejected a well-compressed kernel_noload image whose '
+            'ISIZE trailer records the real uncompressed size: %s' % text)
+
+    @pytest.mark.buildconfigspec('lz4')
+    @pytest.mark.requiredtool('lz4')
+    def test_fit_kernel_noload_decomp_lz4_hdr_sized(self, ubman, fsetup):
+        """A well-compressed lz4 kernel_noload image fits when the frame
+        header carries the content size.
+
+        Same as test_fit_kernel_noload_decomp_gzip_hdr_sized but for lz4: the tool
+        must be invoked with --content-size so the frame's FLG bit is set and
+        bootm can read the size instead of falling back to the 8x heuristic.
+        """
+        sz_1m = 1 << 20
+        bootm_len = int(ubman.config.buildconfig['config_sys_bootm_len'], 0)
+
+        decomp_size = 6 * sz_1m
+        assert decomp_size <= bootm_len, (
+            'Test setup error: decomp_size (%#x) must be <= '
+            'CONFIG_SYS_BOOTM_LEN (%#x)' % (decomp_size, bootm_len))
+        kernel = fit_util.make_fname(ubman, 'test-noload-kernel-lz4.bin')
+        with open(kernel, 'wb') as fd:
+            fd.write(b'\0' * decomp_size)
+        kernel_lz4 = kernel + '.lz4'
+        utils.run_and_log(
+            ubman, ['lz4', '--content-size', '-f', kernel, kernel_lz4])
+
+        image_len = self.filesize(kernel_lz4)
+        heuristic_bound = (image_len * 8 + sz_1m - 1) // sz_1m * sz_1m
+        assert heuristic_bound < decomp_size, (
+            'Test setup error: 8x heuristic bound (%#x) must be < uncompressed '
+            'size (%#x); if this fires, lz4 got less effective and the test '
+            'needs a bigger payload' % (heuristic_bound, decomp_size))
+
+        fit = fit_util.make_fit(ubman, fsetup['mkimage'], NOLOAD_ITS,
+                                {'kernel': kernel_lz4, 'compression': 'lz4'},
+                                basename='test-noload-lz4-hdrsized.fit')
+        fit_addr = fsetup['fit_addr']
+
+        output = ubman.run_command_list([
+            'host load hostfs 0 %x %s' % (fit_addr, fit),
+            'bootm start %x' % fit_addr,
+            'bootm loados',
+        ])
+        text = '\n'.join(output)
+        assert 'Image too large' not in text, (
+            'bootm rejected a well-compressed lz4 kernel_noload image whose '
+            'frame header records the real content size: %s' % text)
+
+    @pytest.mark.buildconfigspec('zstd')
+    @pytest.mark.requiredtool('zstd')
+    def test_fit_kernel_noload_decomp_zstd_hdr_sized(self, ubman, fsetup):
+        """A well-compressed zstd kernel_noload image fits when the frame
+        header carries Frame_Content_Size.
+
+        Same as test_fit_kernel_noload_decomp_gzip_hdr_sized but for zstd. The
+        default zstd encoder embeds Frame_Content_Size for a single-segment
+        frame, so no extra flag is needed; bootm reads it and sizes the
+        buffer accordingly.
+        """
+        sz_1m = 1 << 20
+        bootm_len = int(ubman.config.buildconfig['config_sys_bootm_len'], 0)
+
+        decomp_size = 6 * sz_1m
+        assert decomp_size <= bootm_len, (
+            'Test setup error: decomp_size (%#x) must be <= '
+            'CONFIG_SYS_BOOTM_LEN (%#x)' % (decomp_size, bootm_len))
+        kernel = fit_util.make_fname(ubman, 'test-noload-kernel-zstd.bin')
+        with open(kernel, 'wb') as fd:
+            fd.write(b'\0' * decomp_size)
+        kernel_zstd = kernel + '.zst'
+        utils.run_and_log(ubman, ['zstd', '-f', kernel, '-o', kernel_zstd])
+
+        image_len = self.filesize(kernel_zstd)
+        heuristic_bound = (image_len * 8 + sz_1m - 1) // sz_1m * sz_1m
+        assert heuristic_bound < decomp_size, (
+            'Test setup error: 8x heuristic bound (%#x) must be < uncompressed '
+            'size (%#x); if this fires, zstd got less effective and the test '
+            'needs a bigger payload' % (heuristic_bound, decomp_size))
+
+        fit = fit_util.make_fit(ubman, fsetup['mkimage'], NOLOAD_ITS,
+                                {'kernel': kernel_zstd, 'compression': 'zstd'},
+                                basename='test-noload-zstd-hdrsized.fit')
+        fit_addr = fsetup['fit_addr']
+
+        output = ubman.run_command_list([
+            'host load hostfs 0 %x %s' % (fit_addr, fit),
+            'bootm start %x' % fit_addr,
+            'bootm loados',
+        ])
+        text = '\n'.join(output)
+        assert 'Image too large' not in text, (
+            'bootm rejected a well-compressed zstd kernel_noload image whose '
+            'frame header records the real content size: %s' % text)
+
+    @pytest.mark.buildconfigspec('gzip')
+    def test_fit_kernel_noload_decomp_gzip_boundary(self, ubman, fsetup):
         """Test that decompression succeeds exactly at the buffer limit
 
         For a compressed 'kernel_noload' kernel, bootm_load_os() allocates a
@@ -589,7 +729,7 @@ class TestFitImage:
             % (decomp_size, req_size))
 
         fit = fit_util.make_fit(ubman, fsetup['mkimage'], NOLOAD_ITS,
-                                {'kernel': kernel_gz},
+                                {'kernel': kernel_gz, 'compression': 'gzip'},
                                 basename='test-noload-boundary.fit')
         fit_addr = fsetup['fit_addr']
 
-- 
2.43.0