[llvm] [BOLT][AArch64] Do not crash on authenticated branch instructions (PR #129898)

[Anatoly Trosinenko via llvm-commits]

https://github.com/atrosinenko updated https://github.com/llvm/llvm-project/pull/129898

>From 7a97754a668b69b37041987566ef818df6404392 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Tue, 4 Mar 2025 21:59:38 +0300
Subject: [PATCH 1/2] [BOLT][AArch64] Do not crash on authenticated branch
 instructions

When an indirect branch instruction is decoded, analyzeIndirectBranch
method is asked if this is a well-known code pattern. On AArch64, the
only special pattern which is detected is Jump Table, emitted as a
branch to the sum of a constant base address and a variable offset.
Therefore, `Inst.getOpcode()` being one of `AArch64::BRA*` means Inst
cannot belong to such Jump Table pattern, thus returning early.
---
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   | 12 ++++++++++
 bolt/test/AArch64/test-indirect-branch.s      | 23 ++++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 685b2279e5afb..e073b197e70d2 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -547,6 +547,13 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return false;
   }
 
+  bool isBRA(const MCInst &Inst) const {
+    return (Inst.getOpcode() == AArch64::BRAA ||
+            Inst.getOpcode() == AArch64::BRAB ||
+            Inst.getOpcode() == AArch64::BRAAZ ||
+            Inst.getOpcode() == AArch64::BRABZ);
+  }
+
   bool mayLoad(const MCInst &Inst) const override {
     return isLDRB(Inst) || isLDRH(Inst) || isLDRW(Inst) || isLDRX(Inst) ||
            isLDRQ(Inst) || isLDRD(Inst) || isLDRS(Inst);
@@ -941,6 +948,11 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
       DenseMap<const MCInst *, SmallVector<MCInst *, 4>> &UDChain,
       const MCExpr *&JumpTable, int64_t &Offset, int64_t &ScaleValue,
       MCInst *&PCRelBase) const {
+    // The only kind of indirect branches we match is jump table, thus ignore
+    // authenticating branch instructions early.
+    if (isBRA(Inst))
+      return false;
+
     // Expect AArch64 BR
     assert(Inst.getOpcode() == AArch64::BR && "Unexpected opcode");
 
diff --git a/bolt/test/AArch64/test-indirect-branch.s b/bolt/test/AArch64/test-indirect-branch.s
index 1e16e76b11530..b99737ee97acc 100644
--- a/bolt/test/AArch64/test-indirect-branch.s
+++ b/bolt/test/AArch64/test-indirect-branch.s
@@ -5,7 +5,7 @@
 
 // REQUIRES: system-linux, asserts
 
-// RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+// RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown -mattr=+pauth %s -o %t.o
 // RUN: %clang %cflags --target=aarch64-unknown-linux %t.o -o %t.exe -Wl,-q
 // RUN: llvm-bolt %t.exe -o %t.bolt --print-cfg --strict --debug-only=mcplus \
 // RUN:  -v=1 2>&1 | FileCheck %s
@@ -73,6 +73,27 @@ test2_0:
 test2_1:
   ret
 
+// Make sure BOLT does not crash trying to disassemble BRA* instructions.
+  .globl test_braa
+  .type  test_braa, %function
+test_braa:
+  braa x0, x1
+
+  .globl test_brab
+  .type  test_brab, %function
+test_brab:
+  brab x0, x1
+
+  .globl test_braaz
+  .type  test_braaz, %function
+test_braaz:
+  braaz x0
+
+  .globl test_brabz
+  .type  test_brabz, %function
+test_brabz:
+  brabz x0
+
   .section .rodata,"a", at progbits
 datatable:
   .word test1_0-datatable

>From f3b42ad8116ba0ba6c36d4226c209b24f7b97d8c Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko <atrosinenko at accesssoftek.com>
Date: Mon, 10 Mar 2025 13:15:17 +0300
Subject: [PATCH 2/2] Use switch-case instead of or-ed comparisons

---
 bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index e073b197e70d2..a4a6afe6289bc 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -548,10 +548,15 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
   }
 
   bool isBRA(const MCInst &Inst) const {
-    return (Inst.getOpcode() == AArch64::BRAA ||
-            Inst.getOpcode() == AArch64::BRAB ||
-            Inst.getOpcode() == AArch64::BRAAZ ||
-            Inst.getOpcode() == AArch64::BRABZ);
+    switch (Inst.getOpcode()) {
+    case AArch64::BRAA:
+    case AArch64::BRAB:
+    case AArch64::BRAAZ:
+    case AArch64::BRABZ:
+      return true;
+    default:
+      return false;
+    }
   }
 
   bool mayLoad(const MCInst &Inst) const override {