cfi: add diag mode support

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-fulldeps/session-diagnostic/diagnostic-derive-doc-comment-field.rs

@@ -2,6 +2,7 @@
 // Tests that a doc comment will not preclude a field from being considered a diagnostic argument
 //@ normalize-stderr: "the following other types implement trait `IntoDiagArg`:(?:.*\n){0,9}\s+and \d+ others" -> "normalized in stderr"
 //@ normalize-stderr: "(COMPILER_DIR/.*\.rs):[0-9]+:[0-9]+" -> "$1:LL:CC"
+//@ normalize-stderr: "rustc_errors::Diag::<'a, G>::arg" -> "Diag::<'a, G>::arg"
 
 // The proc_macro2 crate handles spans differently when on beta/stable release rather than nightly,
 // changing the output of this test. Since Subdiagnostic is strictly internal to the compiler

tests/ui-fulldeps/session-diagnostic/diagnostic-derive-doc-comment-field.stderr

@@ -1,5 +1,5 @@
 error[E0277]: the trait bound `NotIntoDiagArg: IntoDiagArg` is not satisfied
-  --> $DIR/diagnostic-derive-doc-comment-field.rs:34:10
+  --> $DIR/diagnostic-derive-doc-comment-field.rs:35:10
    |
 LL | #[derive(Diagnostic)]
    |          ---------- required by a bound introduced by this call
@@ -8,7 +8,7 @@ LL |     arg: NotIntoDiagArg,
    |          ^^^^^^^^^^^^^^ unsatisfied trait bound
    |
 help: the nightly-only, unstable trait `IntoDiagArg` is not implemented for `NotIntoDiagArg`
-  --> $DIR/diagnostic-derive-doc-comment-field.rs:26:1
+  --> $DIR/diagnostic-derive-doc-comment-field.rs:27:1
    |
 LL | struct NotIntoDiagArg;
    | ^^^^^^^^^^^^^^^^^^^^^
@@ -22,7 +22,7 @@ note: required by a bound in `Diag::<'a, G>::arg`
    = note: this error originates in the macro `with_fn` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error[E0277]: the trait bound `NotIntoDiagArg: IntoDiagArg` is not satisfied
-  --> $DIR/diagnostic-derive-doc-comment-field.rs:44:10
+  --> $DIR/diagnostic-derive-doc-comment-field.rs:45:10
    |
 LL | #[derive(Subdiagnostic)]
    |          ------------- required by a bound introduced by this call
@@ -31,7 +31,7 @@ LL |     arg: NotIntoDiagArg,
    |          ^^^^^^^^^^^^^^ unsatisfied trait bound
    |
 help: the nightly-only, unstable trait `IntoDiagArg` is not implemented for `NotIntoDiagArg`
-  --> $DIR/diagnostic-derive-doc-comment-field.rs:26:1
+  --> $DIR/diagnostic-derive-doc-comment-field.rs:27:1
    |
 LL | struct NotIntoDiagArg;
    | ^^^^^^^^^^^^^^^^^^^^^