From e974570c42f4460b30fce11003b7891435931cbd Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Wed, 27 Mar 2024 15:50:53 +0000
Subject: [PATCH] CFI: Only encode Coroutine Parent Args

Fixes #122705
---
 .../src/typeid/typeid_itanium_cxx_abi.rs      | 16 ++++++++--
 tests/ui/sanitizer/cfi-async-closures.rs      | 30 +++++++++++++++++++
 2 files changed, 43 insertions(+), 3 deletions(-)
 create mode 100644 tests/ui/sanitizer/cfi-async-closures.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 aba3212b4ce..82ceccba66a 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
@@ -641,9 +641,7 @@ fn encode_ty<'tcx>(
         }
 
         // Function types
-        ty::FnDef(def_id, args)
-        | ty::Closure(def_id, args)
-        | ty::CoroutineClosure(def_id, args) => {
+        ty::FnDef(def_id, args) | ty::Closure(def_id, args) => {
             // u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
             // as vendor extended type.
             let mut s = String::new();
@@ -654,6 +652,18 @@ fn encode_ty<'tcx>(
             typeid.push_str(&s);
         }
 
+        ty::CoroutineClosure(def_id, args) => {
+            // u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
+            // as vendor extended type.
+            let mut s = String::new();
+            let name = encode_ty_name(tcx, *def_id);
+            let _ = write!(s, "u{}{}", name.len(), &name);
+            let parent_args = tcx.mk_args(args.as_coroutine_closure().parent_args());
+            s.push_str(&encode_args(tcx, parent_args, dict, options));
+            compress(dict, DictKey::Ty(ty, TyQ::None), &mut s);
+            typeid.push_str(&s);
+        }
+
         ty::Coroutine(def_id, args, ..) => {
             // u<length><name>[I<element-type1..element-typeN>E], where <element-type> is <subst>,
             // as vendor extended type.
diff --git a/tests/ui/sanitizer/cfi-async-closures.rs b/tests/ui/sanitizer/cfi-async-closures.rs
new file mode 100644
index 00000000000..50edeb852bd
--- /dev/null
+++ b/tests/ui/sanitizer/cfi-async-closures.rs
@@ -0,0 +1,30 @@
+// Check various forms of dynamic closure calls
+
+//@ edition: 2021
+//@ revisions: cfi kcfi
+// FIXME(#122848) Remove only-linux once OSX CFI binaries work
+//@ only-linux
+//@ [cfi] needs-sanitizer-cfi
+//@ [kcfi] needs-sanitizer-kcfi
+//@ compile-flags: -C target-feature=-crt-static
+//@ [cfi] compile-flags: -C codegen-units=1 -C lto -C prefer-dynamic=off -C opt-level=0
+//@ [cfi] compile-flags: -Z sanitizer=cfi
+//@ [kcfi] compile-flags: -Z sanitizer=kcfi
+//@ run-pass
+
+#![feature(async_closure)]
+#![feature(async_fn_traits)]
+
+use std::ops::AsyncFn;
+
+#[inline(never)]
+fn identity<T>(x: T) -> T { x }
+
+// We can't actually create a `dyn AsyncFn()`, because it's not object-safe, but we should check
+// that we don't bug out when we encounter one.
+
+fn main() {
+   let f = identity(async || ());
+   let _ = f.async_call(());
+   let _ = f();
+}