From 91eadab4a37ac52d3da22a4e9feda88962c2294f Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Mon, 18 Sep 2023 23:54:21 +0000
Subject: [PATCH] cfi: Store type erasure witness for Argument

CFI/KCFI work by enforcing that indirect function calls are always
called at the type they were defined at. The type erasure in Argument
works by casting the reference to the value to be formatted to an Opaque
and also casting the function to format it to take an Opaque reference.

While this is *ABI* safe, which is why we get away with it normally, it
does transform the type that the function is called at. This means that
at the call-site, CFI expects the type of the function to be
`fn(&Opaque, ` even though it is really `fn(&T, ` for some particular
`T`.

This patch avoids this by adding `cast_stub`, a witness to the type
erasure that will cast the `&Opaque` and `fn(&Opaque` back to their
original types before invoking the function.

This change is guarded by the enablement of CFI as it will require an
additional pointer-sized value per `Argument`, and an additional jump
during formatting, and we'd prefer not to pay that if we don't need the
types to be correct at the indirect call invocation.
---
 library/core/src/fmt/rt.rs | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/library/core/src/fmt/rt.rs b/library/core/src/fmt/rt.rs
index 65a4d537cc74d..016de17cb6927 100644
--- a/library/core/src/fmt/rt.rs
+++ b/library/core/src/fmt/rt.rs
@@ -72,6 +72,12 @@ enum ArgumentType<'a> {
         // was derived from a `&'a T`.
         value: NonNull<()>,
         formatter: unsafe fn(NonNull<()>, &mut Formatter<'_>) -> Result,
+        #[cfg(any(sanitize = "cfi", sanitize = "kcfi"))]
+        cast_stub: unsafe fn(
+            unsafe fn(NonNull<()>, &mut Formatter<'_>) -> Result,
+            NonNull<()>,
+            f: &mut Formatter<'_>,
+        ) -> Result,
         _lifetime: PhantomData<&'a ()>,
     },
     Count(usize),
@@ -93,6 +99,21 @@ pub struct Argument<'a> {
     ty: ArgumentType<'a>,
 }
 
+/// This function acts as a witness to an earlier type erasure when constructing an
+/// `ArgumentType`. The type parameter `T` should be instantiated to the erased type.
+/// SAFETY: This function should be called on a value and formatter which underwent a
+/// T->Opaque cast at their definition site.
+#[cfg(any(sanitize = "cfi", sanitize = "kcfi"))]
+unsafe fn cast_stub<T>(
+    formatter: unsafe fn(NonNull<()>, &mut Formatter<'_>) -> Result,
+    value: NonNull<()>,
+    f: &mut Formatter<'_>,
+) -> Result {
+    let value: &T = mem::transmute(value);
+    let formatter: fn(&T, &mut Formatter<'_>) -> Result = mem::transmute(formatter);
+    formatter(value, f)
+}
+
 #[rustc_diagnostic_item = "ArgumentMethods"]
 impl<'a> Argument<'a> {
     #[inline(always)]
@@ -104,6 +125,8 @@ impl<'a> Argument<'a> {
                 value: NonNull::from(x).cast(),
                 // SAFETY: function pointers always have the same layout.
                 formatter: unsafe { mem::transmute(f) },
+                #[cfg(any(sanitize = "cfi", sanitize = "kcfi"))]
+                cast_stub: cast_stub::<T>,
                 _lifetime: PhantomData,
             },
         }
@@ -163,6 +186,15 @@ impl<'a> Argument<'a> {
     #[inline(always)]
     pub(super) unsafe fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         match self.ty {
+            #[cfg(any(sanitize = "cfi", sanitize = "kcfi"))]
+            // SAFETY:
+            // This is the `cast_stub` that was prepared alongside the
+            // formatter and value, so it should have the `T` instantiated to
+            // the erased type that `value` points at.
+            ArgumentType::Placeholder { formatter, value, cast_stub, .. } => unsafe {
+                cast_stub(formatter, value, f)
+            },
+            #[cfg(not(any(sanitize = "cfi", sanitize = "kcfi")))]
             // SAFETY:
             // Because of the invariant that if `formatter` had the type
             // `fn(&T, _) -> _` then `value` has type `&'b T` where `'b` is