From 54b15b0c36d4638167732a0542ed0e34ecf17d7c Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Fri, 15 Mar 2024 18:44:40 +0000
Subject: [PATCH] CFI: Strip auto traits off Self for virtual calls

Additional trait bounds beyond the principal trait and its implications
are not possible in the vtable. This means that if a receiver is
`&dyn Foo + Send`, the function will only be expecting `&dyn Foo`.

This strips those auto traits off before CFI encoding.
---
 .../src/typeid/typeid_itanium_cxx_abi.rs      | 24 +++++++++++++++++++
 .../ui/sanitizer/cfi-marker-trait-objects.rs  | 17 +++++++++++++
 2 files changed, 41 insertions(+)
 create mode 100644 tests/ui/sanitizer/cfi-marker-trait-objects.rs

diff --git a/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs b/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs
index 7be1376bd6362..a69c80236e1e5 100644
--- a/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs
+++ b/compiler/rustc_symbol_mangling/src/typeid/typeid_itanium_cxx_abi.rs
@@ -1153,6 +1153,10 @@ pub fn typeid_for_instance<'tcx>(
         instance.args = tcx.mk_args_trait(invoke_ty, instance.args.into_iter().skip(1));
     }
 
+    if matches!(instance.def, InstanceDef::Virtual(..)) {
+        instance.args = strip_receiver_auto(tcx, instance.args)
+    }
+
     let fn_abi = tcx
         .fn_abi_of_instance(tcx.param_env(instance.def_id()).and((instance, ty::List::empty())))
         .unwrap_or_else(|instance| {
@@ -1161,3 +1165,23 @@ pub fn typeid_for_instance<'tcx>(
 
     typeid_for_fnabi(tcx, fn_abi, options)
 }
+
+fn strip_receiver_auto<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    args: ty::GenericArgsRef<'tcx>,
+) -> ty::GenericArgsRef<'tcx> {
+    let ty = args.type_at(0);
+    let ty::Dynamic(preds, lifetime, kind) = ty.kind() else {
+        bug!("Tried to strip auto traits from non-dynamic type {ty}");
+    };
+    let filtered_preds =
+        if preds.principal().is_some() {
+            tcx.mk_poly_existential_predicates_from_iter(preds.into_iter().filter(|pred| {
+                !matches!(pred.skip_binder(), ty::ExistentialPredicate::AutoTrait(..))
+            }))
+        } else {
+            ty::List::empty()
+        };
+    let new_rcvr = Ty::new_dynamic(tcx, filtered_preds, *lifetime, *kind);
+    tcx.mk_args_trait(new_rcvr, args.into_iter().skip(1))
+}
diff --git a/tests/ui/sanitizer/cfi-marker-trait-objects.rs b/tests/ui/sanitizer/cfi-marker-trait-objects.rs
new file mode 100644
index 0000000000000..cf9be8328076c
--- /dev/null
+++ b/tests/ui/sanitizer/cfi-marker-trait-objects.rs
@@ -0,0 +1,17 @@
+// Test that we can promote closures / fns to trait objects, and call them despite a marker trait.
+
+//@ needs-sanitizer-cfi
+//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
+//@ compile-flags: -C codegen-units=1 -C opt-level=0
+//@ run-pass
+
+
+fn foo() {}
+
+static FOO: &'static (dyn Fn() + Sync) = &foo;
+static BAR: &(dyn Fn() -> i32 + Sync) = &|| 3;
+
+fn main() {
+    FOO();
+    BAR();
+}