KCFI: Add KCFI arity indicator support

Adds KCFI arity indicator support to the Rust compiler (see rust-lang/rust#138311,
https://github.com/llvm/llvm-project/pull/121070, and
https://lore.kernel.org/lkml/CANiq72=3ghFxy8E=AU9p+0imFxKr5iU3sd0hVUXed5BA+KjdNQ@mail.gmail.com/).

3 files changed, 15 insertions, 0 deletions

diff --git a/compiler/rustc_session/src/errors.rs b/compiler/rustc_session/src/errors.rs
index 71d8dbe44fe..3979aea70d3 100644
--- a/compiler/rustc_session/src/errors.rs
+++ b/compiler/rustc_session/src/errors.rs
@@ -148,6 +148,10 @@ pub(crate) struct SanitizerCfiGeneralizePointersRequiresCfi;
 pub(crate) struct SanitizerCfiNormalizeIntegersRequiresCfi;
 
 #[derive(Diagnostic)]
+#[diag(session_sanitizer_kcfi_arity_requires_kcfi)]
+pub(crate) struct SanitizerKcfiArityRequiresKcfi;
+
+#[derive(Diagnostic)]
 #[diag(session_sanitizer_kcfi_requires_panic_abort)]
 pub(crate) struct SanitizerKcfiRequiresPanicAbort;
 
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
index 4f544d2c16b..720bc3f61e4 100644
--- a/compiler/rustc_session/src/options.rs
+++ b/compiler/rustc_session/src/options.rs
@@ -2443,6 +2443,8 @@ written to standard error output)"),
         "enable normalizing integer types (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],
+        "enable KCFI arity indicator (default: no)"),
     sanitizer_memory_track_origins: usize = (0, parse_sanitizer_memory_track_origins, [TRACKED],
         "enable origins tracking in MemorySanitizer"),
     sanitizer_recover: SanitizerSet = (SanitizerSet::empty(), parse_sanitizers, [TRACKED],
diff --git a/compiler/rustc_session/src/session.rs b/compiler/rustc_session/src/session.rs
index a87b1961a99..2c6b6eb619b 100644
--- a/compiler/rustc_session/src/session.rs
+++ b/compiler/rustc_session/src/session.rs
@@ -381,6 +381,10 @@ impl Session {
         self.opts.unstable_opts.sanitizer_cfi_normalize_integers == Some(true)
     }
 
+    pub fn is_sanitizer_kcfi_arity_enabled(&self) -> bool {
+        self.opts.unstable_opts.sanitizer_kcfi_arity == Some(true)
+    }
+
     pub fn is_sanitizer_kcfi_enabled(&self) -> bool {
         self.opts.unstable_opts.sanitizer.contains(SanitizerSet::KCFI)
     }
@@ -1211,6 +1215,11 @@ fn validate_commandline_args_with_session_available(sess: &Session) {
         }
     }
 
+    // KCFI arity indicator requires KCFI.
+    if sess.is_sanitizer_kcfi_arity_enabled() && !sess.is_sanitizer_kcfi_enabled() {
+        sess.dcx().emit_err(errors::SanitizerKcfiArityRequiresKcfi);
+    }
+
     // LLVM CFI pointer generalization requires CFI or KCFI.
     if sess.is_sanitizer_cfi_generalize_pointers_enabled() {
         if !(sess.is_sanitizer_cfi_enabled() || sess.is_sanitizer_kcfi_enabled()) {