[LowerTypeTests] Add ENDBR to .cfi.jumptable for x86 Indirect Branch Tracking

[MaskRay]

Summary

Similar to D81251 for AArch64 BTI. This fixes ./a.out test for

void foo(void) {}
void bar(void) {}
static void (*fptr)(void);
int main(int argc, char **argv) {
  if (argv[1]) fptr = foo;
  else fptr = bar;
  fptr();
}

clang -flto=thin -fvisibility=hidden -fsanitize=cfi-icall -fcf-protection=branch -fuse-ld=lld a.cc

[q66]

tested here and works great

[tejohnson]

I looked up x86 ENDBR and the adding of that instruction seems fine to me, but I have a couple questions on the related changes in the patch.

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1231

Why is this needed only in the endbr case?

1419

Add comment about why this attribute is needed. I see that the attribute is only checked for the x86 backends. But can be queried on the function for any arch - is there any possibility this feature will be implemented for non-x86 in the future? If so, I suggest adding the attribute unconditionally here.

Also, is this needed in the non-endbr case for x86? I.e. it is being added here for all x86 which presumably also changes the pre-existing non-endbr x86 handling.

[MaskRay]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1231

The endbr case has a larger jump table entry size (16 bytes). We can add as many INT3 as needed, but .balign 16, 0xcc seems easier.

1419

This is similar to F->addFnAttr("branch-target-enforcement", "false"); for AArch64 BTI (D81251). Without this attribute, the function will have 2 consecutive endbr64 at the function start (one from the BTI instrumentation, the other from the first instruction of the inline asm instruction).

[tejohnson]

lgtm with comment or change suggested below.

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1419

Please add a comment about why it is needed and that it does not have any effect in the case where CF branch protection is not enabled (alternatively, just guard this with whether the cf-protection-branch module flag is set).

[MaskRay]

add comments for ENDBR and BTI

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

void LowerTypeTestsModule::verifyTypeMDNode(GlobalObject *GO, MDNode *Type) {
   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 kX86IBTJumpTableEntrySize = 16;
 static const unsigned kARMJumpTableEntrySize = 4;
 static const unsigned kARMBTIJumpTableEntrySize = 8;
 static const unsigned kRISCVJumpTableEntrySize = 8;
 
 unsigned LowerTypeTestsModule::getJumpTableEntrySize() {
   switch (Arch) {
     case Triple::x86:
     case Triple::x86_64:
+      if (const auto *MD = mdconst::extract_or_null<ConstantInt>(
+            M.getModuleFlag("cf-protection-branch")))
+        if (MD->getZExtValue())
+          return kX86IBTJumpTableEntrySize;
       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())

 // 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) {
+    bool Endbr = false;
+    if (const auto *MD = mdconst::extract_or_null<ConstantInt>(
+          Dest->getParent()->getModuleFlag("cf-protection-branch")))
+      Endbr = MD->getZExtValue() != 0;
+    if (Endbr)
+      AsmOS << (JumpTableArch == Triple::x86 ? "endbr32\n" : "endbr64\n");
     AsmOS << "jmp ${" << ArgIndex << ":c}@plt\n";
+    if (Endbr)
+      AsmOS << ".balign 16, 0xcc\n";
+    else
       AsmOS << "int3\nint3\nint3\n";
   } 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";
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");
   }
+  // When -mbranch-protection= is used, the inline asm adds a BTI. Suppress BTI
+  // for the function to avoid double BTI. This is a no-op without
+  // -mbranch-protection=.
   if (JumpTableArch == Triple::aarch64) {
     F->addFnAttr("branch-target-enforcement", "false");
     F->addFnAttr("sign-return-address", "none");
   }
   if (JumpTableArch == Triple::riscv32 || JumpTableArch == Triple::riscv64) {
     // Make sure the jump table assembly is not modified by the assembler or
     // the linker.
     F->addFnAttr("target-features", "-c,-relax");
   }
+  // When -fcf-protection= is used, the inline asm adds an ENDBR. Suppress ENDBR
+  // for the function to avoid double ENDBR. This is a no-op without
+  // -fcf-protection=.
+  if (JumpTableArch == Triple::x86 || JumpTableArch == Triple::x86_64)
+    F->addFnAttr(Attribute::NoCfCheck);
   // 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/function.ll

 ; THUMB:      b.w $0
 ; THUMB-SAME: b.w $1
 
 ; RISCV:      tail $0@plt
 ; RISCV-SAME: tail $1@plt
 
 ; NATIVE-SAME: "s,s"(ptr @f.cfi, ptr @g.cfi)
 
-; X86-LINUX: attributes #[[ATTR]] = { naked nounwind }
-; X86-WIN32: attributes #[[ATTR]] = { nounwind }
+; X86-LINUX: attributes #[[ATTR]] = { naked nocf_check nounwind }
+; X86-WIN32: attributes #[[ATTR]] = { nocf_check nounwind }
 ; ARM: attributes #[[ATTR]] = { naked nounwind
 ; THUMB: attributes #[[ATTR]] = { naked nounwind "target-cpu"="cortex-a8" "target-features"="+thumb-mode" }
 ; RISCV: attributes #[[ATTR]] = { naked nounwind "target-features"="-c,-relax" }
 
 ; WASM32: ![[I0]] = !{i64 1}
 ; WASM32: ![[I1]] = !{i64 2}

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

+;; Test jump table generation with Indirect Branch Tracking on x86.
+; RUN: opt -S -passes=lowertypetests -mtriple=i686 %s | FileCheck --check-prefixes=X86,X86_32 %s
+; RUN: opt -S -passes=lowertypetests -mtriple=x86_64 %s | FileCheck --check-prefixes=X86,X86_64 %s
+
+@0 = private unnamed_addr constant [2 x ptr] [ptr @f, ptr @g], align 16
+
+define void @f() !type !0 {
+  ret void
+}
+
+define internal void @g() !type !0 {
+  ret void
+}
+
+declare i1 @llvm.type.test(ptr %ptr, metadata %bitset) nounwind readnone
+
+define i1 @foo(ptr %p) {
+  %x = call i1 @llvm.type.test(ptr %p, metadata !"typeid1")
+  ret i1 %x
+}
+
+!llvm.module.flags = !{!1}
+!0 = !{i32 0, !"typeid1"}
+!1 = !{i32 8, !"cf-protection-branch", i32 1}
+
+; X86:         define private void @.cfi.jumptable() #[[#ATTR:]] align 16 {
+; X86-NEXT:    entry:
+; X86_32-NEXT:   call void asm sideeffect "endbr32\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0Aendbr32\0Ajmp ${1:c}@plt\0A.balign 16, 0xcc\0A", "s,s"(ptr @f.cfi, ptr @g.cfi)
+; X86_64-NEXT:   call void asm sideeffect "endbr64\0Ajmp ${0:c}@plt\0A.balign 16, 0xcc\0Aendbr64\0Ajmp ${1:c}@plt\0A.balign 16, 0xcc\0A", "s,s"(ptr @f.cfi, ptr @g.cfi)
+
+; X86_64: attributes #[[#ATTR]] = { naked nocf_check nounwind }