Fix build error due to chkstk/alloca on x86

The `avb_cert_validate_vbmeta_public_key` API uses more than 8KB of
stack memory and triggers the compiler to insert calls to function
__chkstk() and __alloca() on x86/x86_64 platforms. Although rust
libcompiler_builtins does provide the implementations, there are
two issues:

1. LLVM mangles the symbol and add an additional underscore to the
   symbol name, causing link to fail to find the correct one.
2. The implementation is not enabled for x86 uefi build. x86_64 had the
   same issue but was fixed. But it wasn't ported to x86.

As a workaround, for #1 we define mangled symbols that simply jump to
the correct one. For #2 we apply the same fix from x86_64 for x86 at
build time to and rebuild std from source, until upstream is fixed.

Bug: 337114254

Change-Id: Id7bce71a7d5ca08fc79180733e40fc1ec083da57

9 files changed, 208 insertions, 5 deletions

diff --git a/gbl/efi/BUILD b/gbl/efi/BUILD
index 73f9683..db28634 100644
--- a/gbl/efi/BUILD
+++ b/gbl/efi/BUILD
@@ -60,6 +60,8 @@ rust_binary(
         "@zerocopy",
     ] + select({
         "@gbl//toolchain:gbl_rust_elf_riscv64": ["@gbl//efi/arch/riscv64:efi_arch_deps_riscv64"],
+        "@gbl//toolchain:gbl_rust_uefi_x86_64": ["@gbl//efi/arch/x86_64:efi_arch_deps_x86_64"],
+        "@gbl//toolchain:gbl_rust_uefi_x86_32": ["@gbl//efi/arch/x86:efi_arch_deps_x86"],
         "//conditions:default": [],
     }),
 )
diff --git a/gbl/efi/arch/x86/BUILD b/gbl/efi/arch/x86/BUILD
new file mode 100644
index 0000000..0b69afc
--- /dev/null
+++ b/gbl/efi/arch/x86/BUILD
@@ -0,0 +1,30 @@
+# 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.
+
+load("@gbl//toolchain:gbl_toolchain.bzl", "link_static_cc_library")
+
+package(
+    default_visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "lib_efi_arch_deps_x86",
+    srcs = ["deps.S"],
+    visibility = ["//visibility:public"],
+)
+
+link_static_cc_library(
+    name = "efi_arch_deps_x86",
+    cc_library = ":lib_efi_arch_deps_x86",
+)
diff --git a/gbl/efi/arch/x86/deps.S b/gbl/efi/arch/x86/deps.S
new file mode 100644
index 0000000..9107919
--- /dev/null
+++ b/gbl/efi/arch/x86/deps.S
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+/*
+ * LLVM mangles the "___chkstk()" and "__alloca()" function from rust
+ * libcompiler_builtins with an additional underscore, making them "___chkstk"
+ * and "___alloca and causing the linker to fail finding them. We workaround by
+ * defining the magled symbols that simply jump to the correct target.
+ */
+
+.global ___chkstk
+.global __alloca
+
+___chkstk:
+    jmp ____chkstk
+
+__alloca:
+    jmp ___alloca
diff --git a/gbl/efi/arch/x86_64/BUILD b/gbl/efi/arch/x86_64/BUILD
new file mode 100644
index 0000000..5dda42e
--- /dev/null
+++ b/gbl/efi/arch/x86_64/BUILD
@@ -0,0 +1,30 @@
+# 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.
+
+load("@gbl//toolchain:gbl_toolchain.bzl", "link_static_cc_library")
+
+package(
+    default_visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "lib_efi_arch_deps_x86_64",
+    srcs = ["deps.S"],
+    visibility = ["//visibility:public"],
+)
+
+link_static_cc_library(
+    name = "efi_arch_deps_x86_64",
+    cc_library = ":lib_efi_arch_deps_x86_64",
+)
diff --git a/gbl/efi/arch/x86_64/deps.S b/gbl/efi/arch/x86_64/deps.S
new file mode 100644
index 0000000..dbe78d9
--- /dev/null
+++ b/gbl/efi/arch/x86_64/deps.S
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+/*
+ * LLVM mangles the "void __chkstk()" function from rust libcompiler_builtins
+ * with an additional underscore, making it "___chkstk" and causing the linker
+ * to fail finding "__chkstk". We workaround by defining the magled symbol that
+ * simply jumps to the correct target.
+ */
+
+.global __chkstk
+
+__chkstk:
+    jmp ___chkstk
diff --git a/gbl/patches/BUILD b/gbl/patches/BUILD
new file mode 100644
index 0000000..7509dde
--- /dev/null
+++ b/gbl/patches/BUILD
@@ -0,0 +1,19 @@
+# Copyright (C) 2024 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.
+
+package(
+    default_visibility = ["//visibility:public"],
+)
+
+exports_files(glob(["**/*"]))
diff --git a/gbl/patches/rust-libcompiler-builtins-enable-chkstk-on-uefi.patch b/gbl/patches/rust-libcompiler-builtins-enable-chkstk-on-uefi.patch
new file mode 100644
index 0000000..9153f6e
--- /dev/null
+++ b/gbl/patches/rust-libcompiler-builtins-enable-chkstk-on-uefi.patch
@@ -0,0 +1,32 @@
+--- /tmp/x86.rs	2024-04-25 23:47:26.697256628 +0000
++++ x86.rs	2024-04-25 23:47:33.253251743 +0000
+@@ -11,8 +11,7 @@
+ intrinsics! {
+     #[naked]
+     #[cfg(all(
+-        windows,
+-        target_env = "gnu",
++        any(all(windows, target_env = "gnu"), target_os = "uefi"),
+         not(feature = "no-asm")
+     ))]
+     pub unsafe extern "C" fn ___chkstk_ms() {
+@@ -41,8 +40,7 @@
+     // FIXME: __alloca should be an alias to __chkstk
+     #[naked]
+     #[cfg(all(
+-        windows,
+-        target_env = "gnu",
++        any(all(windows, target_env = "gnu"), target_os = "uefi"),
+         not(feature = "no-asm")
+     ))]
+     pub unsafe extern "C" fn __alloca() {
+@@ -54,8 +52,7 @@
+
+     #[naked]
+     #[cfg(all(
+-        windows,
+-        target_env = "gnu",
++        any(all(windows, target_env = "gnu"), target_os = "uefi"),
+         not(feature = "no-asm")
+     ))]
+     pub unsafe extern "C" fn ___chkstk() {
diff --git a/gbl/toolchain/BUILD b/gbl/toolchain/BUILD
index b4c99d3..882b8ba 100644
--- a/gbl/toolchain/BUILD
+++ b/gbl/toolchain/BUILD
@@ -194,6 +194,8 @@ gbl_clang_cc_toolchain(
         # crash as a result.
         "-march=i686",
     ],
+    # Safe Exception Handlers is not applicable to EFI systems.
+    ld_flags = ["/SAFESEH:no"],
     target_cpu = "x86_32",
     target_system_triple = "i686-unknown-windows-gnu",
 )
@@ -355,7 +357,14 @@ rust_toolchain(
     ],
     global_allocator_library = None,
     rust_doc = "@rust_prebuilts//:bin/rustdoc",
-    rust_std = "@rust_prebuilts//:i686-unknown-uefi_prebuilt_stdlibs",
+    # Need to use our own sysroot because we have a custom patch that needs to be applied for std
+    # in order to enable compiler builtin for chkstk(), alloca(). Once the fix is upstreamed, we
+    # can use the prebuilt.
+    rust_std = select({
+        ":rust_no_sysroot_true": ":rust_stdlib_empty",
+        "//conditions:default": ":rust_std_source_build",
+    }),
+    #rust_std = "@rust_prebuilts//:i686-unknown-uefi_prebuilt_stdlibs",
     rustc = "@rust_prebuilts//:bin/rustc",
     staticlib_ext = ".a",
     stdlib_linkflags = [],
diff --git a/gbl/toolchain/BUILD.android_rust_prebuilts.bazel b/gbl/toolchain/BUILD.android_rust_prebuilts.bazel
index 87b40d1..ee65923 100644
--- a/gbl/toolchain/BUILD.android_rust_prebuilts.bazel
+++ b/gbl/toolchain/BUILD.android_rust_prebuilts.bazel
@@ -63,12 +63,34 @@ rust_library(
     ],
 )
 
+COMPILER_BUILTIN_X86_SRC = "src/stdlibs/vendor/compiler_builtins/src/x86.rs"
+
+COMPILER_BUILTIN_X86_PATCH = "@gbl//patches:rust-libcompiler-builtins-enable-chkstk-on-uefi.patch"
+
+# Apply patches to x86 to enable the compiler built-in for chkstk(), alloca() function.
+# TODO(b/337114254): Remove this patch once upstream is fixed.
+genrule(
+    name = "x86_uefi_chkstk_patch",
+    srcs = [
+        COMPILER_BUILTIN_X86_SRC,
+        COMPILER_BUILTIN_X86_PATCH,
+    ],
+    # Bazel does not let us override the same x86.rs. But module name needs to be the same. Thus we
+    # output as "x86/mod.rs" instead.
+    outs = ["src/stdlibs/vendor/compiler_builtins/src/x86/mod.rs"],
+    cmd = " cp $(location {}) $(OUTS) && patch -u -f $(OUTS) $(location {})"
+        .format(COMPILER_BUILTIN_X86_SRC, COMPILER_BUILTIN_X86_PATCH),
+)
+
 rust_library(
     name = "libcompiler_builtins",
-    srcs = glob([
-        "src/stdlibs/vendor/compiler_builtins/src/**/*.rs",
-        "src/stdlibs/vendor/compiler_builtins/libm/src/**/*.rs",
-    ]),
+    srcs = glob(
+        [
+            "src/stdlibs/vendor/compiler_builtins/src/**/*.rs",
+            "src/stdlibs/vendor/compiler_builtins/libm/src/**/*.rs",
+        ],
+        exclude = [COMPILER_BUILTIN_X86_SRC],
+    ) + [":x86_uefi_chkstk_patch"],
     compile_data = glob(["src/stdlibs/vendor/compiler_builtins/src/**/*.md"]),
     crate_features = [
         "compiler-builtins",
@@ -77,6 +99,7 @@ rust_library(
         "mem",
     ],
     crate_name = "compiler_builtins",
+    data = [":x86_uefi_chkstk_patch"],
     edition = "2015",
     rustc_flags = ["--cap-lints=allow"],
     deps = ["libcore"],