From 233d94e72ff6853b03836b38b093a3d383f131a4 Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Mon, 8 Apr 2024 20:42:18 +0000
Subject: [PATCH] KCFI: Use legal charset in shim encoding

To separate `ReifyReason::FnPtr` from `ReifyReason::VTable`, we
hyphenated the shims. Hyphens are not actually legal, but underscores
are, so use those instead.
---
 compiler/rustc_symbol_mangling/src/v0.rs |  4 ++--
 tests/ui/sanitizer/kcfi-mangling.rs      | 30 ++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100644 tests/ui/sanitizer/kcfi-mangling.rs

diff --git a/compiler/rustc_symbol_mangling/src/v0.rs b/compiler/rustc_symbol_mangling/src/v0.rs
index 8cb5370bb4a..50dad7d6742 100644
--- a/compiler/rustc_symbol_mangling/src/v0.rs
+++ b/compiler/rustc_symbol_mangling/src/v0.rs
@@ -45,8 +45,8 @@ pub(super) fn mangle<'tcx>(
         ty::InstanceDef::ThreadLocalShim(_) => Some("tls"),
         ty::InstanceDef::VTableShim(_) => Some("vtable"),
         ty::InstanceDef::ReifyShim(_, None) => Some("reify"),
-        ty::InstanceDef::ReifyShim(_, Some(ReifyReason::FnPtr)) => Some("reify-fnptr"),
-        ty::InstanceDef::ReifyShim(_, Some(ReifyReason::Vtable)) => Some("reify-vtable"),
+        ty::InstanceDef::ReifyShim(_, Some(ReifyReason::FnPtr)) => Some("reify_fnptr"),
+        ty::InstanceDef::ReifyShim(_, Some(ReifyReason::Vtable)) => Some("reify_vtable"),
 
         ty::InstanceDef::ConstructCoroutineInClosureShim { .. }
         | ty::InstanceDef::CoroutineKindShim { .. } => Some("fn_once"),
diff --git a/tests/ui/sanitizer/kcfi-mangling.rs b/tests/ui/sanitizer/kcfi-mangling.rs
new file mode 100644
index 00000000000..fde7b5451b6
--- /dev/null
+++ b/tests/ui/sanitizer/kcfi-mangling.rs
@@ -0,0 +1,30 @@
+// Check KCFI extra mangling works correctly on v0
+
+//@ needs-sanitizer-kcfi
+//@ no-prefer-dynamic
+//@ compile-flags: -C panic=abort -Zsanitizer=kcfi -C symbol-mangling-version=v0
+//@ build-pass
+
+trait Foo {
+    fn foo(&self);
+}
+
+struct Bar;
+impl Foo for Bar {
+    fn foo(&self) {}
+}
+
+struct Baz;
+impl Foo for Baz {
+    #[track_caller]
+    fn foo(&self) {}
+}
+
+fn main() {
+    // Produces `ReifyShim(_, ReifyReason::FnPtr)`
+    let f: fn(&Bar) = Bar::foo;
+    f(&Bar);
+    // Produces `ReifyShim(_, ReifyReason::Vtable)`
+    let v: &dyn Foo = &Baz as _;
+    v.foo();
+}