Pointer-Authentication Based Forward Edge Protection

Co-authored-by: Schrodinger ZHU Yifan <i@zhuyi.fan>

Author: Pengxiang-Huang

CMakeLists.txt

@@ -374,6 +374,44 @@ if(SNMALLOC_COMPILER_SUPPORT_MCX16)
   target_compile_options(snmalloc INTERFACE $<$<COMPILE_LANGUAGE:CXX>:-mcx16>)
 endif()
 
+check_cxx_compiler_flag(
+  "-Werror -march=armv8.3-a+pauth -fptrauth-intrinsics"
+  SNMALLOC_COMPILER_SUPPORT_MARCH_V83A_PAUTH)
+if(SNMALLOC_COMPILER_SUPPORT_MARCH_V83A_PAUTH)
+  set(SNMALLOC_SAVED_REQUIRED_FLAGS "${CMAKE_REQUIRED_FLAGS}")
+  set(CMAKE_REQUIRED_FLAGS "${CMAKE_REQUIRED_FLAGS} -march=armv8.3-a+pauth -fptrauth-intrinsics")
+  check_cxx_source_compiles("
+#if !defined(__aarch64__) && !defined(_M_ARM64) && !defined(_M_ARM64EC)
+#  error PAC instructions require AArch64
+#endif
+#if !__has_include(<ptrauth.h>)
+#  error Missing ptrauth.h
+#endif
+#include <ptrauth.h>
+#ifndef __ARM_FEATURE_PAUTH
+#  error Compiler target does not expose pointer authentication
+#endif
+static inline void* sign_and_auth(void* ptr, void* storage)
+{
+  auto discriminator = ptrauth_blend_discriminator(storage, 0x5678U);
+  auto signed_ptr = ptrauth_sign_unauthenticated(
+    ptr, ptrauth_key_process_dependent_data, discriminator);
+  return ptrauth_auth_data(
+    signed_ptr, ptrauth_key_process_dependent_data, discriminator);
+}
+int main()
+{
+  return sign_and_auth(nullptr, reinterpret_cast<void*>(0x1234UL)) != nullptr;
+}
+" SNMALLOC_COMPILER_SUPPORT_PACA_PACG)
+  set(CMAKE_REQUIRED_FLAGS "${SNMALLOC_SAVED_REQUIRED_FLAGS}")
+  unset(SNMALLOC_SAVED_REQUIRED_FLAGS)
+endif()
+if(SNMALLOC_COMPILER_SUPPORT_PACA_PACG)
+  target_compile_options(snmalloc INTERFACE
+    $<$<COMPILE_LANGUAGE:CXX>:-march=armv8.3-a+pauth -fptrauth-intrinsics>)
+endif()
+
 if (NOT SNMALLOC_HEADER_ONLY_LIBRARY AND SNMALLOC_IPO)
   check_ipo_supported(RESULT HAS_IPO)
   if (HAS_IPO)
@@ -400,6 +438,7 @@ add_as_define(SNMALLOC_PLATFORM_HAS_GETENTROPY)
 add_as_define(SNMALLOC_PTHREAD_ATFORK_WORKS)
 add_as_define(SNMALLOC_HAS_LINUX_RANDOM_H)
 add_as_define(SNMALLOC_HAS_LINUX_FUTEX_H)
+add_as_define(SNMALLOC_COMPILER_SUPPORT_PACA_PACG)
 if (SNMALLOC_NO_REALLOCARRAY)
   add_as_define(SNMALLOC_NO_REALLOCARRAY)
 endif()

src/snmalloc/aal/aal.h

@@ -180,6 +180,36 @@ namespace snmalloc
 #endif
       }
     }
+
+    static SNMALLOC_FAST_PATH uintptr_t pointer_auth_sign_data(
+      uintptr_t value, address_t storage_addr, uintptr_t tweak) noexcept
+    {
+      if constexpr ((Arch::aal_features & PtrAuthentication) != 0)
+      {
+        return Arch::pointer_auth_sign_data(
+          value, static_cast<uintptr_t>(storage_addr), tweak);
+      }
+      else
+      {
+        UNUSED(storage_addr, tweak);
+        return value;
+      }
+    }
+
+    static SNMALLOC_FAST_PATH uintptr_t pointer_auth_auth_data(
+      uintptr_t value, address_t storage_addr, uintptr_t tweak) noexcept
+    {
+      if constexpr ((Arch::aal_features & PtrAuthentication) != 0)
+      {
+        return Arch::pointer_auth_auth_data(
+          value, static_cast<uintptr_t>(storage_addr), tweak);
+      }
+      else
+      {
+        UNUSED(storage_addr, tweak);
+        return value;
+      }
+    }
   };
 
   template<class Arch>

src/snmalloc/aal/aal_arm.h

@@ -13,6 +13,12 @@
 #endif
 
 #include <stddef.h>
+#include <stdint.h>
+
+#if defined(SNMALLOC_COMPILER_SUPPORT_PACA_PACG) && \
+  defined(__ARM_FEATURE_PAUTH) && __has_include(<ptrauth.h>)
+#  include <ptrauth.h>
+#endif
 
 namespace snmalloc
 {
@@ -28,6 +34,9 @@ namespace snmalloc
     static constexpr uint64_t aal_features = IntegerPointers
 #if defined(SNMALLOC_VA_BITS_32) || !defined(__APPLE__)
       | NoCpuCycleCounters
+#endif
+#if defined(SNMALLOC_COMPILER_SUPPORT_PACA_PACG) && defined(__ARM_FEATURE_PAUTH)
+      | PtrAuthentication
 #endif
       ;
 
@@ -69,6 +78,35 @@ namespace snmalloc
       return t;
     }
 #endif
+
+#if defined(SNMALLOC_COMPILER_SUPPORT_PACA_PACG) && \
+  defined(__ARM_FEATURE_PAUTH) && __has_include(<ptrauth.h>)
+    static SNMALLOC_FAST_PATH uintptr_t pointer_auth_sign_data(
+      uintptr_t value, uintptr_t storage_addr, uintptr_t tweak) noexcept
+    {
+      return unsafe_to_uintptr<void>(ptrauth_sign_unauthenticated(
+        unsafe_from_uintptr<void>(value),
+        ptrauth_key_process_dependent_data,
+        /*
+         * only the lower 16 bits of tweak is used
+         */
+        ptrauth_blend_discriminator(
+          unsafe_from_uintptr<void>(storage_addr), tweak)));
+    }
+
+    static SNMALLOC_FAST_PATH uintptr_t pointer_auth_auth_data(
+      uintptr_t value, uintptr_t storage_addr, uintptr_t tweak) noexcept
+    {
+      return unsafe_to_uintptr<void>(ptrauth_auth_data(
+        unsafe_from_uintptr<void>(value),
+        ptrauth_key_process_dependent_data,
+        /*
+         * only the lower 16 bits of tweak is used
+         */
+        ptrauth_blend_discriminator(
+          unsafe_from_uintptr<void>(storage_addr), tweak)));
+    }
+#endif
   };
 
   using AAL_Arch = AAL_arm;

src/snmalloc/aal/aal_consts.h

@@ -24,6 +24,12 @@ namespace snmalloc
      * internal high-privilege pointers for recycling memory on free().
      */
     StrictProvenance = (1 << 2),
+    /**
+     * This architecture can authenticate stored internal pointers, allowing
+     * allocator metadata edges to be signed before storage and authenticated on
+     * reload.
+     */
+    PtrAuthentication = (1 << 3),
   };
 
   enum AalName : int

src/snmalloc/mem/freelist.h

@@ -290,22 +290,31 @@ namespace snmalloc
         const FreeListKey& key,
         address_t key_tweak)
       {
-        // Note we can consider other encoding schemes here.
-        //   * XORing curr and next.  This doesn't require any key material
-        //   * XORing (curr * key). This makes it harder to guess the underlying
-        //     key, as each location effectively has its own key.
-        // Curr is not used in the current encoding scheme.
-        UNUSED(curr);
-
         if constexpr (
+          mitigations(freelist_forward_edge) && aal_supports<PtrAuthentication>)
+        {
+          return unsafe_from_uintptr<Object::T<BQueue>>(
+            Aal::pointer_auth_sign_data(
+              unsafe_to_uintptr<Object::T<BQueue>>(next),
+              curr,
+              key.key_next ^ key_tweak));
+        }
+        else if constexpr (
           mitigations(freelist_forward_edge) && !aal_supports<StrictProvenance>)
         {
+          // Note we can consider other encoding schemes here.
+          //   * XORing curr and next.  This doesn't require any key material
+          //   * XORing (curr * key). This makes it harder to guess the
+          //     underlying key, as each location effectively has its own key.
+          // Curr is not used in the current encoding scheme.
+          UNUSED(curr);
           return unsafe_from_uintptr<Object::T<BQueue>>(
             unsafe_to_uintptr<Object::T<BQueue>>(next) ^ key.key_next ^
             key_tweak);
         }
         else
         {
+          UNUSED(curr);
           UNUSED(key);
           UNUSED(key_tweak);
           return next;
@@ -364,8 +373,20 @@ namespace snmalloc
         const FreeListKey& key,
         address_t key_tweak)
       {
-        return BHeadPtr<BView, BQueue>::unsafe_from(
-          code_next(curr, next.unsafe_ptr(), key, key_tweak));
+        if constexpr (
+          mitigations(freelist_forward_edge) && aal_supports<PtrAuthentication>)
+        {
+          return BHeadPtr<BView, BQueue>::unsafe_from(
+            unsafe_from_uintptr<Object::T<BQueue>>(Aal::pointer_auth_auth_data(
+              unsafe_to_uintptr<Object::T<BQueue>>(next.unsafe_ptr()),
+              curr,
+              key.key_next ^ key_tweak)));
+        }
+        else
+        {
+          return BHeadPtr<BView, BQueue>::unsafe_from(
+            code_next(curr, next.unsafe_ptr(), key, key_tweak));
+        }
       }
 
       template<