[AArch64] Add BTI to CFI jumptables.

[danielkiss]

Summary

With branch protection the jump to the jump table entries requires a landing pad.

[danielkiss]

This patch depends on D80791 and D75181.

[eugenis]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1220

I'd rather this was a named constant.

1244

Do we need to check the function attribute here as well? What happens if a function opts out of BTI?

[danielkiss]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1220

I'll update it.

1244

No, direct jump does not need a landing pad. So this jump will be fine even that target function opted out.

[eugenis]

LGTM

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1244

What I meant is, if a function has ignore-branch-target-enforcement, do not we want to skip the BTI hint in that function's CFI stub?
Or, rather, replace it with NOP to preserve the stub size.

Thinking about it some more - no, probably not. The attribute can be used on a "naked" function that already has BTI in the assembly; in that case we'd need BTI in the jump table as well.

[danielkiss]

I need to update it due to changes in the dependent patches.

[chill]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1410

Why are we disabling BTI here?

1411

Here we should use sign-return-address=none, and, if needed, I suggest (re-)introducing branch-target-enforcement=true|false ?

[danielkiss]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1410

This function is a jump table and each entry will have a landing pad ( see line 1235).
If it is not disabled then the entry will get an additional landing pad that miss-aligns the table.

Worst case we could add special condition for the first entry but IMHO that would be ugly.

1411

Agree , I will update once the D85649 is updated as proposed there.

[danielkiss]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1245

since D81257 we can use bti c in assembly.

[danielkiss]

rebasing to the new patch series.

[danielkiss]

update module flag handling due to the change is merged in: D85649.

[tamas.petz]

LGTM after the minor change

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

if (!OffsetConstMD)
     report_fatal_error("Type offset must be a constant");
   auto OffsetInt = dyn_cast<ConstantInt>(OffsetConstMD->getValue());
   if (!OffsetInt)
     report_fatal_error("Type offset must be an integer constant");
 }
 
 static const unsigned kX86JumpTableEntrySize = 8;
 static const unsigned kARMJumpTableEntrySize = 4;
+static const unsigned kARMBTIJumpTableEntrySize = 8;
 
 unsigned LowerTypeTestsModule::getJumpTableEntrySize() {
   switch (Arch) {
     case Triple::x86:
     case Triple::x86_64:
       return kX86JumpTableEntrySize;
     case Triple::arm:
     case Triple::thumb:
+      return kARMJumpTableEntrySize;
     case Triple::aarch64:
+      if (const auto *BTE = mdconst::extract_or_null<ConstantInt>(
+            M.getModuleFlag("branch-target-enforcement")))
+        if (BTE->getZExtValue())
+          return kARMBTIJumpTableEntrySize;
       return kARMJumpTableEntrySize;
     default:
       report_fatal_error("Unsupported architecture for jump tables");
   }
 }
 
 // Create a jump table entry for the target. This consists of an instruction
 // sequence containing a relative branch to Dest. Appends inline asm text,
 // constraints and arguments to AsmOS, ConstraintOS and AsmArgs.
 void LowerTypeTestsModule::createJumpTableEntry(
     raw_ostream &AsmOS, raw_ostream &ConstraintOS,
     Triple::ArchType JumpTableArch, SmallVectorImpl<Value *> &AsmArgs,
     Function *Dest) {
   unsigned ArgIndex = AsmArgs.size();
 
   if (JumpTableArch == Triple::x86 || JumpTableArch == Triple::x86_64) {
     AsmOS << "jmp ${" << ArgIndex << ":c}@plt\n";
     AsmOS << "int3\nint3\nint3\n";
-  } else if (JumpTableArch == Triple::arm || JumpTableArch == Triple::aarch64) {
+  } else if (JumpTableArch == Triple::arm) {
+    AsmOS << "b $" << ArgIndex << "\n";
+  } else if (JumpTableArch == Triple::aarch64) {
+    if (const auto *BTE = mdconst::extract_or_null<ConstantInt>(
+          Dest->getParent()->getModuleFlag("branch-target-enforcement")))
+      if (BTE->getZExtValue())
+        AsmOS << "bti c\n";
     AsmOS << "b $" << ArgIndex << "\n";
   } else if (JumpTableArch == Triple::thumb) {
     AsmOS << "b.w $" << ArgIndex << "\n";
   } else {
     report_fatal_error("Unsupported architecture for jump tables");
   }
 
   ConstraintOS << (ArgIndex > 0 ? ",s" : "s");
void LowerTypeTestsModule::createJumpTable(
   if (JumpTableArch == Triple::arm)
     F->addFnAttr("target-features", "-thumb-mode");
   if (JumpTableArch == Triple::thumb) {
     F->addFnAttr("target-features", "+thumb-mode");
     // Thumb jump table assembly needs Thumb2. The following attribute is added
     // by Clang for -march=armv7.
     F->addFnAttr("target-cpu", "cortex-a8");
   }
+  if (JumpTableArch == Triple::aarch64) {
+    F->addFnAttr("branch-target-enforcement", "false");
+    F->addFnAttr("sign-return-address", "none");
+  }
   // Make sure we don't emit .eh_frame for this function.
   F->addFnAttr(Attribute::NoUnwind);
 
   BasicBlock *BB = BasicBlock::Create(M.getContext(), "entry", F);
   IRBuilder<> IRB(BB);
 
   SmallVector<Type *, 16> ArgTypes;
   ArgTypes.reserve(AsmArgs.size());

llvm/test/Transforms/LowerTypeTests/aarch64-jumptable.ll

+; RUN: opt -S -lowertypetests -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck --check-prefixes=AARCH64 %s
+
+; Test for the jump table generation with branch protection on AArch64
+
+target datalayout = "e-p:64:64"
+
+@0 = private unnamed_addr constant [2 x void (...)*] [void (...)* bitcast (void ()* @f to void (...)*), void (...)* bitcast (void ()* @g to void (...)*)], align 16
+
+; AARCH64: @f = alias void (), void ()* @[[JT:.*]]
+
+define void @f() !type !0 {
+  ret void
+}
+
+define internal void @g() !type !0 {
+  ret void
+}
+
+!0 = !{i32 0, !"typeid1"}
+
+declare i1 @llvm.type.test(i8* %ptr, metadata %bitset) nounwind readnone
+
+define i1 @foo(i8* %p) {
+  %x = call i1 @llvm.type.test(i8* %p, metadata !"typeid1")
+  ret i1 %x
+}
+
+!llvm.module.flags = !{!1}
+
+!1 = !{i32 4, !"branch-target-enforcement", i32 1}
+
+; AARCH64:   define private void @[[JT]]() #[[ATTR:.*]] align 8 {
+
+; AARCH64:      bti c
+; AARCH64-SAME: b $0
+; AARCH64-SAME: bti c
+; AARCH64-SAME: b $1
+
+; AARCH64: attributes #[[ATTR]] = { naked nounwind "branch-target-enforcement"="false" "sign-return-address"="none"