From 8e7065d5e18067c7cd1facceda98fe72d71acee7 Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Sun, 24 Mar 2024 18:56:15 +0000
Subject: [PATCH] CFI: Encode Virtual calls as calls through the defining trait

For example, if `trait Foo: Bar`, and we try to call a method from `Bar`
on `dyn Foo`, encode the callsite as passing a `dyn Bar`, not a `dyn
Foo`.
---
 .../src/typeid/typeid_itanium_cxx_abi.rs      | 25 ++++---
 tests/ui/sanitizer/cfi-supertraits.rs         | 68 +++++++++++++++++++
 2 files changed, 84 insertions(+), 9 deletions(-)
 create mode 100644 tests/ui/sanitizer/cfi-supertraits.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 ca571f7ef7d0c..7fc9eae6b5441 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
@@ -1114,8 +1114,10 @@ pub fn typeid_for_instance<'tcx>(
     mut instance: Instance<'tcx>,
     options: TypeIdOptions,
 ) -> String {
-    if matches!(instance.def, ty::InstanceDef::Virtual(..)) {
-        instance.args = strip_receiver_auto(tcx, instance.args)
+    if let ty::InstanceDef::Virtual(def_id, _) = instance.def {
+        let upcast_ty = upcast(tcx, def_id, instance.args);
+        let stripped_ty = strip_receiver_auto(tcx, upcast_ty);
+        instance.args = tcx.mk_args_trait(stripped_ty, instance.args.into_iter().skip(1));
     }
 
     if let Some(impl_id) = tcx.impl_of_method(instance.def_id())
@@ -1146,11 +1148,7 @@ 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);
+fn strip_receiver_auto<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Ty<'tcx> {
     let ty::Dynamic(preds, lifetime, kind) = ty.kind() else {
         bug!("Tried to strip auto traits from non-dynamic type {ty}");
     };
@@ -1162,8 +1160,7 @@ fn strip_receiver_auto<'tcx>(
         } 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))
+    Ty::new_dynamic(tcx, filtered_preds, *lifetime, *kind)
 }
 
 fn trait_object_ty<'tcx>(tcx: TyCtxt<'tcx>, poly_trait_ref: ty::PolyTraitRef<'tcx>) -> Ty<'tcx> {
@@ -1198,3 +1195,13 @@ fn trait_object_ty<'tcx>(tcx: TyCtxt<'tcx>, poly_trait_ref: ty::PolyTraitRef<'tc
     );
     Ty::new_dynamic(tcx, preds, tcx.lifetimes.re_erased, ty::Dyn)
 }
+
+fn upcast<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, args: GenericArgsRef<'tcx>) -> Ty<'tcx> {
+    let self_ty = args.type_at(0);
+    let Some(trait_id) = tcx.trait_of_item(def_id) else {
+        // If it's a virtual call to `drop_in_place`, it won't be on a trait.
+        return self_ty;
+    };
+    let trait_ref = ty::TraitRef::from_method(tcx, trait_id, args);
+    trait_object_ty(tcx, ty::Binder::dummy(trait_ref))
+}
diff --git a/tests/ui/sanitizer/cfi-supertraits.rs b/tests/ui/sanitizer/cfi-supertraits.rs
new file mode 100644
index 0000000000000..0a2d7b060d30a
--- /dev/null
+++ b/tests/ui/sanitizer/cfi-supertraits.rs
@@ -0,0 +1,68 @@
+#![feature(trait_upcasting)]
+// Check that super-traits are callable.
+
+//@ needs-sanitizer-cfi
+// FIXME(#122848) Remove only-linux once OSX CFI binaries work
+//@ only-linux
+//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
+//@ compile-flags: -C target-feature=-crt-static -C codegen-units=1 -C opt-level=0
+//@ run-pass
+
+trait Parent1 {
+    type P1;
+    fn p1(&self) -> Self::P1;
+}
+
+trait Parent2 {
+    type P2;
+    fn p2(&self) -> Self::P2;
+}
+
+trait Child : Parent1 + Parent2 {
+    type C;
+    fn c(&self) -> Self::C;
+}
+
+struct Foo;
+
+impl Parent1 for Foo {
+    type P1 = u16;
+    fn p1(&self) -> Self::P1 {
+        println!("p1");
+        1
+    }
+}
+
+impl Parent2 for Foo {
+    type P2 = u32;
+    fn p2(&self) -> Self::P2 {
+        println!("p2");
+        2
+    }
+}
+
+impl Child for Foo {
+    type C = u8;
+    fn c(&self) -> Self::C {
+        println!("c");
+        0
+    }
+}
+
+fn main() {
+    // Child can access its own methods and super methods.
+    let x = &Foo as &dyn Child<C=u8,P1=u16,P2=u32>;
+    x.c();
+    x.p1();
+    x.p2();
+    // Parents can be created and access their methods.
+    let y = &Foo as &dyn Parent1<P1=u16>;
+    y.p1();
+    let z = &Foo as &dyn Parent2<P2=u32>;
+    z.p2();
+    // Trait upcasting works
+    let x1 = x as &dyn Parent1<P1=u16>;
+    x1.p1();
+    let x2 = x as &dyn Parent2<P2=u32>;
+    x2.p2();
+}