kleaf: system_dlkm: Add flatten images support

Generate flatten images if kernel_images have
build_system_dlkm_flatten is set True. These
images don't have the `uname -r` in the directory
structure and all modules are available as
lib/modules/*.ko with modules.load and modules.dep
depmod artifacts.

Generates images as follow:

If system_dlkm_fs_type is set:
  system_dlkm.flatten.<system_dlkm_fs_type>.img

If system_dlkm_fs_types is set:
  system_dlkm.flatten.<fs_type>.img for each fs types

By default when no fs type is specified:
  system_dlkm.flatten.ext4.img

Bug: 301178664
Test: Manual image verification with different options
Test: TH
Change-Id: I4b7815d131d3d26adaa35c4550e76db71aa33936
Signed-off-by: Ramji Jiyani <ramjiyani@google.com>

3 files changed, 49 insertions, 6 deletions

diff --git a/build_utils.sh b/build_utils.sh
index 7feb3d3e..9a45d75f 100644
--- a/build_utils.sh
+++ b/build_utils.sh
@@ -254,6 +254,23 @@ function build_system_dlkm() {
 
   build_image "${SYSTEM_DLKM_STAGING_DIR}" "${system_dlkm_props_file}" \
     "${DIST_DIR}/${SYSTEM_DLKM_IMAGE_NAME}" /dev/null
+  local generated_images=(${SYSTEM_DLKM_IMAGE_NAME})
+
+  # Build flatten image as /lib/modules/*.ko; if unset or null: default false
+  if [[ ${SYSTEM_DLKM_GEN_FLATTEN_IMAGE:-0} == "1" ]]; then
+    local system_dlkm_flatten_image_name="system_dlkm.flatten.${SYSTEM_DLKM_FS_TYPE}.img"
+    mkdir -p ${SYSTEM_DLKM_STAGING_DIR}/flatten/lib/modules
+    cp $(find ${SYSTEM_DLKM_STAGING_DIR} -type f -name "*.ko") ${SYSTEM_DLKM_STAGING_DIR}/flatten/lib/modules
+    # Copy required depmod artifacts and scrub required files to correct paths
+    cp $(find ${SYSTEM_DLKM_STAGING_DIR} -name "modules.dep") ${SYSTEM_DLKM_STAGING_DIR}/flatten/lib/modules
+    sed -i 's/kernel[^:[:space:]]*\/\([^:[:space:]]*\.ko\)/\1/g' ${SYSTEM_DLKM_STAGING_DIR}/flatten/lib/modules/modules.dep
+    cp $(find ${SYSTEM_DLKM_STAGING_DIR} -name "modules.load") ${SYSTEM_DLKM_STAGING_DIR}/flatten/lib/modules
+    sed -i 's#.*/##' ${SYSTEM_DLKM_STAGING_DIR}/flatten/lib/modules/modules.load
+
+    build_image "${SYSTEM_DLKM_STAGING_DIR}/flatten" "${system_dlkm_props_file}" \
+    "${DIST_DIR}/${system_dlkm_flatten_image_name}" /dev/null
+    generated_images+=(${system_dlkm_flatten_image_name})
+   fi
 
   if [ -z "${SYSTEM_DLKM_PROPS}" ]; then
     rm ${system_dlkm_props_file}
@@ -261,10 +278,13 @@ function build_system_dlkm() {
   fi
 
   # No need to sign the image as modules are signed
-  avbtool add_hashtree_footer \
-    --partition_name system_dlkm \
-    --hash_algorithm sha256 \
-    --image "${DIST_DIR}/${SYSTEM_DLKM_IMAGE_NAME}"
+  for image in "${generated_images[@]}"
+  do
+    avbtool add_hashtree_footer \
+      --partition_name system_dlkm \
+      --hash_algorithm sha256 \
+      --image "${DIST_DIR}/${image}"
+  done
 
   # Archive system_dlkm_staging_dir
   tar -czf "${DIST_DIR}/system_dlkm_staging_archive.tar.gz" -C "${SYSTEM_DLKM_STAGING_DIR}" .
diff --git a/kleaf/impl/image/kernel_images.bzl b/kleaf/impl/image/kernel_images.bzl
index 70e758e9..07c68293 100644
--- a/kleaf/impl/image/kernel_images.bzl
+++ b/kleaf/impl/image/kernel_images.bzl
@@ -35,6 +35,7 @@ def kernel_images(
         build_vendor_boot = None,
         build_vendor_kernel_boot = None,
         build_system_dlkm = None,
+        build_system_dlkm_flatten = None,
         build_dtbo = None,
         dtbo_srcs = None,
         mkbootimg = None,
@@ -109,6 +110,8 @@ def kernel_images(
           - The list contains `dtb.img`
         build_initramfs: Whether to build initramfs. Keep in sync with `BUILD_INITRAMFS`.
         build_system_dlkm: Whether to build system_dlkm.img an image with GKI modules.
+        build_system_dlkm_flatten: Whether to build system_dlkm.flatten.<fs>.img.
+          This image have directory structure as `/lib/modules/*.ko` i.e. no `uname -r` in the path.
         build_vendor_dlkm: Whether to build `vendor_dlkm` image. It must be set if
           `vendor_dlkm_modules_list` is set.
 
@@ -371,6 +374,7 @@ def kernel_images(
             kernel_modules_install = kernel_modules_install,
             # For device system_dlkm, give GKI's system_dlkm_staging_archive.tar.gz
             base_kernel_images = base_kernel_images,
+            build_system_dlkm_flatten_image = build_system_dlkm_flatten,
             deps = deps,
             modules_list = modules_list,
             modules_blocklist = modules_blocklist,
diff --git a/kleaf/impl/image/system_dlkm_image.bzl b/kleaf/impl/image/system_dlkm_image.bzl
index f2462e25..849dd42c 100644
--- a/kleaf/impl/image/system_dlkm_image.bzl
+++ b/kleaf/impl/image/system_dlkm_image.bzl
@@ -110,10 +110,17 @@ def _system_dlkm_image_impl(ctx):
         else:
             system_dlkm_img = ctx.actions.declare_file("{}/system_dlkm.{}.img".format(ctx.label.name, fs_type))
             system_dlkm_img_name = "system_dlkm.{}.img".format(fs_type)
-
         outputs.append(system_dlkm_img)
         outputs_to_compare.append(system_dlkm_img_name)
 
+        system_dlkm_flatten_img = None
+        system_dlkm_flatten_img_name = None
+        if ctx.attr.build_system_dlkm_flatten_image:
+            system_dlkm_flatten_img = ctx.actions.declare_file("{}/system_dlkm.flatten.{}.img".format(ctx.label.name, fs_type))
+            outputs.append(system_dlkm_flatten_img)
+            system_dlkm_flatten_img_name = "system_dlkm.flatten.{}.img".format(fs_type)
+            outputs_to_compare.append(system_dlkm_flatten_img_name)
+
         command += """
                    {extract_staging_archive_cmd}
                  # Build {system_dlkm_img_name}
@@ -123,11 +130,15 @@ def _system_dlkm_image_impl(ctx):
                      SYSTEM_DLKM_FS_TYPE={system_dlkm_fs_type}
                      SYSTEM_DLKM_STAGING_DIR={system_dlkm_staging_dir}
                      SYSTEM_DLKM_IMAGE_NAME={system_dlkm_img_name}
+                     SYSTEM_DLKM_GEN_FLATTEN_IMAGE={build_flatten_image}
                      {extra_flags_cmd}
                      build_system_dlkm
                    )
                  # Move output files into place
                    mv "${{DIST_DIR}}/{system_dlkm_img_name}" {system_dlkm_img}
+                   if [[ {build_flatten_image} == "1" ]]; then
+                     mv "${{DIST_DIR}}/{system_dlkm_flatten_img_name}" {system_dlkm_flatten_img}
+                   fi
                    mv "${{DIST_DIR}}/system_dlkm.modules.load" {system_dlkm_modules_load}
                    mv "${{DIST_DIR}}/system_dlkm_staging_archive.tar.gz" {system_dlkm_staging_archive}
                    if [ -f "${{DIST_DIR}}/system_dlkm.modules.blocklist" ]; then
@@ -139,11 +150,14 @@ def _system_dlkm_image_impl(ctx):
                  # Remove staging directories
                    rm -rf {system_dlkm_staging_dir}
         """.format(
+            build_flatten_image = int(ctx.attr.build_system_dlkm_flatten_image),
             extract_staging_archive_cmd = extract_staging_archive_cmd,
             extra_flags_cmd = extra_flags_cmd,
             modules_staging_dir = modules_staging_dir,
             system_dlkm_fs_type = fs_type,
             system_dlkm_staging_dir = system_dlkm_staging_dir,
+            system_dlkm_flatten_img = system_dlkm_flatten_img.path if system_dlkm_flatten_img else "/dev/null",
+            system_dlkm_flatten_img_name = system_dlkm_flatten_img_name,
             system_dlkm_img = system_dlkm_img.path,
             system_dlkm_img_name = system_dlkm_img_name,
             system_dlkm_modules_load = system_dlkm_modules_load.path,
@@ -180,12 +194,17 @@ system_dlkm_image = rule(
 
 When included in a `copy_to_dist_dir` rule, this rule copies the following to `DIST_DIR`:
 - `system_dlkm.img` if system_dlkm_fs_type is specified
-- `system_dlkm.[erfos|ext4].img` if system_dlkm_fs_types is specified
+- `system_dlkm.[erofs|ext4].img` if system_dlkm_fs_types is specified
+- `system_dlkm.flatten.[erofs|ext4].img` if build_system_dlkm_flatten_image is True
 - `system_dlkm.modules.load`
 
 """,
     attrs = image_utils.build_modules_image_attrs_common({
         "base_kernel_images": attr.label(allow_files = True),
+        "build_system_dlkm_flatten_image": attr.bool(
+            default = False,
+            doc = "When True it builds system_dlkm image with no `uname -r` in the path.",
+        ),
         "modules_list": attr.label(allow_single_file = True),
         "modules_blocklist": attr.label(allow_single_file = True),
         "system_dlkm_fs_type": attr.string(doc = """Deprecated. system_dlkm image fs type""", values = ["ext4", "erofs"]),