Merge changes from topic "microdroid_gki_lz4" into main

* changes:
  [release] Replace initrds in rialto for GKI kernels during resigning
  Compress arm64 GKI with lz4
diff --git a/apex/Android.bp b/apex/Android.bp
index 399ad97..3b5141e 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -254,6 +254,7 @@
         "initrd_bootconfig",
         "lpmake",
         "lpunpack",
+        "lz4",
         "simg2img",
     ],
 }
@@ -274,6 +275,7 @@
         "initrd_bootconfig",
         "lpmake",
         "lpunpack",
+        "lz4",
         "sign_virt_apex",
         "simg2img",
     ],
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index 0b6137b..7c59b54 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -153,12 +153,18 @@
                '--key', key, '--output', output])
 
 
+def is_lz4(args, path):
+    # error 44: Unrecognized header
+    result = RunCommand(args, ['lz4', '-t', path], expected_return_values={0, 44})
+    return result[1] == 0
+
+
 def AvbInfo(args, image_path):
     """Parses avbtool --info image output
 
     Args:
       args: program arguments.
-      image_path: The path to the image.
+      image_path: The path to the image, either raw or lz4 compressed
       descriptor_name: Descriptor name of interest.
 
     Returns:
@@ -169,6 +175,11 @@
     if not os.path.exists(image_path):
         raise ValueError(f'Failed to find image: {image_path}')
 
+    if is_lz4(args, image_path):
+        with tempfile.NamedTemporaryFile() as decompressed_image:
+            RunCommand(args, ['lz4', '-d', '-f', image_path, decompressed_image.name])
+            return AvbInfo(args, decompressed_image.name)
+
     output, ret_code = RunCommand(
         args, ['avbtool', 'info_image', '--image', image_path], expected_return_values={0, 1})
     if ret_code == 1:
@@ -560,11 +571,7 @@
                             wait=[vbmeta_f])
 
     # Re-sign kernel. Note kernel's vbmeta contain addition descriptor from ramdisk(s)
-    def resign_kernel(kernel, initrd_normal, initrd_debug):
-        kernel_file = files[kernel]
-        initrd_normal_file = files[initrd_normal]
-        initrd_debug_file = files[initrd_debug]
-
+    def resign_decompressed_kernel(kernel_file, initrd_normal_file, initrd_debug_file):
         _, kernel_image_descriptors = AvbInfo(args, kernel_file)
         salts = extract_hash_descriptors(
             kernel_image_descriptors, lambda descriptor: descriptor['Salt'])
@@ -580,21 +587,47 @@
               additional_images=[initrd_normal_hashdesc, initrd_debug_hashdesc],
               wait=[initrd_n_f, initrd_d_f])
 
+    def resign_compressed_kernel(kernel_file, initrd_normal_file, initrd_debug_file):
+        # decompress, re-sign, compress again
+        with tempfile.TemporaryDirectory() as work_dir:
+            decompressed_kernel_file = os.path.join(work_dir, os.path.basename(kernel_file))
+            RunCommand(args, ['lz4', '-d', kernel_file, decompressed_kernel_file])
+            resign_decompressed_kernel(decompressed_kernel_file, initrd_normal_file,
+                                       initrd_debug_file).result()
+            RunCommand(args, ['lz4', '-9', '-f', decompressed_kernel_file, kernel_file])
+
+    def resign_kernel(kernel, initrd_normal, initrd_debug):
+        kernel_file = files[kernel]
+        initrd_normal_file = files[initrd_normal]
+        initrd_debug_file = files[initrd_debug]
+
+        # kernel may be compressed with lz4.
+        if is_lz4(args, kernel_file):
+            return Async(resign_compressed_kernel, kernel_file, initrd_normal_file,
+                         initrd_debug_file)
+        else:
+            return resign_decompressed_kernel(kernel_file, initrd_normal_file, initrd_debug_file)
+
     _, original_kernel_descriptors = AvbInfo(args, files['kernel'])
-    resign_kernel_task = resign_kernel('kernel', 'initrd_normal.img', 'initrd_debuggable.img')
+    resign_kernel_tasks = [resign_kernel('kernel', 'initrd_normal.img', 'initrd_debuggable.img')]
+    original_kernels = {"kernel" : original_kernel_descriptors}
 
     for ver in gki_versions:
         if f'gki-{ver}_kernel' in files:
-            resign_kernel(
-                f'gki-{ver}_kernel',
+            kernel_name = f'gki-{ver}_kernel'
+            _, original_kernel_descriptors = AvbInfo(args, files[kernel_name])
+            task = resign_kernel(
+                kernel_name,
                 f'gki-{ver}_initrd_normal.img',
                 f'gki-{ver}_initrd_debuggable.img')
+            resign_kernel_tasks.append(task)
+            original_kernels[kernel_name] = original_kernel_descriptors
 
     # Re-sign rialto if it exists. Rialto only exists in arm64 environment.
     if os.path.exists(files['rialto']):
         update_initrd_digests_task = Async(
-            update_initrd_digests_in_rialto, original_kernel_descriptors, args,
-            files, wait=[resign_kernel_task])
+            update_initrd_digests_of_kernels_in_rialto, original_kernels, args, files,
+            wait=resign_kernel_tasks)
         Async(resign_rialto, args, key, files['rialto'], wait=[update_initrd_digests_task])
 
 def resign_rialto(args, key, rialto_path):
@@ -628,18 +661,7 @@
         f"Value of '{key}' should change for '{context}'" \
         f"Original value: {original[key]}, updated value: {updated[key]}"
 
-def update_initrd_digests_in_rialto(original_descriptors, args, files):
-    _, updated_descriptors = AvbInfo(args, files['kernel'])
-
-    original_digests = extract_hash_descriptors(
-        original_descriptors, lambda x: binascii.unhexlify(x['Digest']))
-    updated_digests = extract_hash_descriptors(
-        updated_descriptors, lambda x: binascii.unhexlify(x['Digest']))
-    assert original_digests.pop("boot") == updated_digests.pop("boot"), \
-        "Hash descriptor of boot should not change for kernel. " \
-        f"Original descriptors: {original_descriptors}, " \
-        f"updated descriptors: {updated_descriptors}"
-
+def update_initrd_digests_of_kernels_in_rialto(original_kernels, args, files):
     # Update the hashes of initrd_normal and initrd_debug in rialto if the
     # bootconfigs in them are updated.
     if args.do_not_update_bootconfigs:
@@ -648,6 +670,26 @@
     with open(files['rialto'], "rb") as file:
         content = file.read()
 
+    for kernel_name, descriptors in original_kernels.items():
+        content = update_initrd_digests_in_rialto(
+            descriptors, args, files, kernel_name, content)
+
+    with open(files['rialto'], "wb") as file:
+        file.write(content)
+
+def update_initrd_digests_in_rialto(
+        original_descriptors, args, files, kernel_name, content):
+    _, updated_descriptors = AvbInfo(args, files[kernel_name])
+
+    original_digests = extract_hash_descriptors(
+        original_descriptors, lambda x: binascii.unhexlify(x['Digest']))
+    updated_digests = extract_hash_descriptors(
+        updated_descriptors, lambda x: binascii.unhexlify(x['Digest']))
+    assert original_digests.pop("boot") == updated_digests.pop("boot"), \
+        "Hash descriptor of boot should not change for " + kernel_name + \
+        f"\nOriginal descriptors: {original_descriptors}, " \
+        f"\nUpdated descriptors: {updated_descriptors}"
+
     # Check that the original and updated digests are different before updating rialto.
     partition_names = {'initrd_normal', 'initrd_debug'}
     assert set(original_digests.keys()) == set(updated_digests.keys()) == partition_names, \
@@ -671,8 +713,7 @@
             f"original digest of the partition {partition_name} not found."
         content = new_content
 
-    with open(files['rialto'], "wb") as file:
-        file.write(content)
+    return content
 
 def extract_hash_descriptors(descriptors, f=lambda x: x):
     return {desc["Partition Name"]: f(desc) for desc in
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 36688fc..999dc52 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -559,7 +559,7 @@
 avb_add_hash_footer {
     name: "microdroid_gki-android14-6.1_kernel_signed",
     defaults: ["microdroid_kernel_signed_defaults"],
-    filename: "microdroid_gki-android14-6.1_kernel",
+    filename: "microdroid_gki-android14-6.1_kernel_signed",
     arch: {
         arm64: {
             src: ":microdroid_gki_kernel_prebuilts-6.1-arm64",
@@ -574,13 +574,29 @@
     ],
 }
 
+// HACK: use cc_genrule for arch-specific properties
+cc_genrule {
+    name: "microdroid_gki-android14-6.1_kernel_signed-lz4",
+    out: ["microdroid_gki-android14-6.1_kernel_signed-lz4"],
+    srcs: [":empty_file"],
+    arch: {
+        arm64: {
+            srcs: [":microdroid_gki-android14-6.1_kernel_signed"],
+            exclude_srcs: [":empty_file"],
+        },
+    },
+    tools: ["lz4"],
+    cmd: "$(location lz4) -9 $(in) $(out)",
+}
+
 prebuilt_etc {
     name: "microdroid_gki-android14-6.1_kernel",
+    filename: "microdroid_gki-android14-6.1_kernel",
     src: ":empty_file",
     relative_install_path: "fs",
     arch: {
         arm64: {
-            src: ":microdroid_gki-android14-6.1_kernel_signed",
+            src: ":microdroid_gki-android14-6.1_kernel_signed-lz4",
         },
         x86_64: {
             src: ":microdroid_gki-android14-6.1_kernel_signed",
@@ -605,21 +621,25 @@
     srcs: ["extract_microdroid_kernel_hashes.py"],
 }
 
-genrule {
+// HACK: use cc_genrule for arch-specific properties
+cc_genrule {
     name: "microdroid_kernel_hashes_rs",
-    srcs: [
-        ":microdroid_kernel",
-        ":microdroid_gki-android14-6.1_kernel",
-    ],
+    srcs: [":microdroid_kernel"],
+    arch: {
+        arm64: {
+            srcs: [":microdroid_gki-android14-6.1_kernel_signed"],
+        },
+        x86_64: {
+            srcs: [":microdroid_gki-android14-6.1_kernel_signed"],
+        },
+    },
     out: ["lib.rs"],
     tools: [
         "extract_microdroid_kernel_hashes",
         "avbtool",
     ],
     cmd: "$(location extract_microdroid_kernel_hashes) --avbtool $(location avbtool) " +
-        "--kernel $(location :microdroid_kernel) " +
-        "$(location :microdroid_gki-android14-6.1_kernel) " +
-        "> $(out)",
+        "--kernel $(in) > $(out)",
 }
 
 rust_library_rlib {
diff --git a/tests/hostside/Android.bp b/tests/hostside/Android.bp
index 13a9925..41d244d 100644
--- a/tests/hostside/Android.bp
+++ b/tests/hostside/Android.bp
@@ -35,6 +35,7 @@
         "initrd_bootconfig",
         "lpmake",
         "lpunpack",
+        "lz4",
         "sign_virt_apex",
         "simg2img",
         "dtdiff",
diff --git a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
index c40a312..4f502ab 100644
--- a/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
+++ b/tests/hostside/java/com/android/microdroid/test/MicrodroidHostTests.java
@@ -954,7 +954,43 @@
         assertThat(hasDebugPolicy).isFalse();
     }
 
+    private boolean isLz4(String path) throws Exception {
+        File lz4tool = findTestFile("lz4");
+        CommandResult result =
+                new RunUtil().runTimedCmd(5000, lz4tool.getAbsolutePath(), "-t", path);
+        return result.getStatus() == CommandStatus.SUCCESS;
+    }
+
+    private void decompressLz4(String inputPath, String outputPath) throws Exception {
+        File lz4tool = findTestFile("lz4");
+        CommandResult result =
+                new RunUtil()
+                        .runTimedCmd(
+                                5000, lz4tool.getAbsolutePath(), "-d", "-f", inputPath, outputPath);
+        String out = result.getStdout();
+        String err = result.getStderr();
+        assertWithMessage(
+                        "lz4 image "
+                                + inputPath
+                                + " decompression failed."
+                                + "\n\tout: "
+                                + out
+                                + "\n\terr: "
+                                + err
+                                + "\n")
+                .about(command_results())
+                .that(result)
+                .isSuccess();
+    }
+
     private String avbInfo(String image_path) throws Exception {
+        if (isLz4(image_path)) {
+            File decompressedImage = FileUtil.createTempFile("decompressed", ".img");
+            decompressedImage.deleteOnExit();
+            decompressLz4(image_path, decompressedImage.getAbsolutePath());
+            image_path = decompressedImage.getAbsolutePath();
+        }
+
         File avbtool = findTestFile("avbtool");
         List<String> command =
                 Arrays.asList(avbtool.getAbsolutePath(), "info_image", "--image", image_path);