Add helper for reading AArch64 thunks from binaries

PiperOrigin-RevId: 703621699

Author: Propeller Team

propeller/BUILD

@@ -269,6 +269,7 @@ cc_library(
     deps = [
         ":addr2cu",
         ":status_macros",
+        "@abseil-cpp//absl/container:btree",
         "@abseil-cpp//absl/container:flat_hash_map",
         "@abseil-cpp//absl/log",
         "@abseil-cpp//absl/log:check",
@@ -280,6 +281,7 @@ cc_library(
         "@llvm-project//llvm:DebugInfo",
         "@llvm-project//llvm:Object",
         "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:TargetParser",
     ],
 )
 
@@ -1080,6 +1082,7 @@ cc_test(
     name = "binary_content_test",
     srcs = ["binary_content_test.cc"],
     data = [
+        "//propeller/testdata:fake_thunks.bin",
         "//propeller/testdata:llvm_function_samples.binary",
         "//propeller/testdata:propeller_barebone_nopie_buildid",
         "//propeller/testdata:propeller_barebone_pie_nobuildid_bin",

propeller/binary_content.cc

@@ -8,6 +8,7 @@
 #include <utility>
 #include <vector>
 
+#include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
@@ -35,6 +36,7 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/TargetParser/Triple.h"
 #include "propeller/addr2cu.h"
 #include "propeller/status_macros.h"
 
@@ -241,6 +243,30 @@ absl::Status ELFFileUtil<ELFT>::InitializeKernelModule(
 }
 }  // namespace
 
+// Read AArch64 thunks from the symbol table and store them in sorted order.
+absl::btree_map<uint64_t, llvm::object::ELFSymbolRef> ReadAArch64ThunkSymbols(
+    const BinaryContent &binary_content) {
+  absl::btree_map<uint64_t, llvm::object::ELFSymbolRef> thunk_map;
+  for (llvm::object::SymbolRef sr : binary_content.object_file->symbols()) {
+    llvm::object::ELFSymbolRef symbol(sr);
+    uint8_t stt = symbol.getELFType();
+    if (stt != llvm::ELF::STT_FUNC) continue;
+    llvm::Expected<uint64_t> address = sr.getAddress();
+    if (!address || !*address) continue;
+
+    llvm::Expected<llvm::StringRef> func_name = symbol.getName();
+    if (!func_name || (!func_name->starts_with("__AArch64ADRPThunk_") &&
+                       !func_name->starts_with("__AArch64AbsLongThunk_")))
+      continue;
+
+    const uint64_t func_size = symbol.getSize();
+    if (func_size == 0) continue;
+
+    thunk_map.insert({*address, sr});
+  }
+  return thunk_map;
+}
+
 namespace propeller {
 absl::flat_hash_map<uint64_t, llvm::SmallVector<llvm::object::ELFSymbolRef>>
 ReadSymbolTable(const BinaryContent &binary_content) {
@@ -279,6 +305,14 @@ ReadSymbolTable(const BinaryContent &binary_content) {
   return symtab;
 }
 
+std::optional<absl::btree_map<uint64_t, llvm::object::ELFSymbolRef>>
+ReadThunkSymbols(const BinaryContent &binary_content) {
+  if (binary_content.object_file->getArch() == llvm::Triple::aarch64)
+    return ::ReadAArch64ThunkSymbols(binary_content);
+
+  return std::nullopt;
+}
+
 absl::StatusOr<std::vector<llvm::object::BBAddrMap>> ReadBbAddrMap(
     const BinaryContent &binary_content) {
   auto *elf_object = llvm::dyn_cast<llvm::object::ELFObjectFileBase>(

propeller/binary_content.h

@@ -7,6 +7,7 @@
 #include <string>
 #include <vector>
 
+#include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -116,6 +117,14 @@ absl::StatusOr<int64_t> GetSymbolAddress(
 absl::flat_hash_map<uint64_t, llvm::SmallVector<llvm::object::ELFSymbolRef>>
 ReadSymbolTable(const BinaryContent &binary_content);
 
+// Returns an AArch64 binary's thunk symbols by reading from its symbol table.
+absl::btree_map<uint64_t, llvm::object::ELFSymbolRef> ReadAArch64ThunkSymbols(
+    const BinaryContent &binary_content);
+
+// Returns the binary's thunk symbols by reading from its symbol table.
+std::optional<absl::btree_map<uint64_t, llvm::object::ELFSymbolRef>>
+ReadThunkSymbols(const BinaryContent &binary_content);
+
 // Returns the binary's `BBAddrMap`s by calling LLVM-side decoding function
 // `ELFObjectFileBase::readBBAddrMap`. Returns error if the call fails or if the
 // result is empty.

propeller/binary_content_test.cc

@@ -69,5 +69,23 @@ TEST(GetSymbolAddressTest, SymbolNotFound) {
       Not(IsOk()));
 }
 
+TEST(ThunkSymbolsTest, AArch64Thunks) {
+  const std::string binary = absl::StrCat(
+      ::testing::SrcDir(),
+      "_main/propeller/testdata/fake_thunks.bin");
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<BinaryContent> binary_content,
+                       GetBinaryContent(binary));
+  EXPECT_THAT(ReadThunkSymbols(*binary_content), Optional(SizeIs(2)));
+}
+
+TEST(ThunkSymbolsTest, x86NoThunks) {
+  const std::string binary = absl::StrCat(
+      ::testing::SrcDir(),
+      "_main/propeller/testdata/propeller_sample_1.bin");
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<BinaryContent> binary_content,
+                       GetBinaryContent(binary));
+  EXPECT_THAT(ReadThunkSymbols(*binary_content), Eq(std::nullopt));
+}
+
 }  // namespace
 }  // namespace propeller

propeller/testdata/BUILD

@@ -36,6 +36,7 @@ exports_files([
     "bimodal_sample.x.bin",
     "call_from_simple_loop.protobuf",
     "clang_v0_labels.binary",
+    "fake_thunks.bin",
     "hot_and_cold_landing_pads.protobuf",
     "libro_sample.so",
     "llvm_function_samples.binary",
@@ -352,3 +353,30 @@ genrule(
     target_compatible_with = ["//third_party/bazel_platforms/cpu:aarch64"],
     tools = [":sample"],
 )
+
+cc_binary(
+    name = "fake_thunks",
+    srcs = ["fake_thunks.c"],
+)
+
+# This rule can be used to manually generate fake_thunks.bin.
+# Do not make this into the deps of tests since regenerating the file every time is cumbersome.
+#
+# To build `fake_thunks_bin`, you need to target Arm. If building with Propeller
+# annotations, pass `--host_features=propeller_annotate` when building the genrule.
+genrule(
+    name = "fake_thunks_bin",
+    srcs = ["fake_thunks.c"],
+    outs = [
+        "fake_thunks.bin.gen",
+    ],
+    cmd = "$(CC) $(CC_FLAGS) -g -O2 -Wl,-build-id -pie -fbasic-block-sections=labels $< -o " +
+          "$(RULEDIR)/fake_thunks.bin.gen",
+    exec_compatible_with = ["//third_party/bazel_platforms/cpu:aarch64"],
+    tags = [
+        "manual",
+        "requires-arch:arm",
+    ],
+    target_compatible_with = ["//third_party/bazel_platforms/cpu:aarch64"],
+    toolchains = _LLVM_PROPELLER_TESTDATA_TOOLCHAINS,
+)

propeller/testdata/fake_thunks.bin

@@ -0,0 +1,19 @@
+/* fake_thunks.c */
+volatile int x = 1;
+
+__attribute__((noinline)) int __AArch64ADRPThunk_test1(int i) {
+  return x + i;
+}
+
+__attribute__((noinline)) int __AArch64ADRPThunk_test2(int i) {
+  return x + i + 1;
+}
+
+int sample1_func() { return 13; }
+
+int main(void) {
+  __AArch64ADRPThunk_test1(x);
+  __AArch64ADRPThunk_test2(x);
+
+  return 0;
+}