cfi: add diag mode support

Currently a Rust CFI failure only inserts a ud2. However, for clang
we have the option for a helpful diagnostic message that explains the
violation and is especially helpful for fixing it. This message works
through hooking the UBSan runtime and calling into it with the
necessary information for a helpful error message. In clang, this is
enabled via -fno-sanitize-trap=cfi.

This change adds the same behavior to rustc's CFI. Instead of a
no-sanitize-trap flag, we added -Z cfi-mode={diag|trap}, with trap
as the default. The diag mode will print the following error message
for a violation:

```
tests/ui/sanitizer/cfi/fn-ptr-type-mismatch.rs:1:1: runtime error: control flow integrity check for type fn(i32, i32) -> i32 failed during indirect function call
fn_ptr_type_mismatch.ecd806f409c5c1fc-cgu.0: note: fn_ptr_type_mismatch::add_one defined here
SUMMARY: UndefinedBehaviorSanitizer: undefined-behavior tests/ui/sanitizer/cfi/fn-ptr-type-mismatch.rs:1:1
```

Author: 1c3t3a

compiler/rustc_codegen_llvm/src/builder.rs

@@ -43,6 +43,7 @@ use crate::type_of::LayoutLlvmExt;
 pub(crate) struct GenericBuilder<'a, 'll, CX: Borrow<SCx<'ll>>> {
     pub llbuilder: &'ll mut llvm::Builder<'ll>,
     pub cx: &'a GenericCx<'ll, CX>,
+    pub span: rustc_span::Span,
 }
 
 pub(crate) type SBuilder<'a, 'll> = GenericBuilder<'a, 'll, SCx<'ll>>;
@@ -93,7 +94,7 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
     fn with_cx(scx: &'a GenericCx<'ll, CX>) -> Self {
         // Create a fresh builder from the simple context.
         let llbuilder = unsafe { llvm::LLVMCreateBuilderInContext(scx.deref().borrow().llcx) };
-        GenericBuilder { llbuilder, cx: scx }
+        GenericBuilder { llbuilder, cx: scx, span: rustc_span::DUMMY_SP }
     }
 
     pub(crate) fn append_block(
@@ -303,7 +304,9 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         unsafe { llvm::LLVMGetInsertBlock(self.llbuilder) }
     }
 
-    fn set_span(&mut self, _span: Span) {}
+    fn set_span(&mut self, span: rustc_span::Span) {
+        self.span = span;
+    }
 
     fn append_block(cx: &'a CodegenCx<'ll, 'tcx>, llfn: &'ll Value, name: &str) -> &'ll BasicBlock {
         unsafe {
@@ -1491,6 +1494,51 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
     pub(crate) fn llfn(&self) -> &'ll Value {
         unsafe { llvm::LLVMGetBasicBlockParent(self.llbb()) }
     }
+
+    fn generate_ubsan_cfi_diag_data(
+        &mut self,
+        span: rustc_span::Span,
+        expected_ty: String,
+        check_kind: u8,
+    ) -> &'ll Value {
+        let cx = self.cx();
+        let tcx = cx.tcx;
+
+        let loc = tcx.sess.source_map().lookup_char_pos(span.lo());
+
+        let filename_str = format!("{}\0", loc.file.name.prefer_local_unconditionally());
+        let filename_val = cx.const_bytes(filename_str.as_bytes());
+        let filename_ptr = cx.static_addr_of_impl(filename_val, Align::ONE, None);
+
+        // SourceLocation UBSan struct: { const char *filename, uint32_t line, uint32_t column }
+        let source_location = cx.const_struct(
+            &[
+                filename_ptr,
+                cx.const_u32(loc.line as u32),
+                // UBSan columns are 1-based
+                cx.const_u32(loc.col.0 as u32 + 1),
+            ],
+            false, // packed = false
+        );
+
+        let ty_name = format!("{}\0", expected_ty);
+        let ty_name_val = cx.const_bytes(ty_name.as_bytes());
+
+        // TypeDescriptor UBSan struct: { uint16_t TypeKind, uint16_t TypeInfo, const char *TypeName }
+        let type_descriptor =
+            cx.const_struct(&[cx.const_i16(0xffffu16 as i16), cx.const_i16(0), ty_name_val], false);
+
+        let type_descriptor_ptr =
+            cx.static_addr_of_impl(type_descriptor, Align::from_bytes(2).unwrap(), None);
+
+        // CFICheckFailData UBSan struct: { uint8_t CheckKind, SourceLocation Loc, TypeDescriptor *Type }
+        let cfi_check_fail_data = cx
+            .const_struct(&[cx.const_u8(check_kind), source_location, type_descriptor_ptr], false);
+        let align = tcx.data_layout.aggregate_align;
+
+        // Returns the final opaque pointer to the struct to be passed to __ubsan_handle_cfi_check_fail
+        cx.static_addr_of_mut(cfi_check_fail_data, align, Some("__ubsan_cfi_check_fail_data"))
+    }
 }
 
 impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
@@ -1905,8 +1953,55 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
             if let Some(dbg_loc) = dbg_loc {
                 self.set_dbg_loc(dbg_loc);
             }
-            self.abort();
-            self.unreachable();
+
+            match self.tcx.sess.opts.unstable_opts.cfi_mode {
+                rustc_session::config::CfiMode::Trap => {
+                    self.abort();
+                    self.unreachable();
+                }
+                rustc_session::config::CfiMode::Diag => {
+                    let fty = self.cx.type_func(
+                        &[self.cx.type_ptr(), self.cx.type_isize(), self.cx.type_isize()],
+                        self.cx.type_void(),
+                    );
+                    let ubsan_handler = self.declare_cfn(
+                        "__ubsan_handle_cfi_check_fail_abort",
+                        llvm::UnnamedAddr::Global,
+                        fty,
+                    );
+
+                    let mut expected_ty = String::from("fn(");
+                    for (i, arg) in fn_abi.args.iter().enumerate() {
+                        if i > 0 {
+                            expected_ty.push_str(", ");
+                        }
+                        use std::fmt::Write;
+                        write!(&mut expected_ty, "{}", arg.layout.ty).unwrap();
+                    }
+                    expected_ty.push(')');
+                    if !fn_abi.ret.layout.ty.is_unit() {
+                        use std::fmt::Write;
+                        write!(&mut expected_ty, " -> {}", fn_abi.ret.layout.ty).unwrap();
+                    }
+
+                    // 4 for cfi-icall (indirect call)
+                    let check_kind = 4;
+                    let diag_data =
+                        self.generate_ubsan_cfi_diag_data(self.span, expected_ty, check_kind);
+
+                    let function_address = self.ptrtoint(llfn, self.cx.type_isize());
+                    self.call(
+                        fty,
+                        None,
+                        None,
+                        ubsan_handler,
+                        &[diag_data, function_address, self.const_usize(0)],
+                        None,
+                        None,
+                    );
+                    self.unreachable();
+                }
+            }
 
             self.switch_to_block(bb_pass);
             if let Some(dbg_loc) = dbg_loc {

compiler/rustc_codegen_ssa/src/back/link.rs

@@ -35,7 +35,7 @@ use rustc_middle::middle::debugger_visualizer::DebuggerVisualizerFile;
 use rustc_middle::middle::dependency_format::Linkage;
 use rustc_middle::middle::exported_symbols::SymbolExportKind;
 use rustc_session::config::{
-    self, CFGuard, CrateType, DebugInfo, LinkerFeaturesCli, OutFileName, OutputFilenames,
+    self, CFGuard, CfiMode, CrateType, DebugInfo, LinkerFeaturesCli, OutFileName, OutputFilenames,
     OutputType, PrintKind, SplitDwarfKind, Strip,
 };
 use rustc_session::lint::builtin::LINKER_MESSAGES;
@@ -1457,6 +1457,9 @@ fn add_sanitizer_libraries(
     if sanitizer.contains(SanitizerSet::REALTIME) {
         link_sanitizer_runtime(sess, flavor, linker, "rtsan");
     }
+    if sanitizer.contains(SanitizerSet::CFI) && sess.opts.unstable_opts.cfi_mode == CfiMode::Diag {
+        link_sanitizer_runtime(sess, flavor, linker, "ubsan");
+    }
 }
 
 fn link_sanitizer_runtime(

compiler/rustc_session/src/config.rs

@@ -88,6 +88,16 @@ pub enum CFProtection {
     Full,
 }
 
+/// The different settings that the `-Z cfi-mode` flag can have.
+#[derive(Clone, Copy, PartialEq, Hash, Debug)]
+pub enum CfiMode {
+    /// Emit a trap instruction on CFI violation (default).
+    Trap,
+
+    /// Emit a diagnostic on CFI violation instead of trapping.
+    Diag,
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Hash, StableHash, Encodable, Decodable)]
 pub enum OptLevel {
     /// `-Copt-level=0`
@@ -3062,14 +3072,14 @@ pub(crate) mod dep_tracking {
     };
 
     use super::{
-        AnnotateMoves, AutoDiff, BranchProtection, CFGuard, CFProtection, CodegenRetagOptions,
-        CoverageOptions, CrateType, DebugInfo, DebugInfoCompression, ErrorOutputType, FmtDebug,
-        FunctionReturn, InliningThreshold, InstrumentCoverage, InstrumentXRay, LinkerPluginLto,
-        LocationDetail, LtoCli, MirStripDebugInfo, NextSolverConfig, Offload, OptLevel,
-        OutFileName, OutputType, OutputTypes, PatchableFunctionEntry, Polonius, ResolveDocLinks,
-        SourceFileHashAlgorithm, SplitDwarfKind, SwitchWithOptPath, SymbolManglingVersion,
-        WasiExecModel,
+        AnnotateMoves, AutoDiff, BranchProtection, CFGuard, CFProtection, CfiMode, CoverageOptions,
+        CrateType, DebugInfo, DebugInfoCompression, ErrorOutputType, FmtDebug, FunctionReturn,
+        InliningThreshold, InstrumentCoverage, InstrumentXRay, LinkerPluginLto, LocationDetail,
+        LtoCli, MirStripDebugInfo, NextSolverConfig, Offload, OptLevel, OutFileName, OutputType,
+        OutputTypes, PatchableFunctionEntry, Polonius, ResolveDocLinks, SourceFileHashAlgorithm,
+        SplitDwarfKind, SwitchWithOptPath, SymbolManglingVersion, WasiExecModel,
     };
+    use crate::config::CodegenRetagOptions;
     use crate::lint;
     use crate::utils::NativeLib;
 
@@ -3147,6 +3157,7 @@ pub(crate) mod dep_tracking {
         SanitizerSet,
         CFGuard,
         CFProtection,
+        CfiMode,
         TargetTuple,
         Edition,
         LinkerPluginLto,

compiler/rustc_session/src/options.rs

@@ -773,6 +773,7 @@ mod desc {
     pub(crate) const parse_cfguard: &str =
         "either a boolean (`yes`, `no`, `on`, `off`, etc), `checks`, or `nochecks`";
     pub(crate) const parse_cfprotection: &str = "`none`|`no`|`n` (default), `branch`, `return`, or `full`|`yes`|`y` (equivalent to `branch` and `return`)";
+    pub(crate) const parse_cfi_mode: &str = "`trap` (default) or `diag`";
     pub(crate) const parse_debuginfo: &str = "either an integer (0, 1, 2), `none`, `line-directives-only`, `line-tables-only`, `limited`, or `full`";
     pub(crate) const parse_debuginfo_compression: &str = "one of `none`, `zlib`, or `zstd`";
     pub(crate) const parse_mir_strip_debuginfo: &str =
@@ -1307,6 +1308,15 @@ pub mod parse {
         true
     }
 
+    pub(crate) fn parse_cfi_mode(slot: &mut CfiMode, v: Option<&str>) -> bool {
+        *slot = match v {
+            Some("trap") => CfiMode::Trap,
+            Some("diag") => CfiMode::Diag,
+            _ => return false,
+        };
+        true
+    }
+
     pub(crate) fn parse_debuginfo(slot: &mut DebugInfo, v: Option<&str>) -> bool {
         match v {
             Some("0") | Some("none") => *slot = DebugInfo::None,
@@ -2261,6 +2271,8 @@ options! {
         "cache the results of derive proc macro invocations (potentially unsound!) (default: no"),
     cf_protection: CFProtection = (CFProtection::None, parse_cfprotection, [TRACKED],
         "instrument control-flow architecture protection"),
+    cfi_mode: CfiMode = (CfiMode::Trap, parse_cfi_mode, [TRACKED],
+        "set the CFI failure mode: `trap` (default) or `diag`"),
     check_cfg_all_expected: bool = (false, parse_bool, [UNTRACKED],
         "show all expected values in check-cfg diagnostics (default: no)"),
     checksum_hash_algorithm: Option<SourceFileHashAlgorithm> = (None, parse_cargo_src_file_hash, [TRACKED],

src/bootstrap/src/core/build_steps/llvm.rs

@@ -1520,21 +1520,28 @@ fn supported_sanitizers(
     let darwin_libs = |os: &str, components: &[&str]| -> Vec<SanitizerRuntime> {
         components
             .iter()
-            .map(move |c| SanitizerRuntime {
-                cmake_target: format!("clang_rt.{c}_{os}_dynamic"),
-                path: out_dir.join(format!("build/lib/darwin/libclang_rt.{c}_{os}_dynamic.dylib")),
-                name: format!("librustc-{channel}_rt.{c}.dylib"),
+            .map(move |c| {
+                let cmake_c = if *c == "ubsan" { "ubsan_standalone" } else { *c };
+                SanitizerRuntime {
+                    cmake_target: format!("clang_rt.{cmake_c}_{os}_dynamic"),
+                    path: out_dir
+                        .join(format!("build/lib/darwin/libclang_rt.{cmake_c}_{os}_dynamic.dylib")),
+                    name: format!("librustc-{channel}_rt.{c}.dylib"),
+                }
             })
             .collect()
     };
 
     let common_libs = |os: &str, arch: &str, components: &[&str]| -> Vec<SanitizerRuntime> {
         components
             .iter()
-            .map(move |c| SanitizerRuntime {
-                cmake_target: format!("clang_rt.{c}-{arch}"),
-                path: out_dir.join(format!("build/lib/{os}/libclang_rt.{c}-{arch}.a")),
-                name: format!("librustc-{channel}_rt.{c}.a"),
+            .map(move |c| {
+                let cmake_c = if *c == "ubsan" { "ubsan_standalone" } else { *c };
+                SanitizerRuntime {
+                    cmake_target: format!("clang_rt.{cmake_c}-{arch}"),
+                    path: out_dir.join(format!("build/lib/{os}/libclang_rt.{cmake_c}-{arch}.a")),
+                    name: format!("librustc-{channel}_rt.{c}.a"),
+                }
             })
             .collect()
     };
@@ -1545,9 +1552,11 @@ fn supported_sanitizers(
         "aarch64-apple-ios-sim" => darwin_libs("iossim", &["asan", "tsan", "rtsan"]),
         "aarch64-apple-ios-macabi" => darwin_libs("osx", &["asan", "lsan", "tsan"]),
         "aarch64-unknown-fuchsia" => common_libs("fuchsia", "aarch64", &["asan"]),
-        "aarch64-unknown-linux-gnu" => {
-            common_libs("linux", "aarch64", &["asan", "lsan", "msan", "tsan", "hwasan", "rtsan"])
-        }
+        "aarch64-unknown-linux-gnu" => common_libs(
+            "linux",
+            "aarch64",
+            &["asan", "lsan", "msan", "tsan", "hwasan", "rtsan", "ubsan"],
+        ),
         "aarch64-unknown-linux-ohos" => {
             common_libs("linux", "aarch64", &["asan", "lsan", "msan", "tsan", "hwasan"])
         }
@@ -1567,7 +1576,7 @@ fn supported_sanitizers(
         "x86_64-unknown-linux-gnu" => common_libs(
             "linux",
             "x86_64",
-            &["asan", "dfsan", "lsan", "msan", "safestack", "tsan", "rtsan"],
+            &["asan", "dfsan", "lsan", "msan", "safestack", "tsan", "rtsan", "ubsan"],
         ),
         "x86_64-unknown-linux-gnuasan" => common_libs("linux", "x86_64", &["asan"]),
         "x86_64-unknown-linux-gnumsan" => common_libs("linux", "x86_64", &["msan"]),

tests/codegen-llvm/sanitizer/cfi/emit-type-checks-diag-mode.rs

@@ -0,0 +1,20 @@
+// Verifies that pointer type membership tests for indirect calls are emitted.
+//
+//@ needs-sanitizer-cfi
+//@ compile-flags: -Clto -Cno-prepopulate-passes -Ctarget-feature=-crt-static -Zsanitizer=cfi -Zcfi-mode=diag -Copt-level=0 -C unsafe-allow-abi-mismatch=sanitizer
+
+#![crate_type = "lib"]
+
+pub fn foo(f: fn(i32) -> i32, arg: i32) -> i32 {
+    // CHECK-LABEL: define{{.*}}foo{{.*}}!type !{{[0-9]+}} !type !{{[0-9]+}} !type !{{[0-9]+}} !type !{{[0-9]+}}
+    // CHECK:       start:
+    // CHECK:       [[TT:%.+]] = call i1 @llvm.type.test(ptr {{%f|%0}}, metadata !"{{[[:print:]]+}}")
+    // CHECK-NEXT:  br i1 [[TT]], label %type_test.pass, label %type_test.fail
+    // CHECK:       type_test.pass:
+    // CHECK-NEXT:  {{%.+}} = call i32 %f(i32{{.*}} %arg)
+    // CHECK:       type_test.fail:
+    // CHECK-NEXT:  {{%.+}} = ptrtoint ptr {{%f|%0}} to i64
+    // CHECK-NEXT:  call void @__ubsan_handle_cfi_check_fail_abort(
+    // CHECK-NEXT:  unreachable
+    f(arg)
+}

tests/ui/sanitizer/cfi/fn-ptr-type-mismatch.rs

@@ -0,0 +1,41 @@
+// Verifies that calling a function pointer with a mismatched type triggers a
+// CFI violation and causes the process to trap.
+
+//@ revisions: cfi kcfi
+// FIXME(#122848) Remove only-linux once OSX CFI binaries work
+//@ only-linux
+//@ ignore-backends: gcc
+//@ [cfi] needs-sanitizer-cfi
+//@ [cfi] needs-sanitizer-support
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C target-feature=-crt-static
+//@ compile-flags: -C unsafe-allow-abi-mismatch=sanitizer
+//@ [cfi] compile-flags: -C opt-level=0 -C codegen-units=1 -C lto
+//@ [cfi] compile-flags: -C prefer-dynamic=off
+//@ [cfi] compile-flags: -Z sanitizer=cfi
+//@ [cfi] compile-flags: -Z cfi-mode=diag
+//@ [kcfi] compile-flags: -Z sanitizer=kcfi
+//@ [kcfi] compile-flags: -C panic=abort -C prefer-dynamic=off
+//@ run-fail-or-crash
+
+use std::hint::black_box;
+use std::mem;
+
+fn add_one(x: i32) -> i32 {
+    x + 1
+}
+
+// Accept a function pointer as a parameter so that the indirect call cannot
+// be devirtualized by the compiler.
+#[inline(never)]
+fn call_with_mismatch(f: fn(i32) -> i32) {
+    // Transmute fn(i32) -> i32 into fn(i32, i32) -> i32, creating a
+    // function pointer type mismatch that CFI should catch.
+    let g: fn(i32, i32) -> i32 = unsafe { mem::transmute(f) };
+    // This indirect call should fail the CFI type check and trap.
+    let _result = g(1, 2);
+}
+
+fn main() {
+    call_with_mismatch(black_box(add_one));
+}