[UBSan] Use -fsanitize-handler-preserve-all-regs in codegen

Pull Request: https://github.com/llvm/llvm-project/pull/168645
2026-01-19 09:31:59 +08:00 · 2025-11-19 17:12:42 -08:00
parent b7eb9883dc
commit e2a29eca56
11 changed files with 84 additions and 14 deletions
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -1134,6 +1134,8 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
                  CodeGenOpts.SanitizeMinimalRuntime),
              /*MayReturn=*/
              CodeGenOpts.SanitizeRecover.has(SanitizerKind::LocalBounds),
+              /*HandlerPreserveAllRegs=*/
+              static_cast<bool>(CodeGenOpts.SanitizeHandlerPreserveAllRegs),
          };
        }
        FPM.addPass(BoundsCheckingPass(Options));
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -3819,6 +3819,8 @@ static void emitCheckHandlerCall(CodeGenFunction &CGF,
  bool NeedsAbortSuffix =
      IsFatal && RecoverKind != CheckRecoverableKind::Unrecoverable;
  bool MinimalRuntime = CGF.CGM.getCodeGenOpts().SanitizeMinimalRuntime;
+  bool HandlerPreserveAllRegs =
+      CGF.CGM.getCodeGenOpts().SanitizeHandlerPreserveAllRegs;
  const SanitizerHandlerInfo &CheckInfo = SanitizerHandlers[CheckHandler];
  const StringRef CheckName = CheckInfo.Name;
  std::string FnName = "__ubsan_handle_" + CheckName.str();
@@ -3828,6 +3830,8 @@ static void emitCheckHandlerCall(CodeGenFunction &CGF,
    FnName += "_minimal";
  if (NeedsAbortSuffix)
    FnName += "_abort";
+  if (HandlerPreserveAllRegs && !NeedsAbortSuffix)
+    FnName += "_preserve";
  bool MayReturn =
      !IsFatal || RecoverKind == CheckRecoverableKind::AlwaysRecoverable;

@@ -3848,6 +3852,10 @@ static void emitCheckHandlerCall(CodeGenFunction &CGF,
            (CGF.CurCodeDecl && CGF.CurCodeDecl->hasAttr<OptimizeNoneAttr>());
  if (NoMerge)
    HandlerCall->addFnAttr(llvm::Attribute::NoMerge);
+  if (HandlerPreserveAllRegs && !NeedsAbortSuffix) {
+    // N.B. there is also a clang::CallingConv which is not what we want here.
+    HandlerCall->setCallingConv(llvm::CallingConv::PreserveAll);
+  }
  if (!MayReturn) {
    HandlerCall->setDoesNotReturn();
    CGF.Builder.CreateUnreachable();
--- a/clang/lib/Driver/SanitizerArgs.cpp
+++ b/clang/lib/Driver/SanitizerArgs.cpp
@@ -419,6 +419,7 @@ SanitizerArgs::SanitizerArgs(const ToolChain &TC,
  const Driver &D = TC.getDriver();
  SanitizerMask TrappingKinds = parseSanitizeTrapArgs(D, Args, DiagnoseErrors);
  SanitizerMask InvalidTrappingKinds = TrappingKinds & NotAllowedWithTrap;
+  const llvm::Triple &Triple = TC.getTriple();

  MinimalRuntime =
      Args.hasFlag(options::OPT_fsanitize_minimal_runtime,
@@ -426,7 +427,8 @@ SanitizerArgs::SanitizerArgs(const ToolChain &TC,
  HandlerPreserveAllRegs =
      Args.hasFlag(options::OPT_fsanitize_handler_preserve_all_regs,
                   options::OPT_fno_sanitize_handler_preserve_all_regs,
-                   HandlerPreserveAllRegs);
+                   HandlerPreserveAllRegs) &&
+      MinimalRuntime && (Triple.isAArch64() || Triple.isX86_64());

  // The object size sanitizer should not be enabled at -O0.
  Arg *OptLevel = Args.getLastArg(options::OPT_O_Group);
@@ -494,7 +496,6 @@ SanitizerArgs::SanitizerArgs(const ToolChain &TC,
      // -fsanitize=function and -fsanitize=kcfi instrument indirect function
      // calls to load a type hash before the function label. Therefore, an
      // execute-only target doesn't support the function and kcfi sanitizers.
-      const llvm::Triple &Triple = TC.getTriple();
      if (isExecuteOnlyTarget(Triple, Args)) {
        if (SanitizerMask KindsToDiagnose =
                Add & NotAllowedWithExecuteOnly & ~DiagnosedKinds) {
--- a/clang/test/CodeGen/cfi-icall-trap-recover-runtime.c
+++ b/clang/test/CodeGen/cfi-icall-trap-recover-runtime.c
@@ -171,7 +171,7 @@ void xf();
 // PRESERVE_MIN-NEXT:    [[TMP3:%.*]] = call i1 @llvm.type.test(ptr [[TMP2]], metadata !"_ZTSFvE"), !nosanitize [[META10:![0-9]+]]
 // PRESERVE_MIN-NEXT:    br i1 [[TMP3]], label %[[CONT:.*]], label %[[HANDLER_CFI_CHECK_FAIL:.*]], !prof [[PROF11:![0-9]+]], !nosanitize [[META10]]
 // PRESERVE_MIN:       [[HANDLER_CFI_CHECK_FAIL]]:
-// PRESERVE_MIN-NEXT:    call void @__ubsan_handle_cfi_check_fail_minimal() #[[ATTR4:[0-9]+]], !nosanitize [[META10]]
+// PRESERVE_MIN-NEXT:    call preserve_allcc void @__ubsan_handle_cfi_check_fail_minimal_preserve() #[[ATTR4:[0-9]+]], !nosanitize [[META10]]
 // PRESERVE_MIN-NEXT:    br label %[[CONT]], !nosanitize [[META10]]
 // PRESERVE_MIN:       [[CONT]]:
 // PRESERVE_MIN-NEXT:    call void (...) [[TMP2]]()
--- a/clang/test/CodeGenCXX/cfi-vcall-trap-recover-runtime.cpp
+++ b/clang/test/CodeGenCXX/cfi-vcall-trap-recover-runtime.cpp
@@ -127,7 +127,7 @@ struct S1 {
 // PRESERVE_MIN-NEXT:    [[TMP2:%.*]] = call i1 @llvm.type.test(ptr [[VTABLE]], metadata !"all-vtables"), !nosanitize [[META5]]
 // PRESERVE_MIN-NEXT:    br i1 [[TMP1]], label %[[CONT:.*]], label %[[HANDLER_CFI_CHECK_FAIL:.*]], !prof [[PROF6:![0-9]+]], !nosanitize [[META5]]
 // PRESERVE_MIN:       [[HANDLER_CFI_CHECK_FAIL]]:
-// PRESERVE_MIN-NEXT:    call void @__ubsan_handle_cfi_check_fail_minimal() #[[ATTR3:[0-9]+]], !nosanitize [[META5]]
+// PRESERVE_MIN-NEXT:    call preserve_allcc void @__ubsan_handle_cfi_check_fail_minimal_preserve() #[[ATTR3:[0-9]+]], !nosanitize [[META5]]
 // PRESERVE_MIN-NEXT:    br label %[[CONT]], !nosanitize [[META5]]
 // PRESERVE_MIN:       [[CONT]]:
 // PRESERVE_MIN-NEXT:    [[VFN:%.*]] = getelementptr inbounds ptr, ptr [[VTABLE]], i64 0
--- a/clang/test/Driver/fsanitize.c
+++ b/clang/test/Driver/fsanitize.c
@@ -984,10 +984,20 @@
 // CHECK-UBSAN-MINIMAL: "-fsanitize={{((signed-integer-overflow|integer-divide-by-zero|shift-base|shift-exponent|unreachable|return|vla-bound|alignment|null|pointer-overflow|float-cast-overflow|array-bounds|enum|bool|builtin|returns-nonnull-attribute|nonnull-attribute|function),?){18}"}}
 // CHECK-UBSAN-MINIMAL: "-fsanitize-minimal-runtime"

-// RUN: %clang --target=x86_64-linux-gnu -fsanitize=undefined -fsanitize-minimal-runtime -fsanitize-handler-preserve-all-regs %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-UBSAN-MINIMAL-PRESERVE
-// CHECK-UBSAN-MINIMAL-PRESERVE: "-fsanitize={{((signed-integer-overflow|integer-divide-by-zero|shift-base|shift-exponent|unreachable|return|vla-bound|alignment|null|pointer-overflow|float-cast-overflow|array-bounds|enum|bool|builtin|returns-nonnull-attribute|nonnull-attribute|function),?){18}"}}
-// CHECK-UBSAN-MINIMAL-PRESERVE: "-fsanitize-minimal-runtime"
-// CHECK-UBSAN-MINIMAL-PRESERVE: "-fsanitize-handler-preserve-all-regs
+// RUN: %clang --target=x86_64-linux-gnu -fsanitize=undefined -fsanitize-minimal-runtime -fsanitize-handler-preserve-all-regs %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-UBSAN-MINIMAL-PRESERVE-X86-64
+// CHECK-UBSAN-MINIMAL-PRESERVE-X86-64: "-fsanitize={{((signed-integer-overflow|integer-divide-by-zero|shift-base|shift-exponent|unreachable|return|vla-bound|alignment|null|pointer-overflow|float-cast-overflow|array-bounds|enum|bool|builtin|returns-nonnull-attribute|nonnull-attribute|function),?){18}"}}
+// CHECK-UBSAN-MINIMAL-PRESERVE-X86-64: "-fsanitize-minimal-runtime"
+// CHECK-UBSAN-MINIMAL-PRESERVE-X86-64: "-fsanitize-handler-preserve-all-regs
+
+// RUN: %clang --target=aarch64-linux-gnu -fsanitize=undefined -fsanitize-minimal-runtime -fsanitize-handler-preserve-all-regs %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-UBSAN-MINIMAL-PRESERVE-AARCH64
+// CHECK-UBSAN-MINIMAL-PRESERVE-AARCH64: "-fsanitize={{((signed-integer-overflow|integer-divide-by-zero|shift-base|shift-exponent|unreachable|return|vla-bound|alignment|null|pointer-overflow|float-cast-overflow|array-bounds|enum|bool|builtin|returns-nonnull-attribute|nonnull-attribute|function),?){18}"}}
+// CHECK-UBSAN-MINIMAL-PRESERVE-AARCH64: "-fsanitize-minimal-runtime"
+// CHECK-UBSAN-MINIMAL-PRESERVE-AARCH64: "-fsanitize-handler-preserve-all-regs
+
+// RUN: %clang --target=i386-linux-gnu -fsanitize=undefined -fsanitize-minimal-runtime -fsanitize-handler-preserve-all-regs %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-UBSAN-MINIMAL-PRESERVE-I386
+// CHECK-UBSAN-MINIMAL-PRESERVE-I386: "-fsanitize={{((signed-integer-overflow|integer-divide-by-zero|shift-base|shift-exponent|unreachable|return|vla-bound|alignment|null|pointer-overflow|float-cast-overflow|array-bounds|enum|bool|builtin|returns-nonnull-attribute|nonnull-attribute|function),?){18}"}}
+// CHECK-UBSAN-MINIMAL-PRESERVE-I386: "-fsanitize-minimal-runtime"
+// CHECK-UBSAN-MINIMAL-PRESERVE-I386-NOT: "-fsanitize-handler-preserve-all-regs

 // RUN: %clang --target=x86_64-linux-gnu -fsanitize=integer -fsanitize-trap=integer %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-INTSAN-TRAP
 // CHECK-INTSAN-TRAP: "-fsanitize-trap=integer-divide-by-zero,shift-base,shift-exponent,signed-integer-overflow,unsigned-integer-overflow,unsigned-shift-base,implicit-unsigned-integer-truncation,implicit-signed-integer-truncation,implicit-integer-sign-change"
--- a/compiler-rt/test/ubsan_minimal/TestCases/override-callback.c
+++ b/compiler-rt/test/ubsan_minimal/TestCases/override-callback.c
@@ -1,6 +1,7 @@
-// RUN: %clang_min_runtime -fsanitize=implicit-integer-sign-change                                        %s -o %t &&             %run %t 2>&1 | FileCheck %s
-// RUN: %clang_min_runtime -fsanitize=implicit-integer-sign-change -fno-sanitize-recover=all              %s -o %t && not --crash %run %t 2>&1 | FileCheck %s
-// RUN: %clang_min_runtime -fsanitize=implicit-integer-sign-change -fno-sanitize-recover=all -DOVERRIDE=1 %s -o %t && not --crash %run %t 2>&1 | FileCheck %s --check-prefixes=FATAL
+// RUN: %clang_min_runtime -fsanitize=implicit-integer-sign-change                                                  %s -o %t &&             %run %t 2>&1 | FileCheck %s
+// RUN: %clang_min_runtime -fsanitize=implicit-integer-sign-change -fsanitize-handler-preserve-all-regs -DPRESERVE  %s -o %t &&             %run %t 2>&1 | FileCheck %s --check-prefixes=PRESERVE
+// RUN: %clang_min_runtime -fsanitize=implicit-integer-sign-change -fno-sanitize-recover=all                        %s -o %t && not --crash %run %t 2>&1 | FileCheck %s
+// RUN: %clang_min_runtime -fsanitize=implicit-integer-sign-change -fno-sanitize-recover=all -DOVERRIDE=1           %s -o %t && not --crash %run %t 2>&1 | FileCheck %s --check-prefixes=FATAL

 #include <stdint.h>
 #include <stdio.h>
@@ -9,9 +10,22 @@
 static int Result;

 void __ubsan_report_error(const char *kind, uintptr_t caller) {
+// -fsanitize-handler-preserve-all-regs is ignored on other architectures.
+// Prented we called to other handler on those.
+#if defined(PRESERVE) && !defined(__aarch64__) && !defined(__x86_64__)
+  fprintf(stderr, "CUSTOM_CALLBACK_PRESERVE: %s\n", kind);
+#else
  fprintf(stderr, "CUSTOM_CALLBACK: %s\n", kind);
+#endif
 }

+#if defined(__aarch64__) || defined(__x86_64__)
+[[clang::preserve_all]] void __ubsan_report_error_preserve(const char *kind,
+                                                           uintptr_t caller) {
+  fprintf(stderr, "CUSTOM_CALLBACK_PRESERVE: %s\n", kind);
+}
+#endif
+
 #if OVERRIDE
 void __ubsan_report_error_fatal(const char *kind, uintptr_t caller) {
  fprintf(stderr, "FATAL_CALLBACK: %s\n", kind);
@@ -21,5 +35,6 @@ void __ubsan_report_error_fatal(const char *kind, uintptr_t caller) {
 int main(int argc, const char **argv) {
  int32_t t0 = (~((uint32_t)0));
  // CHECK: CUSTOM_CALLBACK: implicit-conversion
+  // PRESERVE: CUSTOM_CALLBACK_PRESERVE: implicit-conversion
  // FATAL: FATAL_CALLBACK: implicit-conversion
 }
--- a/llvm/include/llvm/Transforms/Instrumentation/BoundsChecking.h
+++ b/llvm/include/llvm/Transforms/Instrumentation/BoundsChecking.h
@@ -11,6 +11,7 @@

 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/Compiler.h"
+#include "llvm/TargetParser/Triple.h"
 #include <optional>

 namespace llvm {
@@ -23,10 +24,12 @@ class BoundsCheckingPass : public PassInfoMixin<BoundsCheckingPass> {
 public:
  struct Options {
    struct Runtime {
-      Runtime(bool MinRuntime, bool MayReturn)
-          : MinRuntime(MinRuntime), MayReturn(MayReturn) {}
+      Runtime(bool MinRuntime, bool MayReturn, bool HandlerPreserveAllRegs)
+          : MinRuntime(MinRuntime), MayReturn(MayReturn),
+            HandlerPreserveAllRegs(HandlerPreserveAllRegs) {}
      bool MinRuntime;
      bool MayReturn;
+      bool HandlerPreserveAllRegs;
    };
    std::optional<Runtime> Rt; // Trap if empty.
    bool Merge = false;
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1590,24 +1590,31 @@ parseBoundsCheckingOptions(StringRef Params) {
      Options.Rt = {
          /*MinRuntime=*/false,
          /*MayReturn=*/true,
+          /*HandlerPreserveAllRegs=*/false,
      };
    } else if (ParamName == "rt-abort") {
      Options.Rt = {
          /*MinRuntime=*/false,
          /*MayReturn=*/false,
+          /*HandlerPreserveAllRegs=*/false,
      };
    } else if (ParamName == "min-rt") {
      Options.Rt = {
          /*MinRuntime=*/true,
          /*MayReturn=*/true,
+          /*HandlerPreserveAllRegs=*/false,
      };
    } else if (ParamName == "min-rt-abort") {
      Options.Rt = {
          /*MinRuntime=*/true,
          /*MayReturn=*/false,
+          /*HandlerPreserveAllRegs=*/false,
      };
    } else if (ParamName == "merge") {
      Options.Merge = true;
+    } else if (ParamName == "handler-preserve-all-regs") {
+      if (Options.Rt)
+        Options.Rt->HandlerPreserveAllRegs = true;
    } else {
      StringRef ParamEQ;
      StringRef Val;
--- a/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
+++ b/llvm/lib/Transforms/Instrumentation/BoundsChecking.cpp
@@ -178,6 +178,8 @@ getRuntimeCallName(const BoundsCheckingPass::Options::Runtime &Opts) {
    Name += "_minimal";
  if (!Opts.MayReturn)
    Name += "_abort";
+  else if (Opts.HandlerPreserveAllRegs)
+    Name += "_preserve";
  return Name;
 }

@@ -267,7 +269,10 @@ static bool addBoundsChecking(Function &F, TargetLibraryInfo &TLI,
      TrapCall->setDoesNotReturn();
      IRB.CreateUnreachable();
    }
-
+    // The preserve-all logic is somewhat duplicated in CGExpr.cpp for
+    // local-bounds. Make sure to change that too.
+    if (Opts.Rt && Opts.Rt->HandlerPreserveAllRegs && MayReturn)
+      TrapCall->setCallingConv(CallingConv::PreserveAll);
    if (!MayReturn && SingleTrapBB && !DebugTrapBB)
      ReuseTrapBB = TrapBB;

--- a/llvm/test/Instrumentation/BoundsChecking/runtimes.ll
+++ b/llvm/test/Instrumentation/BoundsChecking/runtimes.ll
@@ -8,6 +8,9 @@
 ; RUN: opt < %s -passes='bounds-checking<rt-abort>'     -S | FileCheck %s --check-prefixes=RTABORT-NOMERGE
 ; RUN: opt < %s -passes='bounds-checking<min-rt>'       -S | FileCheck %s --check-prefixes=MINRT-NOMERGE
 ; RUN: opt < %s -passes='bounds-checking<min-rt-abort>' -S | FileCheck %s --check-prefixes=MINRTABORT-NOMERGE
+
+; RUN: opt < %s -passes='bounds-checking<min-rt;handler-preserve-all-regs>'       -S | FileCheck %s --check-prefixes=MINRT-PRESERVE-NOMERGE
+; RUN: opt < %s -passes='bounds-checking<min-rt-abort;handler-preserve-all-regs>' -S | FileCheck %s --check-prefixes=MINRTABORT-NOMERGE
 ;
 ; RUN: opt < %s -passes='bounds-checking<trap;guard=3>'   -S | FileCheck %s --check-prefixes=TR-GUARD-COMMON,TR-GUARD-THREE
 ; RUN: opt < %s -passes='bounds-checking<trap;guard=13>'   -S | FileCheck %s --check-prefixes=TR-GUARD-COMMON,TR-GUARD-THIRTEEN
@@ -95,6 +98,22 @@ define void @f1(i64 %x) nounwind {
 ; RTABORT-NOMERGE-NEXT:    call void @__ubsan_handle_local_out_of_bounds_abort() #[[ATTR2:[0-9]+]], !nosanitize [[META0]]
 ; RTABORT-NOMERGE-NEXT:    unreachable, !nosanitize [[META0]]
 ;
+; MINRT-PRESERVE-NOMERGE-LABEL: define void @f1(
+; MINRT-PRESERVE-NOMERGE-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
+; MINRT-PRESERVE-NOMERGE-NEXT:    [[TMP1:%.*]] = mul i64 16, [[X]]
+; MINRT-PRESERVE-NOMERGE-NEXT:    [[TMP2:%.*]] = alloca i128, i64 [[X]], align 8
+; MINRT-PRESERVE-NOMERGE-NEXT:    [[TMP3:%.*]] = sub i64 [[TMP1]], 0, !nosanitize [[META0:![0-9]+]]
+; MINRT-PRESERVE-NOMERGE-NEXT:    [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 16, !nosanitize [[META0]]
+; MINRT-PRESERVE-NOMERGE-NEXT:    [[TMP5:%.*]] = or i1 false, [[TMP4]], !nosanitize [[META0]]
+; MINRT-PRESERVE-NOMERGE-NEXT:    [[TMP6:%.*]] = or i1 false, [[TMP5]], !nosanitize [[META0]]
+; MINRT-PRESERVE-NOMERGE-NEXT:    br i1 [[TMP6]], label %[[TRAP:.*]], label %[[BB7:.*]]
+; MINRT-PRESERVE-NOMERGE:       [[BB7]]:
+; MINRT-PRESERVE-NOMERGE-NEXT:    [[TMP8:%.*]] = load i128, ptr [[TMP2]], align 4
+; MINRT-PRESERVE-NOMERGE-NEXT:    ret void
+; MINRT-PRESERVE-NOMERGE:       [[TRAP]]:
+; MINRT-PRESERVE-NOMERGE-NEXT:    call preserve_allcc void @__ubsan_handle_local_out_of_bounds_minimal_preserve() #[[ATTR1:[0-9]+]], !nosanitize [[META0]]
+; MINRT-PRESERVE-NOMERGE-NEXT:    br label %[[BB7]], !nosanitize [[META0]]
+;
 ; MINRT-NOMERGE-LABEL: define void @f1(
 ; MINRT-NOMERGE-SAME: i64 [[X:%.*]]) #[[ATTR0:[0-9]+]] {
 ; MINRT-NOMERGE-NEXT:    [[TMP1:%.*]] = mul i64 16, [[X]]