kleaf: Handle lz4 compressed ramdisks in initramfs module lists tests

The initramfs.img can be either gzip or lz4 compressed, so ensure both
compression formats can be handled.

Bug: 324304750
Bug: 324356826
Signed-off-by: Isaac J. Manjarres <isaacmanjarres@google.com>
(cherry picked from https://android-review.googlesource.com/q/commit:c4a8127d9d6606963e4605aed8764dc6c6ca2751)
Merged-In: I18655c85a329ea7a169154213739e7b9850aacd8
Change-Id: I18655c85a329ea7a169154213739e7b9850aacd8

1 files changed, 35 insertions, 4 deletions

diff --git a/kleaf/artifact_tests/initramfs_modules_lists_test.py b/kleaf/artifact_tests/initramfs_modules_lists_test.py
index 7129255..f1d655f 100644
--- a/kleaf/artifact_tests/initramfs_modules_lists_test.py
+++ b/kleaf/artifact_tests/initramfs_modules_lists_test.py
@@ -35,6 +35,36 @@ arguments = None
 
 
 class InitramfsModulesLists(unittest.TestCase):
+    def _detect_decompression_cmd(self, initramfs):
+        """Determines what commands to use for decompressing initramfs.img
+
+        Args:
+            initramfs: The path to the initramfs.img to decompress
+
+        Returns:
+            The command that should be used to decompress the image.
+        """
+        magic_to_decompression_command = {
+            # GZIP
+            b'\x1f\x8b\x08': ["gzip", "-c", "-d"],
+            # LZ4
+            # The kernel build uses legacy LZ4 compression (i.e. lz4 -l ...),
+            # so the legacy LZ4 magic must be used in little-endian format.
+            b'\x02\x21\x4c\x18': ["lz4", "-c", "-d", "-l"],
+        }
+        max_magic_len = max(len(magic) for magic in magic_to_decompression_command)
+
+        with open(initramfs, "rb") as initramfs_file:
+            hdr = initramfs_file.read(max_magic_len)
+
+        self.assertIsNotNone(hdr)
+
+        for magic, command in magic_to_decompression_command.items():
+            if hdr.startswith(magic):
+                return command
+
+        self.fail("No suitable compression method found")
+
     def _decompress_initramfs(self, initramfs, temp_dir):
         """Decompress initramfs into temp_dir.
 
@@ -42,13 +72,14 @@ class InitramfsModulesLists(unittest.TestCase):
             initramfs: path to initramfs.img gzip file to be decompressed
             temp_dir: directory in which to decompress initramfs.img into
         """
+        decompression_cmd = self._detect_decompression_cmd(initramfs)
         with open(initramfs) as initramfs_file:
             with subprocess.Popen(["cpio", "-i"], cwd=temp_dir,
                                   stdin=subprocess.PIPE, stdout=subprocess.PIPE) as cpio_sp:
-                with subprocess.Popen(["gzip", "-c", "-d"], stdin=initramfs_file,
-                                      stdout=cpio_sp.stdin) as gzip_sp:
-                    gzip_sp.communicate()
-                    self.assertEqual(0, gzip_sp.returncode)
+                with subprocess.Popen(decompression_cmd, stdin=initramfs_file,
+                                      stdout=cpio_sp.stdin) as decompress_sp:
+                    decompress_sp.communicate()
+                    self.assertEqual(0, decompress_sp.returncode)
 
     def _diff_modules_lists(self, modules_lists_map, modules_dir):
         """Compares generated modules lists against expected modules lists for equality.