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,64 @@ 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();
+
+            let is_diag = self.tcx.sess.opts.unstable_opts.sanitizer_cfi_diag.unwrap_or(false);
+            let is_recover =
+                self.tcx.sess.opts.unstable_opts.sanitizer_cfi_recover.unwrap_or(false);
+
+            if is_diag || is_recover {
+                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(
+                    if is_recover {
+                        "__ubsan_handle_cfi_check_fail"
+                    } else {
+                        "__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,
+                );
+                if is_recover {
+                    self.br(bb_pass);
+                } else {
+                    self.unreachable();
+                }
+            } else {
+                self.abort();
+                self.unreachable();
+            }
 
             self.switch_to_block(bb_pass);
             if let Some(dbg_loc) = dbg_loc {

compiler/rustc_codegen_ssa/src/back/link.rs

@@ -1457,6 +1457,12 @@ 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.sanitizer_cfi_diag.unwrap_or(false)
+            || sess.opts.unstable_opts.sanitizer_cfi_recover.unwrap_or(false))
+    {
+        link_sanitizer_runtime(sess, flavor, linker, "ubsan");
+    }
 }
 
 fn link_sanitizer_runtime(

compiler/rustc_codegen_ssa/src/back/write.rs

@@ -84,6 +84,8 @@ pub struct ModuleConfig {
     pub instrument_coverage: bool,
 
     pub sanitizer: SanitizerSet,
+    pub sanitizer_cfi_diag: Option<bool>,
+    pub sanitizer_cfi_recover: Option<bool>,
     pub sanitizer_recover: SanitizerSet,
     pub sanitizer_dataflow_abilist: Vec<String>,
     pub sanitizer_memory_track_origins: usize,
@@ -183,6 +185,8 @@ impl ModuleConfig {
             instrument_coverage: if_regular!(sess.instrument_coverage(), false),
 
             sanitizer: if_regular!(sess.sanitizers(), SanitizerSet::empty()),
+            sanitizer_cfi_diag: if_regular!(sess.opts.unstable_opts.sanitizer_cfi_diag, None),
+            sanitizer_cfi_recover: if_regular!(sess.opts.unstable_opts.sanitizer_cfi_recover, None),
             sanitizer_dataflow_abilist: if_regular!(
                 sess.opts.unstable_opts.sanitizer_dataflow_abilist.clone(),
                 Vec::new()

compiler/rustc_session/src/config.rs

@@ -3062,14 +3062,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, 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;
 

compiler/rustc_session/src/options.rs

@@ -2626,6 +2626,10 @@ written to standard error output)"),
         "enable generalizing pointer types (default: no)"),
     sanitizer_cfi_normalize_integers: Option<bool> = (None, parse_opt_bool, [TRACKED] { TARGET_MODIFIER: SanitizerCfiNormalizeIntegers },
         "enable normalizing integer types (default: no)"),
+    sanitizer_cfi_diag: Option<bool> = (None, parse_opt_bool, [TRACKED],
+        "enable CFI diagnostics (default: no)"),
+    sanitizer_cfi_recover: Option<bool> = (None, parse_opt_bool, [TRACKED],
+        "enable CFI recovery (default: no)"),
     sanitizer_dataflow_abilist: Vec<String> = (Vec::new(), parse_comma_list, [TRACKED],
         "additional ABI list files that control how shadow parameters are passed (comma separated)"),
     sanitizer_kcfi_arity: Option<bool> = (None, parse_opt_bool, [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 -Zsanitizer-cfi-diag=true -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/codegen-llvm/sanitizer/cfi/emit-type-checks-recover-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 -Zsanitizer-cfi-recover=true -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(
+    // CHECK-NEXT:  br label %type_test.pass
+    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;
    | ^^^^^^^^^^^^^^^^^^^^^

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 sanitizer-cfi-diag=true
+//@ [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));
+}