[llvm] [CrossDSO CFI] Make sure __cfi_check has BTI attribute. (PR #131224)

[Eli Friedman via llvm-commits]

https://github.com/efriedma-quic updated https://github.com/llvm/llvm-project/pull/131224

>From cbdee73cc70696bdf2de3867f35430f7ccb57862 Mon Sep 17 00:00:00 2001
From: Eli Friedman <efriedma at quicinc.com>
Date: Thu, 13 Mar 2025 10:55:18 -0700
Subject: [PATCH 1/2] [CrossDSO CFI] Make sure __cfi_check has BTI attribute.

When we're doing CrossDSO CFI, clang tries to generate a stub for
__cfi_check so it has the correct attributes.  However, this doesn't
work in some scenarios: if you're doing ThinLTO, and some code is forced
into the "fulllto" module, __cfi_check is generated in that module, and
there is no stub.

In this scenario, there's an inconsistency between the
"branch-target-enforcement" module flag, and the function: the function
is missing the attribute.  So the "bti c" isn't generated, and it faults
at runtime.

This fix works around the issue by forcing the
"branch-target-enforcement" attribute onto the function.

See also c7cacb2f6efe07fa4a1129bb7e2389312670b84a.
---
 llvm/lib/Transforms/IPO/CrossDSOCFI.cpp       |  8 +++
 .../Transforms/CrossDSOCFI/aarch64-bti.ll     | 53 +++++++++++++++++++
 2 files changed, 61 insertions(+)
 create mode 100644 llvm/test/Transforms/CrossDSOCFI/aarch64-bti.ll

diff --git a/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp b/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
index 2d884078940cc..55360756f6a92 100644
--- a/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
+++ b/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
@@ -94,6 +94,14 @@ void CrossDSOCFI::buildCFICheck(Module &M) {
   Triple T(M.getTargetTriple());
   if (T.isARM() || T.isThumb())
     F->addFnAttr("target-features", "+thumb-mode");
+  if (T.isAArch64()) {
+    if (const auto *BTE = mdconst::extract_or_null<ConstantInt>(
+            M.getModuleFlag("branch-target-enforcement"))) {
+      if (BTE->getZExtValue() != 0) {
+        F->addFnAttr("branch-target-enforcement");
+      }
+    }
+  }
 
   auto args = F->arg_begin();
   Value &CallSiteTypeId = *(args++);
diff --git a/llvm/test/Transforms/CrossDSOCFI/aarch64-bti.ll b/llvm/test/Transforms/CrossDSOCFI/aarch64-bti.ll
new file mode 100644
index 0000000000000..baed432f874c8
--- /dev/null
+++ b/llvm/test/Transforms/CrossDSOCFI/aarch64-bti.ll
@@ -0,0 +1,53 @@
+; RUN: opt -S -passes=cross-dso-cfi < %s | FileCheck %s
+
+; CHECK: define void @__cfi_check({{.*}}) [[ATTR:#[0-9]+]] align 4096
+; CHECK: [[ATTR]] = { "branch-target-enforcement" }
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+ at _ZTV1A = constant i8 0, !type !4, !type !5
+ at _ZTV1B = constant i8 0, !type !4, !type !5, !type !6, !type !7
+
+define signext i8 @f11() "branch-target-enforcement" !type !0 !type !1 {
+entry:
+  ret i8 1
+}
+
+define signext i8 @f12() "branch-target-enforcement" !type !0 !type !1 {
+entry:
+  ret i8 2
+}
+
+define signext i8 @f13() "branch-target-enforcement" !type !0 !type !1 {
+entry:
+  ret i8 3
+}
+
+define i32 @f21() "branch-target-enforcement" !type !2 !type !3 {
+entry:
+  ret i32 4
+}
+
+define i32 @f22() "branch-target-enforcement" !type !2 !type !3 {
+entry:
+  ret i32 5
+}
+
+define weak_odr hidden void @__cfi_check_fail(ptr, ptr) "branch-target-enforcement" {
+entry:
+  ret void
+}
+
+!llvm.module.flags = !{!8, !9}
+
+!0 = !{i64 0, !"_ZTSFcvE"}
+!1 = !{i64 0, i64 111}
+!2 = !{i64 0, !"_ZTSFivE"}
+!3 = !{i64 0, i64 222}
+!4 = !{i64 16, !"_ZTS1A"}
+!5 = !{i64 16, i64 333}
+!6 = !{i64 16, !"_ZTS1B"}
+!7 = !{i64 16, i64 444}
+!8 = !{i32 4, !"Cross-DSO CFI", i32 1}
+!9 = !{i32 8, !"branch-target-enforcement", i32 1}

>From 201600fa94449eaa1d7f713dfcb07e66970747c1 Mon Sep 17 00:00:00 2001
From: Eli Friedman <efriedma at quicinc.com>
Date: Fri, 14 Mar 2025 14:00:46 -0700
Subject: [PATCH 2/2] Use createWithDefaultAttr instead of explicitly modifying
 attributes.

Per review comment.
---
 llvm/lib/Transforms/IPO/CrossDSOCFI.cpp | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp b/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
index 55360756f6a92..d94312098182b 100644
--- a/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
+++ b/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
@@ -82,14 +82,26 @@ void CrossDSOCFI::buildCFICheck(Module &M) {
   }
 
   LLVMContext &Ctx = M.getContext();
-  FunctionCallee C = M.getOrInsertFunction(
-      "__cfi_check", Type::getVoidTy(Ctx), Type::getInt64Ty(Ctx),
-      PointerType::getUnqual(Ctx), PointerType::getUnqual(Ctx));
+  FunctionType *CFICheckTy =
+      FunctionType::get(Type::getVoidTy(Ctx),
+                        {Type::getInt64Ty(Ctx), PointerType::getUnqual(Ctx),
+                         PointerType::getUnqual(Ctx)},
+                        false);
+  FunctionCallee C = Function::createWithDefaultAttr(
+      CFICheckTy, GlobalValue::ExternalLinkage, 0, "__cfi_check", &M);
   Function *F = cast<Function>(C.getCallee());
-  // Take over the existing function. The frontend emits a weak stub so that the
-  // linker knows about the symbol; this pass replaces the function body.
-  F->deleteBody();
   F->setAlignment(Align(4096));
+  if (F->getName() != "__cfi_check") {
+    // The frontend might have already created a function named __cfi_check;
+    // delete it.
+    GlobalValue *G = M.getNamedValue("__cfi_check");
+    assert(G && "cfi_check must exist after we constructed it");
+    if (G->getAddressSpace() != F->getAddressSpace())
+      report_fatal_error("__cfi_check with unexpected address space");
+    G->replaceAllUsesWith(F);
+    F->takeName(G);
+    G->eraseFromParent();
+  }
 
   Triple T(M.getTargetTriple());
   if (T.isARM() || T.isThumb())