[libc++] Properly implement array cookies in the ARM ABI

[ldionne]

When we implemented array cookie support for hardening std::unique_ptr, the implementation was only done for the Itanium ABI. I did not initially realize that ARM was using a different ABI for array cookies, so unique_ptr should not have been hardened on ARM.

However, we were also incorrectly setting the ABI-detection macro: we were pretending to be using a vanilla Itanium ABI when in reality the (similar but different) ARM ABI was in use. As a result, unique_ptr was using the wrong representation for array cookies on ARM, which fortunately only mattered in the case of overaligned types.

This patch fixes that.

rdar://160852193

[llvmbot]

@llvm/pr-subscribers-libcxx

Author: Louis Dionne (ldionne)

Changes

When we implemented array cookie support for hardening std::unique_ptr, the implementation was only done for the Itanium ABI. I did not initially realize that ARM was using a different ABI for array cookies, so unique_ptr should not have been hardened on ARM.

However, we were also incorrectly setting the ABI-detection macro: we were pretending to be using a vanilla Itanium ABI when in reality the (similar but different) ARM ABI was in use. As a result, unique_ptr was using the wrong representation for array cookies on ARM, which fortunately only mattered in the case of overaligned types.

This patch fixes that.

rdar://160852193

---

Full diff: https://github.com/llvm/llvm-project/pull/160182.diff

3 Files Affected:

• (modified) libcxx/include/__configuration/abi.h (+2)
• (modified) libcxx/include/__memory/array_cookie.h (+65-5)
• (modified) libcxx/test/std/utilities/smartptr/unique.ptr/unique.ptr.class/unique.ptr.observers/assert.subscript.pass.cpp (+33)

diff --git a/libcxx/include/__configuration/abi.h b/libcxx/include/__configuration/abi.h
index 2d33b9c03090b..e3c7f830358a8 100644
--- a/libcxx/include/__configuration/abi.h
+++ b/libcxx/include/__configuration/abi.h
@@ -32,6 +32,8 @@
 #else
 #  if defined(_WIN32) && defined(_MSC_VER)
 #    define _LIBCPP_ABI_MICROSOFT
+#  elif defined(__arm__) || defined(__aarch64__)
+#    define _LIBCPP_ABI_ARM
 #  else
 #    define _LIBCPP_ABI_ITANIUM
 #  endif
diff --git a/libcxx/include/__memory/array_cookie.h b/libcxx/include/__memory/array_cookie.h
index 806a9e99ecafe..8cc4c0308f1dc 100644
--- a/libcxx/include/__memory/array_cookie.h
+++ b/libcxx/include/__memory/array_cookie.h
@@ -26,12 +26,12 @@ _LIBCPP_BEGIN_NAMESPACE_STD
 // Trait representing whether a type requires an array cookie at the start of its allocation when
 // allocated as `new T[n]` and deallocated as `delete[] array`.
 //
-// Under the Itanium C++ ABI [1], we know that an array cookie is available unless `T` is trivially
-// destructible and the call to `operator delete[]` is not a sized operator delete. Under ABIs other
-// than the Itanium ABI, we assume there are no array cookies.
+// Under the Itanium C++ ABI [1] and the ARM ABI which derives from it, we know that an array cookie is available
+// unless `T` is trivially destructible and the call to `operator delete[]` is not a sized operator delete. Under
+// other ABIs, we assume there are no array cookies.
 //
 // [1]: https://itanium-cxx-abi.github.io/cxx-abi/abi.html#array-cookies
-#ifdef _LIBCPP_ABI_ITANIUM
+#if defined(_LIBCPP_ABI_ITANIUM) || defined(_LIBCPP_ABI_ARM)
 // TODO: Use a builtin instead
 // TODO: We should factor in the choice of the usual deallocation function in this determination.
 template <class _Tp>
@@ -41,13 +41,73 @@ template <class _Tp>
 struct __has_array_cookie : false_type {};
 #endif
 
+// Return the array cookie located before the given pointer.
+//
+// In the Itanium ABI
+// ------------------
+// The array cookie is stored immediately before the first element of the array. If the preferred alignment
+// of array elements (which is different from the ABI alignment) is more than that of size_t, additional
+// padding bytes exist before the array cookie. Assuming array elements of size and alignment 16 bytes, that
+// gives us the following layout:
+//
+// |ooooooooxxxxxxxxaaaaaaaaaaaaaaaabbbbbbbbbbbbbbbbccccccccccccccccdddddddddddddddd|
+//  ^^^^^^^^        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+//     |    ^^^^^^^^                               |
+//     |       |                              array elements
+//  padding    |
+//        array cookie
+//
+// In practice, it is sufficient to read the bytes immediately before the first array element.
+//
+//
+// In the ARM ABI
+// --------------
+// The array cookie is stored at the very start of the allocation and it has the following form:
+//
+//    struct array_cookie {
+//      std::size_t element_size; // element_size != 0
+//      std::size_t element_count;
+//    };
+//
+// Assuming elements of size and alignment 32 bytes, this gives us the following layout:
+//
+//  |xxxxxxxxXXXXXXXXooooooooooooooooaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaabbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb|
+//   ^^^^^^^^        ^^^^^^^^^^^^^^^^
+//      |    ^^^^^^^^        |       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+// element size  |        padding                                 |
+//         element count                                     array elements
+//
+// We calculate the starting address of the allocation by taking into account the ABI (not the preferred)
+// alignment of the type.
 template <class _Tp>
 // Avoid failures when -fsanitize-address-poison-custom-array-cookie is enabled
 _LIBCPP_HIDE_FROM_ABI _LIBCPP_NO_SANITIZE("address") size_t __get_array_cookie(_Tp const* __ptr) {
   static_assert(
       __has_array_cookie<_Tp>::value, "Trying to access the array cookie of a type that is not guaranteed to have one");
-  size_t const* __cookie = reinterpret_cast<size_t const*>(__ptr) - 1; // TODO: Use a builtin instead
+
+#if defined(_LIBCPP_ABI_ITANIUM)
+
+  size_t const* __cookie = reinterpret_cast<size_t const*>(__ptr) - 1;
   return *__cookie;
+
+#elif defined(_LIBCPP_ABI_ARM)
+
+  struct _ArrayCookie {
+    size_t __element_size;
+    size_t __element_count;
+  };
+
+  size_t __cookie_size_with_padding = // max(sizeof(_ArrayCookie), alignof(T))
+      sizeof(_ArrayCookie) < alignof(_Tp) ? alignof(_Tp) : sizeof(_ArrayCookie);
+  char const* __allocation_start = reinterpret_cast<char const*>(__ptr) - __cookie_size_with_padding;
+  _ArrayCookie const* __cookie   = reinterpret_cast<_ArrayCookie const*>(__allocation_start);
+  return __cookie->__element_count;
+
+#else
+
+  static_assert(sizeof(_Tp) == 0, "This function is not implemented for this ABI");
+
+#endif
 }
 
 _LIBCPP_END_NAMESPACE_STD
diff --git a/libcxx/test/std/utilities/smartptr/unique.ptr/unique.ptr.class/unique.ptr.observers/assert.subscript.pass.cpp b/libcxx/test/std/utilities/smartptr/unique.ptr/unique.ptr.class/unique.ptr.observers/assert.subscript.pass.cpp
index b7cc12350027b..2de523dfb25cb 100644
--- a/libcxx/test/std/utilities/smartptr/unique.ptr/unique.ptr.class/unique.ptr.observers/assert.subscript.pass.cpp
+++ b/libcxx/test/std/utilities/smartptr/unique.ptr/unique.ptr.class/unique.ptr.observers/assert.subscript.pass.cpp
@@ -58,15 +58,18 @@ void test() {
   {
     {
       std::unique_ptr<WithCookie[]> ptr(new WithCookie[5]);
+      assert(&ptr[1] == ptr.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(ptr[6], "unique_ptr<T[]>::operator[](index): index out of range");
     }
     {
       std::unique_ptr<WithCookie[]> ptr = std::make_unique<WithCookie[]>(5);
+      assert(&ptr[1] == ptr.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(ptr[6], "unique_ptr<T[]>::operator[](index): index out of range");
     }
 #if TEST_STD_VER >= 20
     {
       std::unique_ptr<WithCookie[]> ptr = std::make_unique_for_overwrite<WithCookie[]>(5);
+      assert(&ptr[1] == ptr.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(ptr[6] = WithCookie(), "unique_ptr<T[]>::operator[](index): index out of range");
     }
 #endif
@@ -82,11 +85,13 @@ void test() {
   {
     {
       std::unique_ptr<NoCookie[]> ptr = std::make_unique<NoCookie[]>(5);
+      assert(&ptr[1] == ptr.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(ptr[6], "unique_ptr<T[]>::operator[](index): index out of range");
     }
 #  if TEST_STD_VER >= 20
     {
       std::unique_ptr<NoCookie[]> ptr = std::make_unique_for_overwrite<NoCookie[]>(5);
+      assert(&ptr[1] == ptr.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(ptr[6] = NoCookie(), "unique_ptr<T[]>::operator[](index): index out of range");
     }
 #  endif
@@ -101,6 +106,7 @@ void test() {
     {
       std::unique_ptr<T[]> ptr = std::make_unique<T[]>(5);
       std::unique_ptr<T[]> other(std::move(ptr));
+      assert(&other[1] == other.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(other[6], "unique_ptr<T[]>::operator[](index): index out of range");
     }
 
@@ -109,6 +115,7 @@ void test() {
       std::unique_ptr<T[]> ptr = std::make_unique<T[]>(5);
       std::unique_ptr<T[]> other;
       other = std::move(ptr);
+      assert(&other[1] == other.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(other[6], "unique_ptr<T[]>::operator[](index): index out of range");
     }
 
@@ -116,6 +123,7 @@ void test() {
     {
       std::unique_ptr<T[]> ptr = std::make_unique<T[]>(5);
       std::unique_ptr<T[], MyDeleter> other(std::move(ptr));
+      assert(&other[1] == other.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(other[6], "unique_ptr<T[]>::operator[](index): index out of range");
     }
 
@@ -124,6 +132,7 @@ void test() {
       std::unique_ptr<T[]> ptr = std::make_unique<T[]>(5);
       std::unique_ptr<T[], MyDeleter> other;
       other = std::move(ptr);
+      assert(&other[1] == other.get() + 1); // ensure no assertion
       TEST_LIBCPP_ASSERT_FAILURE(other[6], "unique_ptr<T[]>::operator[](index): index out of range");
     }
   });
@@ -144,6 +153,20 @@ struct WithCookie {
   char padding[Size];
 };
 
+template <std::size_t Size>
+struct alignas(128) OveralignedNoCookie {
+  char padding[Size];
+};
+
+template <std::size_t Size>
+struct alignas(128) OveralignedWithCookie {
+  OveralignedWithCookie() = default;
+  OveralignedWithCookie(OveralignedWithCookie const&) {}
+  OveralignedWithCookie& operator=(OveralignedWithCookie const&) { return *this; }
+  ~OveralignedWithCookie() {}
+  char padding[Size];
+};
+
 int main(int, char**) {
   test<WithCookie<1>, NoCookie<1>>();
   test<WithCookie<2>, NoCookie<2>>();
@@ -153,6 +176,16 @@ int main(int, char**) {
   test<WithCookie<16>, NoCookie<16>>();
   test<WithCookie<32>, NoCookie<32>>();
   test<WithCookie<256>, NoCookie<256>>();
+
+  test<OveralignedWithCookie<1>, OveralignedNoCookie<1>>();
+  test<OveralignedWithCookie<2>, OveralignedNoCookie<2>>();
+  test<OveralignedWithCookie<3>, OveralignedNoCookie<3>>();
+  test<OveralignedWithCookie<4>, OveralignedNoCookie<4>>();
+  test<OveralignedWithCookie<8>, OveralignedNoCookie<8>>();
+  test<OveralignedWithCookie<16>, OveralignedNoCookie<16>>();
+  test<OveralignedWithCookie<32>, OveralignedNoCookie<32>>();
+  test<OveralignedWithCookie<256>, OveralignedNoCookie<256>>();
+
   test<std::string, int>();
 
   return 0;

[philnik777]

Couldn't we simplify this to

  struct [[__gnu__::__aligned__(alignof(_Tp)), __gnu__::__may_alias__]] _ArrayCookie {
    size_t __element_size;
    size_t __element_count;
  };
  
  auto __cookie = reinterpret_cast<const _ArrayCookie*>(__ptr) - 1;
  return __cookie->__element_count;

[philnik777]

Is there a reference manual for the ARM ABI somewhere? Could we link to that from here?