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/compiler_specific.h"
12 #include "base/dcheck_is_on.h"
13 #include "partition_alloc/partition_alloc_buildflags.h"
14 #include "partition_alloc/partition_alloc_constants.h"
15
16 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
17 #include "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 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
123
124 template <MemorySafetyCheck checks>
HandleMemorySafetyCheckedOperatorNew(std::size_t count)125 NOINLINE void* HandleMemorySafetyCheckedOperatorNew(std::size_t count) {
126 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
127 if constexpr (ShouldUsePartitionAlloc(checks)) {
128 return GetPartitionRootForMemorySafetyCheckedAllocation()
129 ->AllocInline<GetAllocFlags(checks)>(count);
130 }
131 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
132 return ::operator new(count);
133 }
134
135 template <MemorySafetyCheck checks>
HandleMemorySafetyCheckedOperatorNew(std::size_t count,std::align_val_t alignment)136 NOINLINE void* HandleMemorySafetyCheckedOperatorNew(
137 std::size_t count,
138 std::align_val_t alignment) {
139 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
140 if constexpr (ShouldUsePartitionAlloc(checks)) {
141 return GetPartitionRootForMemorySafetyCheckedAllocation()
142 ->AlignedAlloc<GetAllocFlags(checks)>(static_cast<size_t>(alignment),
143 count);
144 }
145 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
146 return ::operator new(count, alignment);
147 }
148
149 template <MemorySafetyCheck checks>
HandleMemorySafetyCheckedOperatorDelete(void * ptr)150 NOINLINE void HandleMemorySafetyCheckedOperatorDelete(void* ptr) {
151 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
152 if constexpr (ShouldUsePartitionAlloc(checks)) {
153 GetPartitionRootForMemorySafetyCheckedAllocation()
154 ->Free<GetFreeFlags(checks)>(ptr);
155 return;
156 }
157 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
158 ::operator delete(ptr);
159 }
160
161 template <MemorySafetyCheck checks>
HandleMemorySafetyCheckedOperatorDelete(void * ptr,std::align_val_t alignment)162 NOINLINE void HandleMemorySafetyCheckedOperatorDelete(
163 void* ptr,
164 std::align_val_t alignment) {
165 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
166 if constexpr (ShouldUsePartitionAlloc(checks)) {
167 GetPartitionRootForMemorySafetyCheckedAllocation()
168 ->Free<GetFreeFlags(checks)>(ptr);
169 return;
170 }
171 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
172 ::operator delete(ptr, alignment);
173 }
174
175 } // namespace base::internal
176
177 // Macros to annotate class/struct's default memory safety check.
178 // ADVANCED_MEMORY_SAFETY_CHECKS(): Enable Check |kAdvancedChecks| for this
179 // object.
180 //
181 // Note that if you use this macro at the top of struct declaration, the
182 // declaration context would be left as |private|. Please switch it back to
183 // |public| manually if needed.
184 //
185 // struct ObjectWithAdvancedChecks {
186 // ADVANCED_MEMORY_SAFETY_CHECKS();
187 // public:
188 // int public_field;
189 // };
190 #define ADVANCED_MEMORY_SAFETY_CHECKS_INTERNAL(SPECIFIER) \
191 public: \
192 static constexpr auto kMemorySafetyChecks = \
193 base::internal::kAdvancedMemorySafetyChecks; \
194 SPECIFIER static void* operator new(std::size_t count) { \
195 return base::internal::HandleMemorySafetyCheckedOperatorNew< \
196 kMemorySafetyChecks>(count); \
197 } \
198 SPECIFIER static void* operator new(std::size_t count, \
199 std::align_val_t alignment) { \
200 return base::internal::HandleMemorySafetyCheckedOperatorNew< \
201 kMemorySafetyChecks>(count, alignment); \
202 } \
203 /* Though we do not hook placement new, we need to define this */ \
204 /* explicitly to allow it. */ \
205 ALWAYS_INLINE static void* operator new(std::size_t, void* ptr) { \
206 return ptr; \
207 } \
208 SPECIFIER static void operator delete(void* ptr) noexcept { \
209 base::internal::HandleMemorySafetyCheckedOperatorDelete< \
210 kMemorySafetyChecks>(ptr); \
211 } \
212 SPECIFIER static void operator delete(void* ptr, \
213 std::align_val_t alignment) noexcept { \
214 base::internal::HandleMemorySafetyCheckedOperatorDelete< \
215 kMemorySafetyChecks>(ptr, alignment); \
216 } \
217 \
218 private: \
219 static_assert(true) /* semicolon here */
220
221 #if DCHECK_IS_ON()
222 // Specify NOINLINE to display the operator on a stack trace.
223 #define ADVANCED_MEMORY_SAFETY_CHECKS() \
224 ADVANCED_MEMORY_SAFETY_CHECKS_INTERNAL(NOINLINE NOT_TAIL_CALLED)
225 #else
226 #define ADVANCED_MEMORY_SAFETY_CHECKS() \
227 ADVANCED_MEMORY_SAFETY_CHECKS_INTERNAL(ALWAYS_INLINE)
228 #endif // DCHECK_IS_ON()
229
230 // When a struct/class with `ADVANCED_MEMORY_SAFETY_CHECKS()` is inherited, a
231 // derived struct/class operator will use customized `operator new()` and
232 // `operator delete()` too. If a class has multiple base classes with the macro,
233 // a compiler may complain ambiguity between multiple `operator new()`s. On the
234 // other hand, if a class uses private inheritance, a compiler may report
235 // private `operator new()` that is making impossible to `new` that class. We
236 // have two utility macros to resolve these issues:
237 // - `INHERIT_MEMORY_SAFETY_CHECKS(BaseClass)`
238 // Explicitly exports operators from given `BaseClass` to re-apply
239 // checks specified in the parent class. This is the recommended option as
240 // a derived class is likely to have the same characteristics to its baes
241 // class.
242 // - `DEFAULT_MEMORY_SAFETY_CHECKS()`
243 // Re-define default `operator new()` and `operator delete()` using
244 // global operators that comes with default checks.
245 //
246 // Note that if you use these macros at the top of struct declaration, the
247 // declaration context would be left as |private|. Please switch it back to
248 // |public| manually if needed.
249 #define INHERIT_MEMORY_SAFETY_CHECKS(BASE_CLASS) \
250 public: \
251 using BASE_CLASS::kMemorySafetyChecks; \
252 using BASE_CLASS::operator new; \
253 using BASE_CLASS::operator delete; \
254 \
255 private: \
256 static_assert(true) /* semicolon here */
257
258 #define DEFAULT_MEMORY_SAFETY_CHECKS() \
259 public: \
260 ALWAYS_INLINE static void* operator new(std::size_t count) { \
261 return ::operator new(count); \
262 } \
263 ALWAYS_INLINE static void* operator new(std::size_t count, \
264 std::align_val_t alignment) { \
265 return ::operator new(count, alignment); \
266 } \
267 ALWAYS_INLINE static void* operator new(std::size_t count, void* ptr) { \
268 return ::operator new(count, ptr); \
269 } \
270 ALWAYS_INLINE static void operator delete(void* ptr) noexcept { \
271 return ::operator delete(ptr); \
272 } \
273 ALWAYS_INLINE static void operator delete( \
274 void* ptr, std::align_val_t alignment) noexcept { \
275 return ::operator delete(ptr, alignment); \
276 } \
277 \
278 private: \
279 static_assert(true) /* semicolon here */
280
281 #endif // BASE_MEMORY_SAFETY_CHECKS_H_
282