Add GKI kernel to virt APEX

This kernel can be used instead of microdroid kernel.

Bug: 305118971
Test: build
Change-Id: Iafc51b4e86ad116b6531eb4e7e7603f5b1ef2705
diff --git a/apex/Android.bp b/apex/Android.bp
index a4c8861..b7fd67e 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -67,7 +67,15 @@
     ],
 }
 
-apex_defaults {
+soong_config_module_type {
+    name: "avf_flag_aware_apex_defaults",
+    module_type: "apex_defaults",
+    config_namespace: "ANDROID",
+    bool_variables: ["release_avf_enable_vendor_modules"],
+    properties: ["prebuilts"],
+}
+
+avf_flag_aware_apex_defaults {
     name: "com.android.virt_avf_enabled",
 
     defaults: ["com.android.virt_common"],
@@ -115,6 +123,15 @@
     apps: [
         "EmptyPayloadApp",
     ],
+    soong_config_variables: {
+        release_avf_enable_vendor_modules: {
+            prebuilts: [
+                "microdroid_gki_initrd_debuggable",
+                "microdroid_gki_initrd_normal",
+                "microdroid_gki_kernel",
+            ],
+        },
+    },
 }
 
 apex_defaults {
diff --git a/apex/sign_virt_apex.py b/apex/sign_virt_apex.py
index 029ac76..7393636 100644
--- a/apex/sign_virt_apex.py
+++ b/apex/sign_virt_apex.py
@@ -413,10 +413,13 @@
 # dict of (key, file) for re-sign/verification. keys are un-versioned for readability.
 virt_apex_files = {
     'kernel': 'etc/fs/microdroid_kernel',
+    'gki_kernel': 'etc/fs/microdroid_gki_kernel',
     'vbmeta.img': 'etc/fs/microdroid_vbmeta.img',
     'super.img': 'etc/fs/microdroid_super.img',
     'initrd_normal.img': 'etc/microdroid_initrd_normal.img',
+    'gki_initrd_normal.img': 'etc/microdroid_gki_initrd_normal.img',
     'initrd_debuggable.img': 'etc/microdroid_initrd_debuggable.img',
+    'gki_initrd_debuggable.img': 'etc/microdroid_gki_initrd_debuggable.img',
 }
 
 
@@ -458,26 +461,40 @@
                      images=images,
                      wait=images_f)
 
+    has_gki_kernel = os.path.isfile(files['gki_kernel'])
+
     vbmeta_bc_f = None
     if not args.do_not_update_bootconfigs:
-        vbmeta_bc_f = Async(UpdateVbmetaBootconfig, args,
-                            [files['initrd_normal.img'],
-                                files['initrd_debuggable.img']], files['vbmeta.img'],
+        initrd_files = [files['initrd_normal.img'], files['initrd_debuggable.img']]
+        if has_gki_kernel:
+            initrd_files += [files['gki_initrd_normal.img'], files['gki_initrd_debuggable.img']]
+        vbmeta_bc_f = Async(UpdateVbmetaBootconfig, args, initrd_files,
+                            files['vbmeta.img'],
                             wait=[vbmeta_f])
 
     # Re-sign kernel. Note kernel's vbmeta contain addition descriptor from ramdisk(s)
-    initrd_normal_hashdesc = tempfile.NamedTemporaryFile(delete=False).name
-    initrd_debug_hashdesc = tempfile.NamedTemporaryFile(delete=False).name
-    initrd_n_f = Async(GenVbmetaImage, args, files['initrd_normal.img'],
-                       initrd_normal_hashdesc, "initrd_normal",
-                       wait=[vbmeta_bc_f] if vbmeta_bc_f is not None else [])
-    initrd_d_f = Async(GenVbmetaImage, args, files['initrd_debuggable.img'],
-                       initrd_debug_hashdesc, "initrd_debug",
-                       wait=[vbmeta_bc_f] if vbmeta_bc_f is not None else [])
-    Async(AddHashFooter, args, key, files['kernel'], partition_name="boot",
-          additional_descriptors=[
-              initrd_normal_hashdesc, initrd_debug_hashdesc],
-          wait=[initrd_n_f, initrd_d_f])
+    def resign_kernel(kernel, initrd_normal, initrd_debug):
+        kernel_file = files[kernel]
+        initrd_normal_file = files[initrd_normal]
+        initrd_debug_file = files[initrd_debug]
+
+        initrd_normal_hashdesc = tempfile.NamedTemporaryFile(delete=False).name
+        initrd_debug_hashdesc = tempfile.NamedTemporaryFile(delete=False).name
+        initrd_n_f = Async(GenVbmetaImage, args, initrd_normal_file,
+                           initrd_normal_hashdesc, "initrd_normal",
+                           wait=[vbmeta_bc_f] if vbmeta_bc_f is not None else [])
+        initrd_d_f = Async(GenVbmetaImage, args, initrd_debug_file,
+                           initrd_debug_hashdesc, "initrd_debug",
+                           wait=[vbmeta_bc_f] if vbmeta_bc_f is not None else [])
+        Async(AddHashFooter, args, key, kernel_file, partition_name="boot",
+              additional_descriptors=[
+                  initrd_normal_hashdesc, initrd_debug_hashdesc],
+              wait=[initrd_n_f, initrd_d_f])
+
+    resign_kernel('kernel', 'initrd_normal.img', 'initrd_debuggable.img')
+
+    if has_gki_kernel:
+        resign_kernel('gki_kernel', 'gki_initrd_normal.img', 'gki_initrd_debuggable.img')
 
 
 def VerifyVirtApex(args):
@@ -502,7 +519,8 @@
         assert info['Public key (sha1)'] == pubkey_digest, f'pubkey mismatch: {file}'
 
     for f in files.values():
-        if f in (files['initrd_normal.img'], files['initrd_debuggable.img']):
+        if f in (files['initrd_normal.img'], files['initrd_debuggable.img'],
+                 files['gki_initrd_normal.img'], files['gki_initrd_debuggable.img']):
             # TODO(b/245277660): Verify that ramdisks contain the correct vbmeta digest
             continue
         if f == files['super.img']:
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 4e735e6..a1ce594 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -330,6 +330,22 @@
     ],
 }
 
+android_filesystem {
+    name: "microdroid_gki_modules-6.1-arm64",
+    deps: [
+        "microdroid_gki_kernel_modules-6.1-arm64",
+    ],
+    type: "compressed_cpio",
+}
+
+android_filesystem {
+    name: "microdroid_gki_modules-6.1-x86_64",
+    deps: [
+        "microdroid_gki_kernel_modules-6.1-x86_64",
+    ],
+    type: "compressed_cpio",
+}
+
 genrule {
     name: "microdroid_bootconfig_arm64_gen",
     srcs: [
@@ -444,6 +460,23 @@
     },
 }
 
+avb_gen_vbmeta_image {
+    name: "microdroid_gki_initrd_normal_hashdesc",
+    src: ":microdroid_gki_initrd_normal",
+    partition_name: "initrd_normal",
+    salt: initrd_normal_salt,
+    enabled: false,
+    arch: {
+        // Microdroid kernel is only available in these architectures.
+        arm64: {
+            enabled: true,
+        },
+        x86_64: {
+            enabled: true,
+        },
+    },
+}
+
 // python -c "import hashlib; print(hashlib.sha256(b'initrd_debug').hexdigest())"
 initrd_debug_salt = "8ab9dc9cb7e6456700ff6ef18c6b4c3acc24c5fa5381b829563f8d7a415d869a"
 
@@ -464,6 +497,23 @@
     },
 }
 
+avb_gen_vbmeta_image {
+    name: "microdroid_gki_initrd_debug_hashdesc",
+    src: ":microdroid_gki_initrd_debuggable",
+    partition_name: "initrd_debug",
+    salt: initrd_debug_salt,
+    enabled: false,
+    arch: {
+        // Microdroid kernel is only available in these architectures.
+        arm64: {
+            enabled: true,
+        },
+        x86_64: {
+            enabled: true,
+        },
+    },
+}
+
 soong_config_module_type {
     name: "flag_aware_avb_add_hash_footer",
     module_type: "avb_add_hash_footer",
@@ -513,6 +563,42 @@
     },
 }
 
+flag_aware_avb_add_hash_footer {
+    name: "microdroid_gki_kernel_signed",
+    src: ":empty_file",
+    filename: "microdroid_gki_kernel",
+    partition_name: "boot",
+    private_key: ":microdroid_sign_key",
+    salt: bootloader_salt,
+    enabled: false,
+    arch: {
+        arm64: {
+            src: ":microdroid_gki_kernel_prebuilts-6.1-arm64",
+            enabled: true,
+        },
+        x86_64: {
+            src: ":microdroid_gki_kernel_prebuilts-6.1-x86_64",
+            enabled: true,
+        },
+    },
+    include_descriptors_from_images: [
+        ":microdroid_gki_initrd_normal_hashdesc",
+        ":microdroid_gki_initrd_debug_hashdesc",
+    ],
+    // Below are properties that are conditionally set depending on value of build flags.
+    soong_config_variables: {
+        release_avf_enable_llpvm_changes: {
+            rollback_index: 1,
+            props: [
+                {
+                    name: "com.android.virt.cap",
+                    value: "secretkeeper_protection",
+                },
+            ],
+        },
+    },
+}
+
 prebuilt_etc {
     name: "microdroid_kernel",
     src: ":empty_file",
@@ -526,3 +612,17 @@
         },
     },
 }
+
+prebuilt_etc {
+    name: "microdroid_gki_kernel",
+    src: ":empty_file",
+    relative_install_path: "fs",
+    arch: {
+        arm64: {
+            src: ":microdroid_gki_kernel_signed",
+        },
+        x86_64: {
+            src: ":microdroid_gki_kernel_signed",
+        },
+    },
+}
diff --git a/microdroid/initrd/Android.bp b/microdroid/initrd/Android.bp
index de28d8a..6cd84fa 100644
--- a/microdroid/initrd/Android.bp
+++ b/microdroid/initrd/Android.bp
@@ -40,6 +40,28 @@
     cmd: "cat $(in) > $(out)",
 }
 
+genrule {
+    name: "microdroid_gki_initrd_gen_arm64",
+    srcs: [
+        ":microdroid_ramdisk",
+        ":microdroid_fstab_ramdisk",
+        ":microdroid_gki_modules-6.1-arm64",
+    ],
+    out: ["microdroid_initrd.img"],
+    cmd: "cat $(in) > $(out)",
+}
+
+genrule {
+    name: "microdroid_gki_initrd_gen_x86_64",
+    srcs: [
+        ":microdroid_ramdisk",
+        ":microdroid_fstab_ramdisk",
+        ":microdroid_gki_modules-6.1-x86_64",
+    ],
+    out: ["microdroid_initrd.img"],
+    cmd: "cat $(in) > $(out)",
+}
+
 // This contains vbmeta hashes & related (boot)configs which are passed to kernel/init
 genrule {
     name: "microdroid_vbmeta_bootconfig_gen",
@@ -74,6 +96,17 @@
 }
 
 genrule {
+    name: "microdroid_gki_initrd_debuggable_arm64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_gki_initrd_gen_arm64",
+        ":microdroid_bootconfig_debuggable_src",
+    ] + bootconfigs_arm64,
+    out: ["microdroid_gki_initrd_debuggable_arm64"],
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
+}
+
+genrule {
     name: "microdroid_initrd_debuggable_x86_64",
     tools: ["initrd_bootconfig"],
     srcs: [
@@ -85,6 +118,17 @@
 }
 
 genrule {
+    name: "microdroid_gki_initrd_debuggable_x86_64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_gki_initrd_gen_x86_64",
+        ":microdroid_bootconfig_debuggable_src",
+    ] + bootconfigs_x86_64,
+    out: ["microdroid_gki_initrd_debuggable_x86_64"],
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
+}
+
+genrule {
     name: "microdroid_initrd_normal_arm64",
     tools: ["initrd_bootconfig"],
     srcs: [
@@ -96,6 +140,17 @@
 }
 
 genrule {
+    name: "microdroid_gki_initrd_normal_arm64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_gki_initrd_gen_arm64",
+        ":microdroid_bootconfig_normal_src",
+    ] + bootconfigs_arm64,
+    out: ["microdroid_gki_initrd_normal_arm64"],
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
+}
+
+genrule {
     name: "microdroid_initrd_normal_x86_64",
     tools: ["initrd_bootconfig"],
     srcs: [
@@ -106,6 +161,17 @@
     cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
+genrule {
+    name: "microdroid_gki_initrd_normal_x86_64",
+    tools: ["initrd_bootconfig"],
+    srcs: [
+        ":microdroid_gki_initrd_gen_x86_64",
+        ":microdroid_bootconfig_normal_src",
+    ] + bootconfigs_x86_64,
+    out: ["microdroid_gki_initrd_normal_x86_64"],
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
+}
+
 prebuilt_etc {
     name: "microdroid_initrd_debuggable",
     // We don't have ramdisk for architectures other than x86_64 & arm64
@@ -122,6 +188,21 @@
 }
 
 prebuilt_etc {
+    name: "microdroid_gki_initrd_debuggable",
+    // We don't have ramdisk for architectures other than x86_64 & arm64
+    src: ":empty_file",
+    arch: {
+        x86_64: {
+            src: ":microdroid_gki_initrd_debuggable_x86_64",
+        },
+        arm64: {
+            src: ":microdroid_gki_initrd_debuggable_arm64",
+        },
+    },
+    filename: "microdroid_gki_initrd_debuggable.img",
+}
+
+prebuilt_etc {
     name: "microdroid_initrd_normal",
     // We don't have ramdisk for architectures other than x86_64 & arm64
     src: ":empty_file",
@@ -135,3 +216,18 @@
     },
     filename: "microdroid_initrd_normal.img",
 }
+
+prebuilt_etc {
+    name: "microdroid_gki_initrd_normal",
+    // We don't have ramdisk for architectures other than x86_64 & arm64
+    src: ":empty_file",
+    arch: {
+        x86_64: {
+            src: ":microdroid_gki_initrd_normal_x86_64",
+        },
+        arm64: {
+            src: ":microdroid_gki_initrd_normal_arm64",
+        },
+    },
+    filename: "microdroid_gki_initrd_normal.img",
+}