1 // Copyright 2023 The Chromium Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #ifndef BASE_MEMORY_SAFETY_CHECKS_H_
6 #define BASE_MEMORY_SAFETY_CHECKS_H_
7
8 #include <new>
9 #include <type_traits>
10
11 #include "base/allocator/partition_allocator/src/partition_alloc/partition_alloc_buildflags.h"
12 #include "base/allocator/partition_allocator/src/partition_alloc/partition_alloc_constants.h"
13 #include "base/compiler_specific.h"
14 #include "base/dcheck_is_on.h"
15
16 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
17 #include "base/allocator/partition_allocator/src/partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc.h"
18 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
19
20 // This header defines following macros:
21 // - ADVANCED_MEMORY_SAFETY_CHECKS()
22 // They can be used to specify a class/struct that is targeted to perform
23 // additional CHECKS across variety of memory safety mechanisms such as
24 // PartitionAllocator.
25 //
26 // class Foo {
27 // ADVANCED_MEMORY_SAFETY_CHECKS();
28 // }
29 //
30 // Checks here are disabled by default because of their performance cost.
31 // Currently, the macro is managed by the memory safety team internally and
32 // you should not add / remove it manually.
33
34 // We cannot hide things behind anonymous namespace because they are referenced
35 // via macro, which can be defined anywhere.
36 // To avoid tainting ::base namespace, define things inside this namespace.
37 namespace base::internal {
38
39 enum class MemorySafetyCheck : uint32_t {
40 kForcePartitionAlloc = (1u << 0),
41 // Enables |FreeFlags::kSchedulerLoopQuarantine|.
42 // Requires PA-E.
43 kSchedulerLoopQuarantine = (1u << 1),
44
45 // Enables |FreeFlags::kZap|.
46 // Requires PA-E.
47 kZapOnFree = (1u << 2),
48 };
49
50 constexpr MemorySafetyCheck operator|(MemorySafetyCheck a,
51 MemorySafetyCheck b) {
52 return static_cast<MemorySafetyCheck>(static_cast<uint32_t>(a) |
53 static_cast<uint32_t>(b));
54 }
55
56 constexpr MemorySafetyCheck operator&(MemorySafetyCheck a,
57 MemorySafetyCheck b) {
58 return static_cast<MemorySafetyCheck>(static_cast<uint32_t>(a) &
59 static_cast<uint32_t>(b));
60 }
61
62 // Set of checks for ADVANCED_MEMORY_SAFETY_CHECKS() annotated objects.
63 constexpr auto kAdvancedMemorySafetyChecks =
64 MemorySafetyCheck::kForcePartitionAlloc |
65 MemorySafetyCheck::kSchedulerLoopQuarantine | MemorySafetyCheck::kZapOnFree;
66
67 // Define type traits to determine type |T|'s memory safety check status.
68 namespace {
69
70 // Allocator type traits.
ShouldUsePartitionAlloc(MemorySafetyCheck checks)71 constexpr bool ShouldUsePartitionAlloc(MemorySafetyCheck checks) {
72 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
73 return static_cast<bool>(checks &
74 (MemorySafetyCheck::kForcePartitionAlloc |
75 MemorySafetyCheck::kSchedulerLoopQuarantine |
76 MemorySafetyCheck::kZapOnFree));
77 #else
78 return false;
79 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
80 }
81
82 // Returns |partition_alloc::AllocFlags| corresponding to |checks|.
GetAllocFlags(MemorySafetyCheck checks)83 constexpr partition_alloc::AllocFlags GetAllocFlags(MemorySafetyCheck checks) {
84 return partition_alloc::AllocFlags::kReturnNull |
85 partition_alloc::AllocFlags::kNoHooks;
86 }
87
88 // Returns |partition_alloc::FreeFlags| corresponding to |checks|.
GetFreeFlags(MemorySafetyCheck checks)89 constexpr partition_alloc::FreeFlags GetFreeFlags(MemorySafetyCheck checks) {
90 auto flags = partition_alloc::FreeFlags::kNone;
91 if (static_cast<bool>(checks & MemorySafetyCheck::kSchedulerLoopQuarantine)) {
92 flags |= partition_alloc::FreeFlags::kSchedulerLoopQuarantine;
93 }
94 if (static_cast<bool>(checks & MemorySafetyCheck::kZapOnFree)) {
95 flags |= partition_alloc::FreeFlags::kZap;
96 }
97 return flags;
98 }
99
100 } // namespace
101
102 // Public utility type traits.
103 template <typename T>
104 inline constexpr MemorySafetyCheck get_memory_safety_checks = [] {
105 if constexpr (requires { T::kMemorySafetyChecks; }) {
106 return T::kMemorySafetyChecks;
107 } else {
108 return static_cast<MemorySafetyCheck>(0);
109 }
110 }();
111
112 template <typename T, MemorySafetyCheck c>
113 inline constexpr bool is_memory_safety_checked =
114 (get_memory_safety_checks<T> & c) == c;
115
116 // Allocator functions.
117 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
118 ALWAYS_INLINE partition_alloc::PartitionRoot*
GetPartitionRootForMemorySafetyCheckedAllocation()119 GetPartitionRootForMemorySafetyCheckedAllocation() {
120 return allocator_shim::internal::PartitionAllocMalloc::Allocator();
121 }
122
123 ALWAYS_INLINE partition_alloc::PartitionRoot*
GetAlignedPartitionRootForMemorySafetyCheckedAllocation()124 GetAlignedPartitionRootForMemorySafetyCheckedAllocation() {
125 return allocator_shim::internal::PartitionAllocMalloc::AlignedAllocator();
126 }
127 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
128
129 template <MemorySafetyCheck checks>
HandleMemorySafetyCheckedOperatorNew(std::size_t count)130 NOINLINE void* HandleMemorySafetyCheckedOperatorNew(std::size_t count) {
131 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
132 if constexpr (ShouldUsePartitionAlloc(checks)) {
133 return GetPartitionRootForMemorySafetyCheckedAllocation()
134 ->AllocInline<GetAllocFlags(checks)>(count);
135 } else
136 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
137 {
138 return ::operator new(count);
139 }
140 }
141
142 template <MemorySafetyCheck checks>
HandleMemorySafetyCheckedOperatorNew(std::size_t count,std::align_val_t alignment)143 NOINLINE void* HandleMemorySafetyCheckedOperatorNew(
144 std::size_t count,
145 std::align_val_t alignment) {
146 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
147 if constexpr (ShouldUsePartitionAlloc(checks)) {
148 return GetAlignedPartitionRootForMemorySafetyCheckedAllocation()
149 ->AlignedAlloc<GetAllocFlags(checks)>(static_cast<size_t>(alignment),
150 count);
151 } else
152 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
153 {
154 return ::operator new(count, alignment);
155 }
156 }
157
158 template <MemorySafetyCheck checks>
HandleMemorySafetyCheckedOperatorDelete(void * ptr)159 NOINLINE void HandleMemorySafetyCheckedOperatorDelete(void* ptr) {
160 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
161 if constexpr (ShouldUsePartitionAlloc(checks)) {
162 GetPartitionRootForMemorySafetyCheckedAllocation()
163 ->Free<GetFreeFlags(checks)>(ptr);
164 } else
165 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
166 {
167 ::operator delete(ptr);
168 }
169 }
170
171 template <MemorySafetyCheck checks>
HandleMemorySafetyCheckedOperatorDelete(void * ptr,std::align_val_t alignment)172 NOINLINE void HandleMemorySafetyCheckedOperatorDelete(
173 void* ptr,
174 std::align_val_t alignment) {
175 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
176 if constexpr (ShouldUsePartitionAlloc(checks)) {
177 GetAlignedPartitionRootForMemorySafetyCheckedAllocation()
178 ->Free<GetFreeFlags(checks)>(ptr);
179 } else
180 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
181 {
182 ::operator delete(ptr, alignment);
183 }
184 }
185
186 } // namespace base::internal
187
188 // Macros to annotate class/struct's default memory safety check.
189 // ADVANCED_MEMORY_SAFETY_CHECKS(): Enable Check |kAdvancedChecks| for this
190 // object.
191 //
192 // Note that if you use this macro at the top of struct declaration, the
193 // declaration context would be left as |private|. Please switch it back to
194 // |public| manually if needed.
195 //
196 // struct ObjectWithAdvancedChecks {
197 // ADVANCED_MEMORY_SAFETY_CHECKS();
198 // public:
199 // int public_field;
200 // };
201 #define ADVANCED_MEMORY_SAFETY_CHECKS_INTERNAL(SPECIFIER) \
202 public: \
203 static constexpr auto kMemorySafetyChecks = \
204 base::internal::kAdvancedMemorySafetyChecks; \
205 SPECIFIER static void* operator new(std::size_t count) { \
206 return base::internal::HandleMemorySafetyCheckedOperatorNew< \
207 kMemorySafetyChecks>(count); \
208 } \
209 SPECIFIER static void* operator new(std::size_t count, \
210 std::align_val_t alignment) { \
211 return base::internal::HandleMemorySafetyCheckedOperatorNew< \
212 kMemorySafetyChecks>(count, alignment); \
213 } \
214 /* Though we do not hook placement new, we need to define this */ \
215 /* explicitly to allow it. */ \
216 ALWAYS_INLINE static void* operator new(std::size_t, void* ptr) { \
217 return ptr; \
218 } \
219 SPECIFIER static void operator delete(void* ptr) noexcept { \
220 base::internal::HandleMemorySafetyCheckedOperatorDelete< \
221 kMemorySafetyChecks>(ptr); \
222 } \
223 SPECIFIER static void operator delete(void* ptr, \
224 std::align_val_t alignment) noexcept { \
225 base::internal::HandleMemorySafetyCheckedOperatorDelete< \
226 kMemorySafetyChecks>(ptr, alignment); \
227 } \
228 \
229 private: \
230 static_assert(true) /* semicolon here */
231
232 #if DCHECK_IS_ON()
233 // Specify NOINLINE to display the operator on a stack trace.
234 #define ADVANCED_MEMORY_SAFETY_CHECKS() \
235 ADVANCED_MEMORY_SAFETY_CHECKS_INTERNAL(NOINLINE NOT_TAIL_CALLED)
236 #else
237 #define ADVANCED_MEMORY_SAFETY_CHECKS() \
238 ADVANCED_MEMORY_SAFETY_CHECKS_INTERNAL(ALWAYS_INLINE)
239 #endif // DCHECK_IS_ON()
240
241 #endif // BASE_MEMORY_SAFETY_CHECKS_H_
242