kleaf: Add support for specifying modules for recovery/charger

Currently, the same modules.load file is used for normal boot,
recovery/fastbootd and charger modes during first-stage init. This
means that during normal boot, the modules that are needed for
recovery and charger modes are loaded during first-stage init, which
is not required.

Add support for specifying a list of modules that are required for
booting into recovery or charger mode. This introduces two new
attributes to the kernel_images() macro called: modules_recovery_list
and modules_charger_list. Each attribute points to a list of modules
that are needed to boot into each mode respectively. This list can
be generated dynamically by concatenating the list of vendor_boot or
vendor_kernel_boot modules required for first stage init, along with
the list of modules needed to boot into each mode.

Defining these attributes also produces a file named
modules.load.[recovery/charger] in the initramfs, which init uses
instead of modules.load when booting into recovery or charger mode
respectively to load the first-stage init and recovery and charger
modules. In case of normal boot, init will use the modules.load
file, which will only contain the modules needed for first
stage init, and the recovery and charger modules will be loaded
from the vendor_dlkm partition during second-stage init.

Bug: 266752750
Change-Id: I11503b11683ef64cf0933b8641959ea5acd63ab1
Signed-off-by: Isaac J. Manjarres <isaacmanjarres@google.com>
(cherry picked from commit b9ba4b67e90b85d01eb2f9cebf62830a73574cd9)
[isaacmanjarres: Re-wrote some of the commit message/implemented it
as part of kleaf. Also added tests for the new attributes.]
(cherry picked from commit d08605f4c795e05d8569e975affd2736e430a555)
[isaacmanjarres: Resolved trivial merge conflicts when backporting
from main.]

8 files changed, 413 insertions, 49 deletions

diff --git a/build_utils.sh b/build_utils.sh
index de5251a2..3e8ddf8d 100644
--- a/build_utils.sh
+++ b/build_utils.sh
@@ -25,14 +25,27 @@ function rel_path() {
 # $1 directory of kernel modules ($1/lib/modules/x.y)
 # $2 flags to pass to depmod
 # $3 kernel version
+# $4 Optional: File with list of modules to run depmod on.
+#              If left empty, depmod will run on all modules
+#              under $1/lib/modules/x.y
 function run_depmod() {
   (
     local ramdisk_dir=$1
     local depmod_stdout
     local depmod_stderr=$(mktemp)
+    local version=$3
+    local modules_list_file=$4
+    local modules_list=""
+
+    if [[ -n "${modules_list_file}" ]]; then
+      while read -r line; do
+        # depmod expects absolute paths for module files
+        modules_list+="$(realpath ${ramdisk_dir}/lib/modules/${version}/${line}) "
+      done <${modules_list_file}
+    fi
 
     cd ${ramdisk_dir}
-    if ! depmod_stdout="$(depmod $2 -F ${DIST_DIR}/System.map -b . $3 \
+    if ! depmod_stdout="$(depmod $2 -F ${DIST_DIR}/System.map -b . ${version} ${modules_list} \
         2>${depmod_stderr})"; then
       echo "$depmod_stdout"
       cat ${depmod_stderr} >&2
@@ -50,12 +63,116 @@ function run_depmod() {
   )
 }
 
+# $1 MODULES_LIST, <File containing the list of modules that should go in the
+#                   ramdisk.>
+# $2 MODULES_RECOVERY_LIST, <File containing the list of modules that should
+#                            go in the ramdisk and be loaded when booting into
+#                            recovery mode during first stage init.
+#
+#                            This parameter is optional, and if not used, should
+#                            be passed as an empty string to ensure that
+#                            subsequent parameters are treated correctly.>
+# $3 MODULES_CHARGER_LIST, <File containing the list of modules that should
+#                           go in the ramdisk and be loaded when booting into
+#                           charger mode during first stage init.
+#
+#                           This parameter is optional, and if not used, should
+#                           be passed as an empty string to ensure that
+#                           subsequent paratmers are treated correctly.>
+# $4 MODULES_ORDER_LIST, <The modules.order file that contains all of the
+#                         modules that were built.>
+#
+# This function creates new modules.order* files by filtering the module lists
+# through the set of modules that were built ($MODULES_ORDER_LIST).
+#
+# Each modules.order* file is created by filtering each list as follows:
+#
+# Let f be the filter_module_list function, which filters arg 1 through arg 2.
+#
+# f(MODULES_LIST, MODULES_ORDER_LIST) ==> modules.order
+#
+# f(MODULES_RECOVERY_LIST, MODULES_ORDER_LIST) ==> modules.order.recovery
+#
+# f(MODULES_CHARGER_LIST, MODULES_ORDER_LIST) ==> modules.order.charger
+#
+# Filtering ensures that only the modules in MODULES_LIST end up in the
+# respective partition that create_modules_staging() is invoked for.
+#
+# Note: This function overwrites the original file pointed to by
+# MODULES_ORDER_LIST when MODULES_LIST is set.
+function create_modules_order_lists() {
+  local modules_list_file="${1}"
+  local modules_recovery_list_file="${2}"
+  local modules_charger_list_file="${3}"
+  local modules_order_list_file="${4}"
+  local dest_dir=$(dirname $(realpath ${modules_order_list_file}))
+  local tmp_modules_order_file=$(mktemp)
+
+  cp ${modules_order_list_file} ${tmp_modules_order_file}
+
+  declare -A module_lists_arr
+  module_lists_arr["modules.order"]=${modules_list_file}
+  module_lists_arr["modules.order.recovery"]=${modules_recoverylist_file}
+  module_lists_arr["modules.order.charger"]=${modules_chargerlist_file}
+
+  for mod_order_file in ${!module_lists_arr[@]}; do
+    local mod_list_file=${module_lists_arr[${mod_order_file}]}
+    local dest_file=${dest_dir}/${mod_order_file}
+
+    # Need to make sure we can find modules_list_file from the staging dir
+    if [[ -n "${mod_list_file}" ]]; then
+      if [[ -f "${ROOT_DIR}/${mod_list_file}" ]]; then
+        modules_list_file="${ROOT_DIR}/${mod_list_file}"
+      elif [[ "${mod_list_file}" != /* ]]; then
+        echo "ERROR: modules list must be an absolute path or relative to ${ROOT_DIR}: ${mod_list_file}" >&2
+        rm -f ${tmp_modules_order_file}
+        exit 1
+      elif [[ ! -f "${mod_list_file}" ]]; then
+        echo "ERROR: Failed to find modules list: ${mod_list_file}" >&2
+        rm -f ${tmp_modules_order_file}
+        exit 1
+      fi
+
+      local modules_list_filter=$(mktemp)
+
+      # Remove all lines starting with "#" (comments)
+      # Exclamation point makes interpreter ignore the exit code under set -e
+      ! grep -v "^\#" ${mod_list_file} > ${modules_list_filter}
+
+      # Append a new line at the end of file
+      # If file doesn't end in newline the last module is skipped from filter
+      echo >> ${modules_list_filter}
+
+      # grep the modules.order for any KOs in the modules list
+      ! grep -w -f ${modules_list_filter} ${tmp_modules_order_file} > ${dest_file}
+
+      rm -f ${modules_list_filter}
+    fi
+  done
+
+  rm -f ${tmp_modules_order_file}
+}
+
 # $1 MODULES_LIST, <File contains the list of modules that should go in the ramdisk>
 # $2 MODULES_STAGING_DIR    <The directory to look for all the compiled modules>
 # $3 IMAGE_STAGING_DIR  <The destination directory in which MODULES_LIST is
 #                        expected, and it's corresponding modules.* files>
 # $4 MODULES_BLOCKLIST, <File contains the list of modules to prevent from loading>
-# $5 flags to pass to depmod
+# $5 MODULES_RECOVERY_LIST <File contains the list of modules that should go in
+#                           the ramdisk but should only be loaded when booting
+#                           into recovery.
+#
+#                           This parameter is optional, and if not used, should
+#                           be passed as an empty string to ensure that the depmod
+#                           flags are assigned correctly.>
+# $6 MODULES_CHARGER_LIST <File contains the list of modules that should go in
+#                          the ramdisk but should only be loaded when booting
+#                          into charger mode.
+#
+#                          This parameter is optional, and if not used, should
+#                          be passed as an empty string to ensure that the
+#                          depmod flags are assigned correctly.>
+# $7 flags to pass to depmod
 function create_modules_staging() {
   local modules_list_file=$1
   local src_dir=$(echo $2/lib/modules/*)
@@ -63,7 +180,9 @@ function create_modules_staging() {
   local dest_dir=$3/lib/modules/${version}
   local dest_stage=$3
   local modules_blocklist_file=$4
-  local depmod_flags=$5
+  local modules_recoverylist_file=$5
+  local modules_chargerlist_file=$6
+  local depmod_flags=$7
 
   rm -rf ${dest_dir}
   mkdir -p ${dest_dir}/kernel
@@ -117,34 +236,10 @@ function create_modules_staging() {
       -exec ${OBJCOPY:-${CROSS_COMPILE}objcopy} --strip-debug {} \;
   fi
 
-  if [ -n "${modules_list_file}" ]; then
-    # Need to make sure we can find modules_list_file from the staging dir
-    if [[ -f "${ROOT_DIR}/${modules_list_file}" ]]; then
-      modules_list_file="${ROOT_DIR}/${modules_list_file}"
-    elif [[ "${modules_list_file}" != /* ]]; then
-      echo "modules list must be an absolute path or relative to ${ROOT_DIR}: ${modules_list_file}"
-      exit 1
-    elif [[ ! -f "${modules_list_file}" ]]; then
-      echo "Failed to find modules list: ${modules_list_file}"
-      exit 1
-    fi
-
-    local modules_list_filter=$(mktemp)
-    local old_modules_list=$(mktemp)
-
-    # Remove all lines starting with "#" (comments)
-    # Exclamation point makes interpreter ignore the exit code under set -e
-    ! grep -v "^\#" ${modules_list_file} > ${modules_list_filter}
-
-    # Append a new line at the end of file
-    # If file doesn't end in newline the last module is skipped from filter
-    echo >> ${modules_list_filter}
-
-    # grep the modules.order for any KOs in the modules list
-    cp ${dest_dir}/modules.order ${old_modules_list}
-    ! grep -w -f ${modules_list_filter} ${old_modules_list} > ${dest_dir}/modules.order
-    rm -f ${modules_list_filter} ${old_modules_list}
-  fi
+  # create_modules_order_lists() will overwrite modules.order if MODULES_LIST is
+  # set.
+  create_modules_order_lists "${modules_list_file:-""}" "${modules_recovery_list_file:-""}" \
+	                     "${modules_charger_list_file:-""}" ${dest_dir}/modules.order
 
   if [ -n "${modules_blocklist_file}" ]; then
     # Need to make sure we can find modules_blocklist_file from the staging dir
@@ -173,15 +268,39 @@ function create_modules_staging() {
     # Trim modules from tree that aren't mentioned in modules.order
     (
       cd ${dest_dir}
-      find * -type f -name "*.ko" | (grep -v -w -f modules.order -f $used_blocklist_modules - || true) | xargs -r rm
+      local grep_flags="-v -w -f modules.order -f ${used_blocklist_modules} "
+      if [[ -f modules.order.recovery ]]; then
+        grep_flags+="-f modules.order.recovery "
+      fi
+      if [[ -f modules.order.charger ]]; then
+        grep_flags+="-f modules.order.charger "
+      fi
+      find * -type f -name "*.ko" | (grep ${grep_flags} - || true) | xargs -r rm
     )
     rm $used_blocklist_modules
   fi
 
   # Re-run depmod to detect any dependencies between in-kernel and external
-  # modules. Then, create modules.order based on all the modules compiled.
-  run_depmod ${dest_stage} "${depmod_flags}" "${version}"
-  cp ${dest_dir}/modules.order ${dest_dir}/modules.load
+  # modules, as well as recovery and charger modules. Then, create the
+  # modules.order files based on all the modules compiled.
+  declare -A module_load_lists_arr
+  module_load_lists_arr["modules.order"]="modules.load"
+  module_load_lists_arr["modules.order.recovery"]="modules.load.recovery"
+  module_load_lists_arr["modules.order.charger"]="modules.load.charger"
+
+  for mod_order_file in ${!module_load_lists_arr[@]}; do
+    local mod_order_filepath=${dest_dir}/${mod_order_file}
+    local mod_load_filepath=${dest_dir}/${module_load_lists_arr[${mod_order_file}]}
+
+    if [[ -f ${mod_order_filepath} ]]; then
+      if [[ "${mod_order_file}" == "modules.order" ]]; then
+        run_depmod ${dest_stage} "${depmod_flags}" "${version}"
+      else
+        run_depmod ${dest_stage} "${depmod_flags}" "${version}" "${mod_order_filepath}"
+      fi
+      cp ${mod_order_filepath} ${mod_load_filepath}
+    fi
+  done
 }
 
 function build_system_dlkm() {
@@ -190,7 +309,8 @@ function build_system_dlkm() {
 
   rm -rf ${SYSTEM_DLKM_STAGING_DIR}
   create_modules_staging "${SYSTEM_DLKM_MODULES_LIST:-${MODULES_LIST}}" "${MODULES_STAGING_DIR}" \
-    ${SYSTEM_DLKM_STAGING_DIR} "${SYSTEM_DLKM_MODULES_BLOCKLIST:-${MODULES_BLOCKLIST}}" "-e"
+    ${SYSTEM_DLKM_STAGING_DIR} "${SYSTEM_DLKM_MODULES_BLOCKLIST:-${MODULES_BLOCKLIST}}" \
+    "${MODULES_RECOVERY_LIST:-""}" "${MODULES_CHARGER_LIST:-""}" "-e"
 
   local system_dlkm_root_dir=$(echo ${SYSTEM_DLKM_STAGING_DIR}/lib/modules/*)
   cp ${system_dlkm_root_dir}/modules.load ${DIST_DIR}/system_dlkm.modules.load
diff --git a/kleaf/artifact_tests/BUILD.bazel b/kleaf/artifact_tests/BUILD.bazel
index 30306ef3..a0368f7c 100644
--- a/kleaf/artifact_tests/BUILD.bazel
+++ b/kleaf/artifact_tests/BUILD.bazel
@@ -79,3 +79,14 @@ py_binary(
         "@io_abseil_py//absl/testing:absltest",
     ],
 )
+
+py_binary(
+    name = "initramfs_modules_lists_test",
+    srcs = ["initramfs_modules_lists_test.py"],
+    python_version = "PY3",
+    # All packages that uses kernel_module must be able to see this.
+    visibility = ["//visibility:public"],
+    deps = [
+        "@io_abseil_py//absl/testing:absltest",
+    ],
+)
diff --git a/kleaf/artifact_tests/initramfs_modules_lists_test.py b/kleaf/artifact_tests/initramfs_modules_lists_test.py
new file mode 100644
index 00000000..ac31d0c8
--- /dev/null
+++ b/kleaf/artifact_tests/initramfs_modules_lists_test.py
@@ -0,0 +1,102 @@
+# Copyright (C) 2023 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import subprocess
+import sys
+import tempfile
+import unittest
+
+from absl.testing import absltest
+
+
+def load_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--expected_modules_list")
+    parser.add_argument("--expected_modules_recovery_list")
+    parser.add_argument("--expected_modules_charger_list")
+    parser.add_argument("files", nargs="*", default=[])
+    return parser.parse_known_args()
+
+
+arguments = None
+
+
+class InitramfsModulesLists(unittest.TestCase):
+    def _decompress_initramfs(self, initramfs, temp_dir):
+        """Decompress initramfs into temp_dir.
+
+        Args:
+            initramfs: path to initramfs.img gzip file to be decompressed
+            temp_dir: directory in which to decompress initramfs.img into
+        """
+        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)
+
+    def _diff_modules_lists(self, modules_lists_map, modules_dir):
+        """Compares generated modules lists against expected modules lists for equality.
+
+        Given a dictionary of modules.load* files as keys, and expected modules
+        lists as values, compares each key value pair for equality.
+
+        Args:
+            modules_lists_map: dictionary with modules lists to compare
+            modules_dir:       directory that contains the modules.load* files
+        """
+        for modules_load, expected_modules_list_path in modules_lists_map.items():
+            modules_load_path = os.path.join(modules_dir, modules_load)
+            self.assertTrue(os.path.isfile(modules_load_path), f"Can't find {modules_load_path}")
+
+            with open(modules_load_path) as modules_load_file, \
+                 open(expected_modules_list_path) as expected_modules_list_file:
+                modules_load_lines = [os.path.basename(f) for f in modules_load_file.readlines()]
+                expected_modules_list_lines = expected_modules_list_file.readlines()
+                self.assertEqual(modules_load_lines.sort(), expected_modules_list_lines.sort())
+
+    def test_diff(self):
+        initramfs_list = [f for f in arguments.files if os.path.basename(f) == "initramfs.img"]
+        self.assertEqual(1, len(initramfs_list))
+        initramfs = initramfs_list[0]
+        modules_lists_map = {}
+
+        if arguments.expected_modules_list:
+            modules_lists_map["modules.load"] = arguments.expected_modules_list
+
+        if arguments.expected_modules_recovery_list:
+            modules_lists_map["modules.load.recovery"] = arguments.expected_modules_recovery_list
+
+        if arguments.expected_modules_charger_list:
+            modules_lists_map["modules.load.charger"] = arguments.expected_modules_charger_list
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            self._decompress_initramfs(initramfs, temp_dir)
+
+            lib_modules = os.path.join(temp_dir, "lib/modules")
+            self.assertTrue(os.path.isdir(lib_modules))
+
+            kernel_versions = os.listdir(lib_modules)
+            for v in kernel_versions:
+                modules_dir = os.path.join(lib_modules, v)
+                self._diff_modules_lists(modules_lists_map, modules_dir)
+
+if __name__ == '__main__':
+    arguments, unknown = load_arguments()
+    sys.argv[1:] = unknown
+    absltest.main()
diff --git a/kleaf/artifact_tests/kernel_test.bzl b/kleaf/artifact_tests/kernel_test.bzl
index cf76035b..2631ab61 100644
--- a/kleaf/artifact_tests/kernel_test.bzl
+++ b/kleaf/artifact_tests/kernel_test.bzl
@@ -117,3 +117,61 @@ def initramfs_modules_options_test(
         timeout = "short",
         **kwargs
     )
+
+def initramfs_modules_lists_test(
+        name,
+        kernel_images,
+        expected_modules_list = None,
+        expected_modules_recovery_list = None,
+        expected_modules_charger_list = None,
+        **kwargs):
+    """Tests that the initramfs has modules.load* files with the given content.
+
+    Args:
+        name: name of the test
+        kernel_images: name of the `kernel_images` target. It must build initramfs.
+        expected_modules_list: file with the expected content for `modules.load`
+        expected_modules_recovery_list: file with the expected content for `modules.load.recovery`
+        expected_modules_charger_list: file with the expected content for `modules.load.charger`
+        **kwargs: Additional attributes to the internal rule, e.g.
+          [`visibility`](https://docs.bazel.build/versions/main/visibility.html).
+          See complete list
+          [here](https://docs.bazel.build/versions/main/be/common-definitions.html#common-attributes).
+    """
+    test_binary = Label("//build/kernel/kleaf/artifact_tests:initramfs_modules_lists_test")
+    args = []
+
+    if expected_modules_list:
+        args += [
+            "--expected_modules_list",
+            "$(rootpath {})".format(expected_modules_list),
+        ]
+
+    if expected_modules_recovery_list:
+        args += [
+            "--expected_modules_recovery_list",
+            "$(rootpath {})".format(expected_modules_recovery_list),
+        ]
+
+    if expected_modules_charger_list:
+        args += [
+            "--expected_modules_charger_list",
+            "$(rootpath {})".format(expected_modules_charger_list),
+        ]
+
+    args.append("$(rootpaths {})".format(kernel_images))
+
+    hermetic_exec_test(
+        name = name,
+        data = [
+            expected_modules_list,
+            expected_modules_recovery_list,
+            expected_modules_charger_list,
+            kernel_images,
+            test_binary,
+        ],
+        script = run_py_binary_cmd(test_binary),
+        args = args,
+        timeout = "short",
+        **kwargs
+    )
diff --git a/kleaf/impl/image/image_utils.bzl b/kleaf/impl/image/image_utils.bzl
index 2f53f340..c577e9d9 100644
--- a/kleaf/impl/image/image_utils.bzl
+++ b/kleaf/impl/image/image_utils.bzl
@@ -116,6 +116,8 @@ def _build_modules_image_impl_common(
 
     for attr_name in (
         "modules_list",
+        "modules_recovery_list",
+        "modules_charger_list",
         "modules_blocklist",
         "vendor_dlkm_fs_type",
         "vendor_dlkm_modules_list",
diff --git a/kleaf/impl/image/initramfs.bzl b/kleaf/impl/image/initramfs.bzl
index bf5cf095..783f942a 100644
--- a/kleaf/impl/image/initramfs.bzl
+++ b/kleaf/impl/image/initramfs.bzl
@@ -31,6 +31,8 @@ def _initramfs_impl(ctx):
     initramfs_img = ctx.actions.declare_file("{}/initramfs.img".format(ctx.label.name))
     modules_load = ctx.actions.declare_file("{}/modules.load".format(ctx.label.name))
     vendor_boot_modules_load = ctx.outputs.vendor_boot_modules_load
+    vendor_boot_modules_load_recovery = ctx.outputs.vendor_boot_modules_load_recovery
+    vendor_boot_modules_load_charger = ctx.outputs.vendor_boot_modules_load_charger
     initramfs_staging_archive = ctx.actions.declare_file("{}/initramfs_staging_archive.tar.gz".format(ctx.label.name))
 
     outputs = [
@@ -40,6 +42,12 @@ def _initramfs_impl(ctx):
     if vendor_boot_modules_load:
         outputs.append(vendor_boot_modules_load)
 
+    if vendor_boot_modules_load_recovery:
+        outputs.append(vendor_boot_modules_load_recovery)
+
+    if vendor_boot_modules_load_charger:
+        outputs.append(vendor_boot_modules_load_charger)
+
     modules_staging_dir = initramfs_img.dirname + "/staging"
     initramfs_staging_dir = modules_staging_dir + "/initramfs_staging"
 
@@ -51,6 +59,30 @@ def _initramfs_impl(ctx):
             vendor_boot_modules_load = vendor_boot_modules_load.path,
         )
 
+    cp_vendor_boot_modules_load_recovery_cmd = ""
+    if vendor_boot_modules_load_recovery:
+        modules_load_recovery = ctx.actions.declare_file("{}/modules.load.recovery".format(ctx.label.name))
+        outputs.append(modules_load_recovery)
+        cp_vendor_boot_modules_load_recovery_cmd = """
+               cp ${{modules_root_dir}}/modules.load.recovery {modules_load_recovery}
+               cp ${{modules_root_dir}}/modules.load.recovery {vendor_boot_modules_load_recovery}
+        """.format(
+            modules_load_recovery = modules_load_recovery.path,
+            vendor_boot_modules_load_recovery = vendor_boot_modules_load_recovery.path,
+        )
+
+    cp_vendor_boot_modules_load_charger_cmd = ""
+    if vendor_boot_modules_load_charger:
+        modules_load_charger = ctx.actions.declare_file("{}/modules.load.charger".format(ctx.label.name))
+        outputs.append(modules_load_charger)
+        cp_vendor_boot_modules_load_charger_cmd = """
+               cp ${{modules_root_dir}}/modules.load.charger {modules_load_charger}
+               cp ${{modules_root_dir}}/modules.load.charger {vendor_boot_modules_load_charger}
+        """.format(
+            modules_load_charger = modules_load_charger.path,
+            vendor_boot_modules_load_charger = vendor_boot_modules_load_charger.path,
+        )
+
     additional_inputs = []
     if ctx.file.modules_options:
         cp_modules_options_cmd = """
@@ -75,10 +107,13 @@ def _initramfs_impl(ctx):
                mkdir -p {initramfs_staging_dir}
              # Build initramfs
                create_modules_staging "${{MODULES_LIST}}" {modules_staging_dir} \
-                 {initramfs_staging_dir} "${{MODULES_BLOCKLIST}}" "-e"
+                       {initramfs_staging_dir} "${{MODULES_BLOCKLIST}}" \
+                       "${{MODULES_RECOVERY_LIST:-""}}" "${{MODULES_CHARGER_LIST:-""}}" "-e"
                modules_root_dir=$(readlink -e {initramfs_staging_dir}/lib/modules/*) || exit 1
                cp ${{modules_root_dir}}/modules.load {modules_load}
                {cp_vendor_boot_modules_load_cmd}
+               {cp_vendor_boot_modules_load_recovery_cmd}
+               {cp_vendor_boot_modules_load_charger_cmd}
                {cp_modules_options_cmd}
                mkbootfs "{initramfs_staging_dir}" >"{modules_staging_dir}/initramfs.cpio"
                {ramdisk_compress} "{modules_staging_dir}/initramfs.cpio" >"{initramfs_img}"
@@ -94,6 +129,8 @@ def _initramfs_impl(ctx):
         initramfs_img = initramfs_img.path,
         initramfs_staging_archive = initramfs_staging_archive.path,
         cp_vendor_boot_modules_load_cmd = cp_vendor_boot_modules_load_cmd,
+        cp_vendor_boot_modules_load_recovery_cmd = cp_vendor_boot_modules_load_recovery_cmd,
+        cp_vendor_boot_modules_load_charger_cmd = cp_vendor_boot_modules_load_charger_cmd,
         cp_modules_options_cmd = cp_modules_options_cmd,
     )
 
@@ -125,7 +162,11 @@ initramfs = rule(
 When included in a `copy_to_dist_dir` rule, this rule copies the following to `DIST_DIR`:
 - `initramfs.img`
 - `modules.load`
+- `modules.load.recovery`
+- `modules.load.charger`
 - `vendor_boot.modules.load`
+- `vendor_boot.modules.load.recovery`
+- `vendor_boot.modules.load.charger`
 
 An additional label, `{name}/vendor_boot.modules.load`, is declared to point to the
 corresponding files.
@@ -134,7 +175,15 @@ corresponding files.
         "vendor_boot_modules_load": attr.output(
             doc = "`vendor_boot.modules.load` or `vendor_kernel_boot.modules.load`",
         ),
+        "vendor_boot_modules_load_recovery": attr.output(
+            doc = "`vendor_boot.modules.load.recovery` or `vendor_kernel_boot.modules.load.recovery`",
+        ),
+        "vendor_boot_modules_load_charger": attr.output(
+            doc = "`vendor_boot.modules.load.charger` or `vendor_kernel_boot.modules.load.charger`",
+        ),
         "modules_list": attr.label(allow_single_file = True),
+        "modules_recovery_list": attr.label(allow_single_file = True),
+        "modules_charger_list": attr.label(allow_single_file = True),
         "modules_blocklist": attr.label(allow_single_file = True),
         "modules_options": attr.label(allow_single_file = True),
         "ramdisk_compression": attr.string(
diff --git a/kleaf/impl/image/kernel_images.bzl b/kleaf/impl/image/kernel_images.bzl
index 07c68293..1d0b4ce3 100644
--- a/kleaf/impl/image/kernel_images.bzl
+++ b/kleaf/impl/image/kernel_images.bzl
@@ -42,6 +42,8 @@ def kernel_images(
         deps = None,
         boot_image_outs = None,
         modules_list = None,
+        modules_recovery_list = None,
+        modules_charger_list = None,
         modules_blocklist = None,
         modules_options = None,
         vendor_ramdisk_binaries = None,
@@ -86,6 +88,8 @@ def kernel_images(
           - For `initramfs`:
             - The file specified by `MODULES_LIST`
             - The file specified by `MODULES_BLOCKLIST`, if `MODULES_BLOCKLIST` is set
+            - The file containing the list of modules needed for booting into recovery.
+            - The file containing the list of modules needed for booting into charger mode.
           - For `vendor_dlkm` image:
             - The file specified by `VENDOR_DLKM_MODULES_LIST`
             - The file specified by `VENDOR_DLKM_MODULES_BLOCKLIST`, if set
@@ -180,6 +184,10 @@ def kernel_images(
         modules_list: A file containing list of modules to use for `vendor_boot.modules.load`.
 
           This corresponds to `MODULES_LIST` in `build.config` for `build.sh`.
+        modules_recovery_list: A file containing a list of modules to load when booting into
+          recovery.
+        modules_charger_list: A file containing a list of modules to load when booting into
+          charger mode.
         modules_blocklist: A file containing a list of modules which are
           blocked from being loaded.
 
@@ -339,12 +347,23 @@ def kernel_images(
     if build_vendor_kernel_boot and "vendor_kernel_boot.img" not in boot_image_outs:
         boot_image_outs.append("vendor_kernel_boot.img")
 
+    vendor_boot_name = None
+    if build_vendor_boot:
+        vendor_boot_name = "vendor_boot"
+    elif build_vendor_kernel_boot:
+        vendor_boot_name = "vendor_kernel_boot"
+
     vendor_boot_modules_load = None
+    vendor_boot_modules_load_recovery = None
+    vendor_boot_modules_load_charger = None
     if build_initramfs:
-        if build_vendor_boot:
-            vendor_boot_modules_load = "{}_initramfs/vendor_boot.modules.load".format(name)
-        elif build_vendor_kernel_boot:
-            vendor_boot_modules_load = "{}_initramfs/vendor_kernel_boot.modules.load".format(name)
+        vendor_boot_modules_load = "{}_initramfs/{}.modules.load".format(name, vendor_boot_name)
+
+        if modules_recovery_list:
+            vendor_boot_modules_load_recovery = "{}_initramfs/{}.modules.load.recovery".format(name, vendor_boot_name)
+
+        if modules_charger_list:
+            vendor_boot_modules_load_charger = "{}_initramfs/{}.modules.load.charger".format(name, vendor_boot_name)
 
         if ramdisk_compression_args and ramdisk_compression != "lz4":
             fail(
@@ -358,7 +377,11 @@ def kernel_images(
             kernel_modules_install = kernel_modules_install,
             deps = deps,
             vendor_boot_modules_load = vendor_boot_modules_load,
+            vendor_boot_modules_load_recovery = vendor_boot_modules_load_recovery,
+            vendor_boot_modules_load_charger = vendor_boot_modules_load_charger,
             modules_list = modules_list,
+            modules_recovery_list = modules_recovery_list,
+            modules_charger_list = modules_charger_list,
             modules_blocklist = modules_blocklist,
             modules_options = modules_options,
             ramdisk_compression = ramdisk_compression,
@@ -410,12 +433,6 @@ def kernel_images(
         all_rules.append(":{}_vendor_dlkm_image".format(name))
 
     if build_any_boot_image:
-        if build_vendor_kernel_boot:
-            vendor_boot_name = "vendor_kernel_boot"
-        elif build_vendor_boot:
-            vendor_boot_name = "vendor_boot"
-        else:
-            vendor_boot_name = None
         boot_images(
             name = "{}_boot_images".format(name),
             kernel_build = kernel_build,
diff --git a/kleaf/kernel.bzl b/kleaf/kernel.bzl
index 42ee2ffb..ffae3adf 100644
--- a/kleaf/kernel.bzl
+++ b/kleaf/kernel.bzl
@@ -18,8 +18,11 @@
 # rules and macros. The implementations stays in sub-extensions,
 # which is not expected to be loaded directly by users.
 
-load("//build/kernel/kleaf/impl:gki_artifacts.bzl", _gki_artifacts = "gki_artifacts", _gki_artifacts_prebuilts = "gki_artifacts_prebuilts")
-load("//build/kernel/kleaf/artifact_tests:kernel_test.bzl", _kernel_module_test = "kernel_module_test")
+load(
+    "//build/kernel/kleaf/artifact_tests:kernel_test.bzl",
+    _initramfs_modules_lists_test = "initramfs_modules_lists_test",
+    _kernel_module_test = "kernel_module_test",
+)
 load("//build/kernel/kleaf/impl:abi/extracted_symbols.bzl", _extract_symbols = "extracted_symbols")
 load("//build/kernel/kleaf/impl:abi/kernel_abi.bzl", _kernel_abi = "kernel_abi")
 load("//build/kernel/kleaf/impl:abi/kernel_abi_dist.bzl", _kernel_abi_dist = "kernel_abi_dist")
@@ -28,6 +31,7 @@ load("//build/kernel/kleaf/impl:checkpatch.bzl", _checkpatch = "checkpatch")
 load("//build/kernel/kleaf/impl:ddk/ddk_headers.bzl", _ddk_headers = "ddk_headers")
 load("//build/kernel/kleaf/impl:ddk/ddk_module.bzl", _ddk_module = "ddk_module")
 load("//build/kernel/kleaf/impl:ddk/ddk_submodule.bzl", _ddk_submodule = "ddk_submodule")
+load("//build/kernel/kleaf/impl:gki_artifacts.bzl", _gki_artifacts = "gki_artifacts", _gki_artifacts_prebuilts = "gki_artifacts_prebuilts")
 load("//build/kernel/kleaf/impl:image/kernel_images.bzl", _kernel_images = "kernel_images")
 load("//build/kernel/kleaf/impl:kernel_build.bzl", _kernel_build_macro = "kernel_build")
 load("//build/kernel/kleaf/impl:kernel_build_config.bzl", _kernel_build_config = "kernel_build_config")
@@ -51,6 +55,7 @@ ddk_submodule = _ddk_submodule
 extract_symbols = _extract_symbols
 gki_artifacts = _gki_artifacts
 gki_artifacts_prebuilts = _gki_artifacts_prebuilts
+initramfs_modules_lists_test = _initramfs_modules_lists_test
 kernel_abi = _kernel_abi
 kernel_abi_dist = _kernel_abi_dist
 kernel_build = _kernel_build_macro