kleaf: extract module_env.tar.gz to KLEAF_REPO_DIR

This extracts module_env.tar.gz, a sub-tree that contains
all necessary scripts, to KLEAF_REPO_DIR.

This resolves the following error:

/bin/bash: line 35: external/kleaf~override/common/scripts/kconfig/merge_config.sh: No such file or directory

Bug: 291918087
Bug: 308404299
Test: manual
Change-Id: I06fc1c4a17e5fb58ab429c71f722f601c84dd4b7

5 files changed, 25 insertions, 2 deletions

diff --git a/kleaf/docs/api_reference/kernel.md b/kleaf/docs/api_reference/kernel.md
index d2a9640..1dc9ffc 100644
--- a/kleaf/docs/api_reference/kernel.md
+++ b/kleaf/docs/api_reference/kernel.md
@@ -261,8 +261,8 @@ Define an executable that creates `compile_commands.json` from a `kernel_build`.
 kernel_filegroup(<a href="#kernel_filegroup-name">name</a>, <a href="#kernel_filegroup-deps">deps</a>, <a href="#kernel_filegroup-srcs">srcs</a>, <a href="#kernel_filegroup-collect_unstripped_modules">collect_unstripped_modules</a>, <a href="#kernel_filegroup-config_out_dir">config_out_dir</a>, <a href="#kernel_filegroup-config_out_dir_files">config_out_dir_files</a>,
                  <a href="#kernel_filegroup-ddk_module_defconfig_fragments">ddk_module_defconfig_fragments</a>, <a href="#kernel_filegroup-debug">debug</a>, <a href="#kernel_filegroup-env_setup_script">env_setup_script</a>, <a href="#kernel_filegroup-exec_platform">exec_platform</a>,
                  <a href="#kernel_filegroup-gki_artifacts">gki_artifacts</a>, <a href="#kernel_filegroup-images">images</a>, <a href="#kernel_filegroup-internal_outs">internal_outs</a>, <a href="#kernel_filegroup-kasan">kasan</a>, <a href="#kernel_filegroup-kasan_generic">kasan_generic</a>, <a href="#kernel_filegroup-kasan_sw_tags">kasan_sw_tags</a>, <a href="#kernel_filegroup-kcsan">kcsan</a>,
-                 <a href="#kernel_filegroup-kernel_release">kernel_release</a>, <a href="#kernel_filegroup-kernel_uapi_headers">kernel_uapi_headers</a>, <a href="#kernel_filegroup-lto">lto</a>, <a href="#kernel_filegroup-module_outs_file">module_outs_file</a>, <a href="#kernel_filegroup-modules_prepare_archive">modules_prepare_archive</a>,
-                 <a href="#kernel_filegroup-protected_modules_list">protected_modules_list</a>, <a href="#kernel_filegroup-target_platform">target_platform</a>, <a href="#kernel_filegroup-trim_nonlisted_kmi">trim_nonlisted_kmi</a>)
+                 <a href="#kernel_filegroup-kernel_release">kernel_release</a>, <a href="#kernel_filegroup-kernel_uapi_headers">kernel_uapi_headers</a>, <a href="#kernel_filegroup-lto">lto</a>, <a href="#kernel_filegroup-module_env_archive">module_env_archive</a>, <a href="#kernel_filegroup-module_outs_file">module_outs_file</a>,
+                 <a href="#kernel_filegroup-modules_prepare_archive">modules_prepare_archive</a>, <a href="#kernel_filegroup-protected_modules_list">protected_modules_list</a>, <a href="#kernel_filegroup-target_platform">target_platform</a>, <a href="#kernel_filegroup-trim_nonlisted_kmi">trim_nonlisted_kmi</a>)
 </pre>
 
 **EXPERIMENTAL.** The API of `kernel_filegroup` rapidly changes and
@@ -303,6 +303,7 @@ It can be used in the `base_kernel` attribute of a [`kernel_build`](#kernel_buil
 | <a id="kernel_filegroup-kernel_release"></a>kernel_release |  A file providing the kernel release string. This is preferred over `gki_artifacts`.   | <a href="https://bazel.build/concepts/labels">Label</a> | optional |  `None`  |
 | <a id="kernel_filegroup-kernel_uapi_headers"></a>kernel_uapi_headers |  The label pointing to `kernel-uapi-headers.tar.gz`.<br><br>This attribute should be set to the `kernel-uapi-headers.tar.gz` artifact built by the [`kernel_build`](#kernel_build) macro if the `kernel_filegroup` rule were a `kernel_build`.<br><br>Setting this attribute allows [`merged_kernel_uapi_headers`](#merged_kernel_uapi_headers) to work properly when this `kernel_filegroup` is set to the `base_kernel`.<br><br>For example: <pre><code>kernel_filegroup(&#10;    name = "kernel_aarch64_prebuilts",&#10;    srcs = [&#10;        "vmlinux",&#10;        # ...&#10;    ],&#10;    kernel_uapi_headers = "kernel-uapi-headers.tar.gz",&#10;)&#10;&#10;kernel_build(&#10;    name = "tuna",&#10;    base_kernel = ":kernel_aarch64_prebuilts",&#10;    # ...&#10;)&#10;&#10;merged_kernel_uapi_headers(&#10;    name = "tuna_merged_kernel_uapi_headers",&#10;    kernel_build = "tuna",&#10;    # ...&#10;)</code></pre>   | <a href="https://bazel.build/concepts/labels">Label</a> | optional |  `None`  |
 | <a id="kernel_filegroup-lto"></a>lto |  -   | String | optional |  `"default"`  |
+| <a id="kernel_filegroup-module_env_archive"></a>module_env_archive |  Archive from `kernel_build.pack_module_env` that contains necessary files to build external modules.   | <a href="https://bazel.build/concepts/labels">Label</a> | optional |  `None`  |
 | <a id="kernel_filegroup-module_outs_file"></a>module_outs_file |  A file containing `module_outs` of the original [`kernel_build`](#kernel_build) target.   | <a href="https://bazel.build/concepts/labels">Label</a> | required |  |
 | <a id="kernel_filegroup-modules_prepare_archive"></a>modules_prepare_archive |  Archive from `modules_prepare`   | <a href="https://bazel.build/concepts/labels">Label</a> | optional |  `None`  |
 | <a id="kernel_filegroup-protected_modules_list"></a>protected_modules_list |  -   | <a href="https://bazel.build/concepts/labels">Label</a> | optional |  `None`  |
diff --git a/kleaf/impl/common_providers.bzl b/kleaf/impl/common_providers.bzl
index 917ef4a..999dd56 100644
--- a/kleaf/impl/common_providers.bzl
+++ b/kleaf/impl/common_providers.bzl
@@ -315,6 +315,7 @@ KernelBuildFilegroupDeclInfo = provider(
         "internal_outs": """[depset](https://bazel.build/extending/depsets) of `kernel_build`'s
             `internal_outs`. May be `None`.""",
         "ruledir": "`ruledir` from `kernel_build` that signifies the root for `internal_outs`.",
+        "module_env_archive": "Archive preparing an environment to build modules. May be `None`.",
     },
 )
 
diff --git a/kleaf/impl/kernel_build.bzl b/kleaf/impl/kernel_build.bzl
index 03ed488..2f382f0 100644
--- a/kleaf/impl/kernel_build.bzl
+++ b/kleaf/impl/kernel_build.bzl
@@ -2014,6 +2014,7 @@ def _create_infos(
         config_out_dir = ctx.file.config,
         internal_outs = depset(all_output_files["internal_outs"].values()),
         ruledir = main_action_ret.ruledir,
+        module_env_archive = module_scripts_archive,
     )
 
     default_info_files = all_output_files["outs"].values() + all_output_files["module_outs"].values()
@@ -2540,6 +2541,7 @@ def _repack_modules_staging_archive(
     )
     return modules_staging_archive
 
+# TODO(b/291918087): Merge into filegroup_decl.tar.gz to flatten the archive.
 def _create_module_scripts_archive(
         ctx,
         module_srcs):
diff --git a/kleaf/impl/kernel_filegroup.bzl b/kleaf/impl/kernel_filegroup.bzl
index 51d8626..1339367 100644
--- a/kleaf/impl/kernel_filegroup.bzl
+++ b/kleaf/impl/kernel_filegroup.bzl
@@ -124,9 +124,19 @@ def _get_ddk_config_env(ctx):
         out_dir = ctx.file.config_out_dir,
     )
 
+    ddk_config_env_setup_command += """
+        # Restore module sources
+        {check_sandbox_cmd}
+        tar xf {module_env_archive} -C ${{KLEAF_REPO_DIR}}
+    """.format(
+        module_env_archive = ctx.file.module_env_archive.path,
+        check_sandbox_cmd = utils.get_check_sandbox_cmd(),
+    )
+
     ddk_config_env_setup_script = ctx.actions.declare_file(
         "{name}/{name}_ddk_config_setup.sh".format(name = ctx.attr.name),
     )
+
     ctx.actions.write(
         output = ddk_config_env_setup_script,
         content = ddk_config_env_setup_command,
@@ -135,6 +145,7 @@ def _get_ddk_config_env(ctx):
         setup_script = ddk_config_env_setup_script,
         inputs = depset([
             ddk_config_env_setup_script,
+            ctx.file.module_env_archive,
             ctx.file.env_setup_script,
             ctx.version_file,
         ], transitive = [target.files for target in ctx.attr.config_out_dir_files]),
@@ -435,6 +446,11 @@ default, which in turn sets `collect_unstripped_modules` to `True` by default.
             allow_single_file = True,
             doc = "Archive from `modules_prepare`",
         ),
+        "module_env_archive": attr.label(
+            allow_single_file = True,
+            doc = """Archive from `kernel_build.pack_module_env` that contains
+                necessary files to build external modules.""",
+        ),
         "internal_outs": attr.label_keyed_string_dict(
             allow_files = True,
             doc = "Keys: from `_kernel_build.internal_outs`. Values: path under `$OUT_DIR`.",
diff --git a/kleaf/impl/kernel_filegroup_declaration.bzl b/kleaf/impl/kernel_filegroup_declaration.bzl
index d6f5580..b8cec83 100644
--- a/kleaf/impl/kernel_filegroup_declaration.bzl
+++ b/kleaf/impl/kernel_filegroup_declaration.bzl
@@ -100,6 +100,7 @@ kernel_filegroup(
     config_out_dir = {config_out_dir_repr},
     env_setup_script = {env_setup_script_repr},
     modules_prepare_archive = {modules_prepare_archive_repr},
+    module_env_archive = {module_env_archive_repr},
     internal_outs = {internal_outs_repr},
     target_platform = {target_platform_repr},
     exec_platform = {exec_platform_repr},
@@ -163,6 +164,7 @@ def _write_filegroup_decl_file(ctx, info, deps_files, kernel_uapi_headers, templ
         depset([info.modules_prepare_archive]),
         **(one | pkg)
     )
+    sub.add_joined("{module_env_archive_repr}", depset([info.module_env_archive]), **(one | pkg))
 
     # {":bazel-out/k8-fastbuild/bin/common/kernel_aarch64/Module.symvers": "Module.symvers", ...}
     sub.add_joined(
@@ -207,6 +209,7 @@ def _create_archive(ctx, info, deps_files, kernel_uapi_headers, filegroup_decl_f
         info.config_out_dir,
         info.env_setup_script,
         info.modules_prepare_archive,
+        info.module_env_archive,
     ]
     if info.src_protected_modules_list:
         direct_inputs.append(info.src_protected_modules_list)