[BOLT][AArch64] Do not crash on authenticated branch instructions #129898

atrosinenko · 2025-03-05T16:32:14Z

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.

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.

llvmbot · 2025-03-05T16:32:54Z

@llvm/pr-subscribers-bolt

Author: Anatoly Trosinenko (atrosinenko)

Changes

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.

Full diff: https://github.com/llvm/llvm-project/pull/129898.diff

2 Files Affected:

(modified) bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp (+12)
(modified) bolt/test/AArch64/test-indirect-branch.s (+22-1)

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",@progbits
 datatable:
   .word test1_0-datatable

yota9

LGTM Thanks

yavtuk · 2025-03-06T09:01:58Z

folks how about to use isIndirectBranchOpcode func from AArch64InstrInfo.h file?
Do we need own implementation ?

yavtuk · 2025-03-06T10:42:07Z

One more point from my side, jump tables for aarch64 can be implemented over br & blr instructions. We should also include the checking for indirect call instructions
case AArch64::BLRAA:
case AArch64::BLRAB:
case AArch64::BLRAAZ:
case AArch64::BLRABZ:

atrosinenko · 2025-03-06T12:35:33Z

folks how about to use isIndirectBranchOpcode func from AArch64InstrInfo.h file?
Do we need own implementation ?

@yavtuk Do you suggest replacing isBRA with a call to isIndirectBranchOpcode? The difference is that isBRA only checks for Branch to register, with pointer authentication instructions defined in Armv8.3-a (FEAT_PAUTH), while isIndirectBranchOpcode checks for both AArch64::BR and its "authenticated" variants, as expected:

static inline bool isIndirectBranchOpcode(int Opc) {
  switch (Opc) {
  case AArch64::BR:
  case AArch64::BRAA:
  case AArch64::BRAB:
  case AArch64::BRAAZ:
  case AArch64::BRABZ:
    return true;
  }
  return false;
}

That is, isIndirectBranchOpcode intentionally includes AArch64::BR and isBRA intentionally excludes AArch64::BR.

Though, I agree that there may exist some redundancy in AArch64MCPlusBuilder: for example, AArch64MCPlusBuilder::mayLoad is implemented like this:

  bool mayLoad(const MCInst &Inst) const override {
    return isLDRB(Inst) || isLDRH(Inst) || isLDRW(Inst) || isLDRX(Inst) ||
           isLDRQ(Inst) || isLDRD(Inst) || isLDRS(Inst);
  }

TableGen-erated descriptions from AArch64 backend could probably be re-used instead in mayLoad, but I'm not sure.

atrosinenko · 2025-03-06T12:44:12Z

One more point from my side, jump tables for aarch64 can be implemented over br & blr instructions. We should also include the checking for indirect call instructions
case AArch64::BLRAA:
case AArch64::BLRAB:
case AArch64::BLRAAZ:
case AArch64::BLRABZ:

@yavtuk Could you clarify this? If I understand the existing analyzeIndirectBranchFragment function correctly, it only matches the sequences ending in br $xN where $xN contains a sum of some base address and an addend, so I doubt authenticated branch instruction can be used instead of a plain AArch64::BR as the last instruction.

yavtuk · 2025-03-06T14:04:59Z

folks how about to use isIndirectBranchOpcode func from AArch64InstrInfo.h file?
Do we need own implementation ?

@yavtuk Do you suggest replacing isBRA with a call to isIndirectBranchOpcode? The difference is that isBRA only checks for Branch to register, with pointer authentication instructions defined in Armv8.3-a (FEAT_PAUTH), while isIndirectBranchOpcode checks for both AArch64::BR and its "authenticated" variants, as expected:
static inline bool isIndirectBranchOpcode(int Opc) {
  switch (Opc) {
  case AArch64::BR:
  case AArch64::BRAA:
  case AArch64::BRAB:
  case AArch64::BRAAZ:
  case AArch64::BRABZ:
    return true;
  }
  return false;
}
That is, isIndirectBranchOpcode intentionally includes AArch64::BR and isBRA intentionally excludes AArch64::BR.

Though, I agree that there may exist some redundancy in AArch64MCPlusBuilder: for example, AArch64MCPlusBuilder::mayLoad is implemented like this:
  bool mayLoad(const MCInst &Inst) const override {
    return isLDRB(Inst) || isLDRH(Inst) || isLDRW(Inst) || isLDRX(Inst) ||
           isLDRQ(Inst) || isLDRD(Inst) || isLDRS(Inst);
  }
TableGen-erated descriptions from AArch64 backend could probably be re-used instead in mayLoad, but I'm not sure.

as example I think we can create the different function isIndirectAuthBranch or we can use if (isIndirectBranchOpcode() & !Inst.getOpcode() != AArch64::BR)

If we can avoid duplication then let try to do it, if not you can leave it as is

yavtuk · 2025-03-06T14:08:37Z

One more point from my side, jump tables for aarch64 can be implemented over br & blr instructions. We should also include the checking for indirect call instructions
case AArch64::BLRAA:
case AArch64::BLRAB:
case AArch64::BLRAAZ:
case AArch64::BLRABZ:

@yavtuk Could you clarify this? If I understand the existing analyzeIndirectBranchFragment function correctly, it only matches the sequences ending in br $xN where $xN contains a sum of some base address and an addend, so I doubt authenticated branch instruction can be used instead of a plain AArch64::BR as the last instruction.

Current implementation for aarch64 supports only BR as indirect instruction for Jump tables, but I am working on the jump tables functionality right now, and I see a lot of JT with BLR instruction. If you are going to use BLRAA or other instructions it's worth to add checking for them as well

atrosinenko · 2025-03-10T10:47:35Z

Current implementation for aarch64 supports only BR as indirect instruction for Jump tables, but I am working on the jump tables functionality right now, and I see a lot of JT with BLR instruction. If you are going to use BLRAA or other instructions it's worth to add checking for them as well

Now got it, thank you! I thought these were the same four AArch64::BRA* instructions, but these were their "with link" counterparts. Could you please provide an example of such JT lowering - using "branch with link" looks quite counter-intuitive at first.

as example I think we can create the different function isIndirectAuthBranch or we can use if (isIndirectBranchOpcode() & !Inst.getOpcode() != AArch64::BR)

If we can avoid duplication then let try to do it, if not you can leave it as is

Testing for isIndirectBranchOpcode() & !Inst.getOpcode() != AArch64::BR looks like replacing

    assert(Inst.getOpcode() == AArch64::BR && "Unexpected opcode");

with

    if (Inst.getOpcode() != AArch64::BR)
      return;

This could make sense as "there are many variants of Inst.getOpcode(), and we only handle this specific one", but an assertion looks reasonable as it somewhat documents "we handle this opcode, we know about these opcodes and intentionally ignore them and we do not expect anything else here". Considering defining yet another isIndirectAuthBranch function - I would prefer to stick to the existing approach of defining BOLT-specific isXYZ(Opcode) for this PR (I briefly grepped the existing existing sources and did not find other possible users of this function so far), but I spotted that multiple styles for isXYZ(Opcode) functions are already used in AArch64MCPlusBuilder.cpp, so I changed isBRA to use more concise switch-case style.

yavtuk · 2025-03-10T14:27:39Z

Could you please provide an example of such JT lowering - using "branch with link" looks quite counter-intuitive at first.

@atrosinenko the simplified version with array of functions pointers,
pattern can be different but blr the latest instruction, sorry if source code format is uncontrolled

clang -fjump-tables
gcc --param=case-values-threshold=4


int __attribute__ ((noinline)) option0() {
       printf("Function 1 called\n");
       return 1;
}

int (*jumpTable[6])(void) = { option0, option1, option2, option3, option4, option5 };

void __attribute__ ((noinline)) handleOptionJumpTable(int option) {
       const int result = jumpTable[option]();
       printf("result %d\n", result);
}

objdump --disassemble=handleOptionJumpTable --no-addresses --no-show-raw-insn ./jt

<main>
        ldur    w0, [x29, #-8]
        bl      <handleOptionJumpTable>

<handleOptionJumpTable>:
       prologue

        adrp    x8, jumpTable
        ldr     x8, [x8, #3656]
        ldr     x8, [x8, x0, lsl #3]
        blr     x8

        epilogue
        ret

atrosinenko · 2025-03-11T09:59:07Z

        adrp    x8, jumpTable
        ldr     x8, [x8, #3656]
        ldr     x8, [x8, x0, lsl #3]
        blr     x8

looks like a natural way to express

       const int result = jumpTable[option]();
       printf("result %d\n", result);

in machine code, as blr calls a function here. I think such hand-written emulation of jump tables is far out of scope for analyzeIndirectBranchFragment.

atrosinenko · 2025-03-14T09:52:35Z

I assume all the discussions on this PR are resolved now. @yavtuk I will merge this if you don't mind?

yavtuk · 2025-03-16T09:12:45Z

Yes, you can, thanks

atrosinenko requested review from asl and kbeyls March 5, 2025 16:32

atrosinenko requested review from aaupov, maksfb, rafaelauler, ayermolo, dcci and yota9 as code owners March 5, 2025 16:32

llvmbot added the BOLT label Mar 5, 2025

yota9 approved these changes Mar 5, 2025

View reviewed changes

Use switch-case instead of or-ed comparisons

f3b42ad

atrosinenko merged commit 4f2ee07 into llvm:main Mar 17, 2025
10 checks passed

atrosinenko deleted the bolt-pauth-analyze-indirect-branch-crash branch March 17, 2025 09:00

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[BOLT][AArch64] Do not crash on authenticated branch instructions #129898

[BOLT][AArch64] Do not crash on authenticated branch instructions #129898

Uh oh!

atrosinenko commented Mar 5, 2025

Uh oh!

llvmbot commented Mar 5, 2025

Uh oh!

yota9 left a comment

Uh oh!

yavtuk commented Mar 6, 2025

Uh oh!

yavtuk commented Mar 6, 2025

Uh oh!

atrosinenko commented Mar 6, 2025

Uh oh!

atrosinenko commented Mar 6, 2025

Uh oh!

yavtuk commented Mar 6, 2025 •

edited

Loading

Uh oh!

yavtuk commented Mar 6, 2025

Uh oh!

atrosinenko commented Mar 10, 2025

Uh oh!

yavtuk commented Mar 10, 2025 •

edited

Loading

Uh oh!

atrosinenko commented Mar 11, 2025

Uh oh!

atrosinenko commented Mar 14, 2025

Uh oh!

yavtuk commented Mar 16, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

[BOLT][AArch64] Do not crash on authenticated branch instructions #129898

[BOLT][AArch64] Do not crash on authenticated branch instructions #129898

Uh oh!

Conversation

atrosinenko commented Mar 5, 2025

Uh oh!

llvmbot commented Mar 5, 2025

Uh oh!

yota9 left a comment

Choose a reason for hiding this comment

Uh oh!

yavtuk commented Mar 6, 2025

Uh oh!

yavtuk commented Mar 6, 2025

Uh oh!

atrosinenko commented Mar 6, 2025

Uh oh!

atrosinenko commented Mar 6, 2025

Uh oh!

yavtuk commented Mar 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yavtuk commented Mar 6, 2025

Uh oh!

atrosinenko commented Mar 10, 2025

Uh oh!

yavtuk commented Mar 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

atrosinenko commented Mar 11, 2025

Uh oh!

atrosinenko commented Mar 14, 2025

Uh oh!

yavtuk commented Mar 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yavtuk commented Mar 6, 2025 •

edited

Loading

yavtuk commented Mar 10, 2025 •

edited

Loading

yavtuk commented Mar 16, 2025 •

edited

Loading