From 73b583cd41df2448d97c34475ad3e696882f6945 Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Sun, 3 Mar 2024 18:38:32 +0000
Subject: [PATCH] CFI: Strip auto traits from Virtual receivers

As the instance being called is behind a vtable, it cannot depend on
auto traits on the receiver (unless the principal trait requires them,
in which case the additional constraint is not needed).

Removing this causes the type signature of the `Virtual` instance to
match the type signature of the `CfiShim`-wrapped entry in the vtable.
---
 compiler/rustc_codegen_ssa/src/mir/block.rs   |  4 +--
 compiler/rustc_middle/src/query/mod.rs        |  7 +++++
 compiler/rustc_ty_utils/src/instance.rs       | 26 ++++++++++++++++---
 .../ui/sanitizer/cfi-marker-trait-objects.rs  | 17 ++++++++++++
 4 files changed, 49 insertions(+), 5 deletions(-)
 create mode 100644 tests/ui/sanitizer/cfi-marker-trait-objects.rs

diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
index 9c7aadb81f828..4271f510d6148 100644
--- a/compiler/rustc_codegen_ssa/src/mir/block.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -495,7 +495,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     //
                     let virtual_drop = Instance {
                         def: ty::InstanceDef::Virtual(drop_fn.def_id(), 0),
-                        args: drop_fn.args,
+                        args: bx.tcx().strip_receiver_auto(drop_fn.args),
                     };
                     debug!("ty = {:?}", ty);
                     debug!("drop_fn = {:?}", drop_fn);
@@ -535,7 +535,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     // SO THEN WE CAN USE THE ABOVE CODE.
                     let virtual_drop = Instance {
                         def: ty::InstanceDef::Virtual(drop_fn.def_id(), 0),
-                        args: drop_fn.args,
+                        args: bx.tcx().strip_receiver_auto(drop_fn.args),
                     };
                     debug!("ty = {:?}", ty);
                     debug!("drop_fn = {:?}", drop_fn);
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index eb647d1ca8333..9bc3e1ed63b7a 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -2256,6 +2256,13 @@ rustc_queries! {
     query trait_object_ty(trait_ref: ty::PolyTraitRef<'tcx>) -> Ty<'tcx> {
         desc { "Compute the trait object type for calling a method on a trait" }
     }
+
+    /// Strip auto traits off the first parameter in the parametr list. Intended for use when
+    /// constructing `InstanceDef::Virtual`, as auto traits won't be part of the vtable's `Self`
+    /// types.
+    query strip_receiver_auto(args: ty::GenericArgsRef<'tcx>) -> ty::GenericArgsRef<'tcx> {
+        desc { "Strip auto traits off the first type arg" }
+    }
 }
 
 rustc_query_append! { define_callbacks! }
diff --git a/compiler/rustc_ty_utils/src/instance.rs b/compiler/rustc_ty_utils/src/instance.rs
index 3e3bccce47fe4..eb2273b58452d 100644
--- a/compiler/rustc_ty_utils/src/instance.rs
+++ b/compiler/rustc_ty_utils/src/instance.rs
@@ -5,7 +5,7 @@ use rustc_infer::infer::TyCtxtInferExt;
 use rustc_middle::query::Providers;
 use rustc_middle::traits::{BuiltinImplSource, CodegenObligationError};
 use rustc_middle::ty::GenericArgsRef;
-use rustc_middle::ty::{self, Instance, TyCtxt, TypeVisitableExt};
+use rustc_middle::ty::{self, Instance, Ty, TyCtxt, TypeVisitableExt};
 use rustc_span::sym;
 use rustc_trait_selection::traits;
 use traits::{translate_args, Reveal};
@@ -199,7 +199,7 @@ fn resolve_associated_item<'tcx>(
             traits::get_vtable_index_of_object_method(tcx, *vtable_base, trait_item_id).map(
                 |index| Instance {
                     def: ty::InstanceDef::Virtual(trait_item_id, index),
-                    args: rcvr_args,
+                    args: tcx.strip_receiver_auto(rcvr_args),
                 },
             )
         }
@@ -339,6 +339,26 @@ fn resolve_associated_item<'tcx>(
     })
 }
 
+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))
+}
+
 pub(crate) fn provide(providers: &mut Providers) {
-    *providers = Providers { resolve_instance, ..*providers };
+    *providers = Providers { resolve_instance, strip_receiver_auto, ..*providers };
 }
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();
+}