[CrossDSO CFI] Make sure __cfi_check has BTI attribute.

[efriedma-quic]

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 attribute: the function attribute is missing. 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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Eli Friedman (efriedma-quic)

Changes

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 attribute: the function attribute is missing. 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.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/IPO/CrossDSOCFI.cpp (+9)
• (added) llvm/test/Transforms/CrossDSOCFI/aarch64-bti.ll (+53)

diff --git a/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp b/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
index 2d884078940cc..813adeba75564 100644
--- a/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
+++ b/llvm/lib/Transforms/IPO/CrossDSOCFI.cpp
@@ -94,6 +94,15 @@ 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"
+
+@_ZTV1A = constant i8 0, !type !4, !type !5
+@_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}

[efriedma-quic]

I'm not confident this patch is the best solution, but I'm not sure what a better solution looks like; should clang be setting the relevant function attributes somehow?

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[pirama-arumuga-nainar]

Adding @pcc as well.

[DanielKristofKiss]

M.getOrInsertFunction could be replaced with the Function::createWithDefaultAttr to create the functions with the right set of attributes. It is used in other places in the sanitisers to create functions. It will solve this and future problems.

[efriedma-quic]

Gave this a shot. Seems to be okay for aarch64, but I'm a little concerned this will cause issues for other targets (see #98673).

[pcc]

The way I think this should work is that the function is created by the frontend with the correct attributes and then filled in by the LTO backend. See #100833 for details. That approach is compatible with #98673 because the attributes are still coming from the frontend.

[efriedma-quic]

I briefly described why that doesn't work in the commit message, but I'll go into more detail:

In ThinLTO, the code is in modules which correspond to translation units in the original source code. However, depending on the exact structure of the inputs, you can end up with an additional "base" module that only contains some global variables, and no code. LTO chooses to generate __cfi_check into that module; when it does, there is no function to take the attributes from.

So there's a function named __cfi_check, created by the frontend, in every module except the one that actually matters. So I'm just trying to repair the IR the best I can once we've reached this point.

If you think something is supposed to happen earlier in the pipeline, can you describe what, exactly, that looks like?

[pcc]

My expectation would be that __cfi_check always appears on the full LTO side and has the correct attributes. If we end up generating code for that function without attributes, it sounds like a bug. Do you have a reproducer showing how __cfi_check can end up being code generated without attributes?

[efriedma-quic]

$ cat a.c
int a = 3;
void f(){}
$ cat b.c
void g(){}
$ clang a.c b.c -flto=thin -fsanitize=cfi -fsanitize-cfi-cross-dso -O2 -c -fvisibility=default
$ clang a.o b.o -flto=thin -fsanitize=cfi -fsanitize-cfi-cross-dso -O2 -fuse-ld=lld -shared -Wl,-save-temps -fvisibility=default

__cfi_check is generated in a.out.0.5.precodegen.bc, which doesn't contain a stub.

[pcc]

Thanks. With the attached patch (not extensively tested) it looks like __cfi_check now goes to the correct module. Next step to fix this bug would be to resolve issues pointed out in #100833 so that the frontend gives it the correct attributes.

0001-ThinLTOBitcodeWriter-Move-__cfi_check-to-merged-modu.patch.txt

[DanielKristofKiss]

Thanks, LGTM

[teresajohnson]

Added @vitalybuka and @pcc who are more familiar with CFI. The change looks ok to me but I am not sure whether this should have been handled elsewhere.