Add rule for ASB consumption via APKs.

PiperOrigin-RevId: 558059461
Change-Id: I872764f21c662ae101a766a9890bcd97ba09af52
diff --git a/rules/acls.bzl b/rules/acls.bzl
index b6e174c..d252e99 100644
--- a/rules/acls.bzl
+++ b/rules/acls.bzl
@@ -42,6 +42,7 @@
 load("//rules/acls:android_instrumentation_binary_starlark_resources.bzl", "ANDROID_INSTRUMENTATION_BINARY_STARLARK_RESOURCES_FALLBACK", "ANDROID_INSTRUMENTATION_BINARY_STARLARK_RESOURCES_ROLLOUT")
 load("//rules/acls:android_binary_starlark_javac.bzl", "ANDROID_BINARY_STARLARK_JAVAC_FALLBACK", "ANDROID_BINARY_STARLARK_JAVAC_ROLLOUT")
 load("//rules/acls:android_binary_starlark_split_transition.bzl", "ANDROID_BINARY_STARLARK_SPLIT_TRANSITION_FALLBACK", "ANDROID_BINARY_STARLARK_SPLIT_TRANSITION_ROLLOUT")
+load("//rules/acls:android_binary_with_sandboxed_sdks_allowlist.bzl", "ANDROID_BINARY_WITH_SANDBOXED_SDKS_ALLOWLIST")
 load("//rules/acls:android_feature_splits_dogfood.bzl", "ANDROID_FEATURE_SPLITS_DOGFOOD")
 load("//rules/acls:android_library_resources_without_srcs.bzl", "ANDROID_LIBRARY_RESOURCES_WITHOUT_SRCS", "ANDROID_LIBRARY_RESOURCES_WITHOUT_SRCS_GENERATOR_FUNCTIONS")
 load("//rules/acls:android_library_starlark_resource_outputs.bzl", "ANDROID_LIBRARY_STARLARK_RESOURCE_OUTPUTS_FALLBACK", "ANDROID_LIBRARY_STARLARK_RESOURCE_OUTPUTS_ROLLOUT")
@@ -112,6 +113,9 @@
 def _in_android_binary_starlark_split_transition(fqn):
     return not matches(fqn, ANDROID_BINARY_STARLARK_SPLIT_TRANSITION_FALLBACK_DICT) and matches(fqn, ANDROID_BINARY_STARLARK_SPLIT_TRANSITION_ROLLOUT_DICT)
 
+def _in_android_binary_with_sandboxed_sdks_allowlist(fqn):
+    return matches(fqn, ANDROID_BINARY_WITH_SANDBOXED_SDKS_ALLOWLIST_DICT)
+
 def _in_android_feature_splits_dogfood(fqn):
     return matches(fqn, ANDROID_FEATURE_SPLITS_DOGFOOD_DICT)
 
@@ -257,6 +261,7 @@
 ANDROID_BINARY_STARLARK_JAVAC_FALLBACK_DICT = make_dict(ANDROID_BINARY_STARLARK_JAVAC_FALLBACK)
 ANDROID_BINARY_STARLARK_SPLIT_TRANSITION_ROLLOUT_DICT = make_dict(ANDROID_BINARY_STARLARK_SPLIT_TRANSITION_ROLLOUT)
 ANDROID_BINARY_STARLARK_SPLIT_TRANSITION_FALLBACK_DICT = make_dict(ANDROID_BINARY_STARLARK_SPLIT_TRANSITION_FALLBACK)
+ANDROID_BINARY_WITH_SANDBOXED_SDKS_ALLOWLIST_DICT = make_dict(ANDROID_BINARY_WITH_SANDBOXED_SDKS_ALLOWLIST)
 ANDROID_FEATURE_SPLITS_DOGFOOD_DICT = make_dict(ANDROID_FEATURE_SPLITS_DOGFOOD)
 ANDROID_LIBRARY_RESOURCES_WITHOUT_SRCS_DICT = make_dict(ANDROID_LIBRARY_RESOURCES_WITHOUT_SRCS)
 ANDROID_LIBRARY_RESOURCES_WITHOUT_SRCS_GENERATOR_FUNCTIONS_DICT = make_dict(ANDROID_LIBRARY_RESOURCES_WITHOUT_SRCS_GENERATOR_FUNCTIONS)
@@ -372,6 +377,7 @@
     in_android_instrumentation_binary_starlark_resources = _in_android_instrumentation_binary_starlark_resources,
     in_android_binary_starlark_javac = _in_android_binary_starlark_javac,
     in_android_binary_starlark_split_transition = _in_android_binary_starlark_split_transition,
+    in_android_binary_with_sandboxed_sdks_allowlist = _in_android_binary_with_sandboxed_sdks_allowlist,
     in_android_feature_splits_dogfood = _in_android_feature_splits_dogfood,
     in_android_library_starlark_resource_outputs_rollout = _in_android_library_starlark_resource_outputs_rollout,
     in_android_library_resources_without_srcs = _in_android_library_resources_without_srcs,
diff --git a/rules/acls/android_binary_with_sandboxed_sdks_allowlist.bzl b/rules/acls/android_binary_with_sandboxed_sdks_allowlist.bzl
new file mode 100644
index 0000000..b926437
--- /dev/null
+++ b/rules/acls/android_binary_with_sandboxed_sdks_allowlist.bzl
@@ -0,0 +1,20 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# 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.
+
+"""Allow list of android_binary_with_sandboxed_sdks rule."""
+
+# keep sorted
+ANDROID_BINARY_WITH_SANDBOXED_SDKS_ALLOWLIST = [
+    "//:__subpackages__",
+]
diff --git a/rules/android_sandboxed_sdk/android_binary_with_sandboxed_sdks_macro.bzl b/rules/android_sandboxed_sdk/android_binary_with_sandboxed_sdks_macro.bzl
new file mode 100644
index 0000000..3214f34
--- /dev/null
+++ b/rules/android_sandboxed_sdk/android_binary_with_sandboxed_sdks_macro.bzl
@@ -0,0 +1,220 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# 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.
+
+"""Bazel rule for defining an Android binary that depends on sandboxed SDKs."""
+
+load(":providers.bzl", "AndroidSandboxedSdkBundleInfo")
+load("//rules:acls.bzl", "acls")
+load("//rules:bundletool.bzl", _bundletool = "bundletool")
+load("//rules:common.bzl", _common = "common")
+load(
+    "//rules:utils.bzl",
+    _get_android_toolchain = "get_android_toolchain",
+)
+load("//rules:java.bzl", _java = "java")
+
+def _gen_sdk_dependencies_manifest_impl(ctx):
+    manifest = ctx.actions.declare_file(ctx.label.name + "_sdk_dep_manifest.xml")
+
+    module_configs = [
+        bundle[AndroidSandboxedSdkBundleInfo].sdk_info.sdk_module_config
+        for bundle in ctx.attr.sdk_bundles
+    ]
+
+    args = ctx.actions.args()
+    args.add("generate-sdk-dependencies-manifest")
+    args.add("--manifest-package", ctx.attr.package)
+    args.add("--sdk-module-configs", ",".join([config.path for config in module_configs]))
+    args.add("--debug-keystore", ctx.file.debug_key.path)
+    args.add("--debug-keystore-pass", "android")
+    args.add("--debug-keystore-alias", "androiddebugkey")
+    args.add("--output-manifest", manifest)
+    _java.run(
+        ctx = ctx,
+        host_javabase = _common.get_host_javabase(ctx),
+        executable = _get_android_toolchain(ctx).sandboxed_sdk_toolbox.files_to_run,
+        arguments = [args],
+        inputs = module_configs + [ctx.file.debug_key],
+        outputs = [manifest],
+        mnemonic = "GenSdkDepManifest",
+        progress_message = "Generate SDK dependencies manifest %s" % manifest.short_path,
+    )
+
+    return [
+        DefaultInfo(
+            files = depset([manifest]),
+        ),
+    ]
+
+_gen_sdk_dependencies_manifest = rule(
+    attrs = dict(
+        package = attr.string(),
+        sdk_bundles = attr.label_list(
+            providers = [
+                [AndroidSandboxedSdkBundleInfo],
+            ],
+        ),
+        debug_key = attr.label(
+            allow_single_file = True,
+            default = Label("//tools/android:debug_keystore"),
+        ),
+        _host_javabase = attr.label(
+            cfg = "exec",
+            default = Label("//tools/jdk:current_java_runtime"),
+        ),
+    ),
+    executable = False,
+    implementation = _gen_sdk_dependencies_manifest_impl,
+    toolchains = [
+        "//toolchains/android:toolchain_type",
+    ],
+)
+
+def _android_binary_with_sandboxed_sdks_impl(ctx):
+    sdk_apks = []
+    for idx, sdk_bundle_target in enumerate(ctx.attr.sdk_bundles):
+        apk_out = ctx.actions.declare_file("%s/sdk_dep_apks/%s.apk" % (
+            ctx.label.name,
+            idx,
+        ))
+        _bundletool.build_sdk_apks(
+            ctx,
+            out = apk_out,
+            aapt2 = _get_android_toolchain(ctx).aapt2.files_to_run,
+            sdk_bundle = sdk_bundle_target[AndroidSandboxedSdkBundleInfo].asb,
+            debug_key = ctx.file.debug_key,
+            bundletool = _get_android_toolchain(ctx).bundletool.files_to_run,
+            host_javabase = _common.get_host_javabase(ctx),
+        )
+        sdk_apks.append(apk_out)
+
+    app_apk = ctx.attr.internal_android_binary[ApkInfo].signed_apk
+    adb = _get_android_toolchain(ctx).adb.files_to_run.executable
+    substitutions = {
+        "%adb%": adb.short_path,
+        "%app_apk%": app_apk.short_path,
+        "%sdk_apks%": ",".join([apk.short_path for apk in sdk_apks]),
+    }
+
+    install_script = ctx.actions.declare_file("%s_install_script.sh" % ctx.label.name)
+    ctx.actions.expand_template(
+        template = ctx.file._install_script_template,
+        output = install_script,
+        substitutions = substitutions,
+        is_executable = True,
+    )
+
+    return [
+        DefaultInfo(
+            executable = install_script,
+            files = depset([app_apk] + sdk_apks),
+            runfiles = ctx.runfiles([
+                adb,
+                app_apk,
+            ] + sdk_apks),
+        ),
+    ]
+
+_android_binary_with_sandboxed_sdks = rule(
+    attrs = dict(
+        internal_android_binary = attr.label(
+            providers = [
+                [ApkInfo],
+            ],
+        ),
+        debug_key = attr.label(
+            allow_single_file = True,
+            default = Label("//tools/android:debug_keystore"),
+        ),
+        sdk_bundles = attr.label_list(
+            providers = [
+                [AndroidSandboxedSdkBundleInfo],
+            ],
+        ),
+        _install_script_template = attr.label(
+            allow_single_file = True,
+            default = ":install_script.sh_template",
+        ),
+        _host_javabase = attr.label(
+            cfg = "exec",
+            default = Label("//tools/jdk:current_java_runtime"),
+        ),
+    ),
+    executable = True,
+    implementation = _android_binary_with_sandboxed_sdks_impl,
+    toolchains = [
+        "//toolchains/android:toolchain_type",
+    ],
+)
+
+def android_binary_with_sandboxed_sdks_macro(
+        _android_binary,
+        _android_library,
+        **attrs):
+    """android_binary_with_sandboxed_sdks.
+
+    Args:
+      _android_binary: The android_binary rule to use.
+      _android_library: The android_library rule to use.
+      **attrs: android_binary attributes.
+    """
+
+    name = attrs.pop("name", None)
+    fully_qualified_name = "//%s:%s" % (native.package_name(), name)
+    if (not acls.in_android_binary_with_sandboxed_sdks_allowlist(fully_qualified_name)):
+        fail("%s is not allowed to use the android_binary_with_sandboxed_sdks macro." %
+             fully_qualified_name)
+
+    sdk_bundles = attrs.pop("sdk_bundles", None)
+    debug_keystore = getattr(attrs, "debug_keystore", None)
+
+    bin_package = _java.resolve_package_from_label(
+        Label(fully_qualified_name),
+        getattr(attrs, "custom_package", None),
+    )
+
+    # Generate a manifest that lists all the SDK dependencies with <uses-sdk-library> tags.
+    sdk_dependencies_manifest_name = "%s_sdk_dependencies_manifest" % name
+    _gen_sdk_dependencies_manifest(
+        name = sdk_dependencies_manifest_name,
+        package = "%s.internalsdkdependencies" % bin_package,
+        sdk_bundles = sdk_bundles,
+    )
+
+    # Use the manifest in a normal android_library. This will later be added as a dependency to the
+    # binary, so the manifest is merged with the app's.
+    sdk_dependencies_lib_name = "%s_sdk_dependencies_lib" % name
+    _android_library(
+        name = sdk_dependencies_lib_name,
+        exports_manifest = 1,
+        manifest = ":%s" % sdk_dependencies_manifest_name,
+    )
+    deps = attrs.pop("deps", [])
+    deps.append(":%s" % sdk_dependencies_lib_name)
+
+    # Generate the android_binary as normal, passing the extra flags.
+    bin_label = Label("%s_app_bin" % fully_qualified_name)
+    _android_binary(
+        name = bin_label.name,
+        deps = deps,
+        **attrs
+    )
+
+    # This final rule will call Bundletool to generate the SDK APKs and provide the install script.
+    _android_binary_with_sandboxed_sdks(
+        name = name,
+        sdk_bundles = sdk_bundles,
+        debug_key = debug_keystore,
+        internal_android_binary = bin_label,
+    )
diff --git a/rules/bundletool.bzl b/rules/bundletool.bzl
index 1869535..7988bac 100644
--- a/rules/bundletool.bzl
+++ b/rules/bundletool.bzl
@@ -75,6 +75,60 @@
     ))
     ctx.actions.write(out, json_content)
 
+def _build_sdk_apks(
+        ctx,
+        out = None,
+        aapt2 = None,
+        sdk_bundle = None,
+        debug_key = None,
+        bundletool = None,
+        host_javabase = None):
+    apks_out = ctx.actions.declare_directory(ctx.label.name + "_sdk_apks")
+    args = ctx.actions.args()
+    args.add("build-sdk-apks")
+    args.add("--aapt2", aapt2.executable.path)
+    args.add("--sdk-bundle", sdk_bundle)
+    args.add("--ks", debug_key)
+    args.add("--ks-pass=pass:android")
+    args.add("--ks-key-alias=androiddebugkey")
+    args.add("--key-pass=pass:android")
+    args.add("--output-format=DIRECTORY")
+    args.add("--output", apks_out.path)
+    _java.run(
+        ctx = ctx,
+        host_javabase = host_javabase,
+        executable = bundletool,
+        arguments = [args],
+        inputs = [
+            sdk_bundle,
+            debug_key,
+        ],
+        tools = [aapt2],
+        outputs = [apks_out],
+        mnemonic = "BuildSdkApksDir",
+        progress_message = "Building SDK APKs directory %s" % apks_out.short_path,
+    )
+
+    # Now move standalone APK out of bundletool output dir.
+    ctx.actions.run_shell(
+        command = """
+set -e
+APKS_OUT_DIR=%s
+DEBUG_APK_PATH=%s
+
+mv "${APKS_OUT_DIR}/standalones/standalone.apk" "${DEBUG_APK_PATH}"
+""" % (
+            apks_out.path,
+            out.path,
+        ),
+        tools = [],
+        arguments = [],
+        inputs = [apks_out],
+        outputs = [out],
+        mnemonic = "ExtractDebugSdkApk",
+        progress_message = "Extract debug SDK APK to %s" % out.short_path,
+    )
+
 def _build_sdk_bundle(
         ctx,
         out = None,
@@ -324,6 +378,7 @@
 bundletool = struct(
     build = _build,
     build_device_json = _build_device_json,
+    build_sdk_apks = _build_sdk_apks,
     build_sdk_bundle = _build_sdk_bundle,
     build_sdk_module = _build_sdk_module,
     bundle_to_apks = _bundle_to_apks,