1 //===-- dfsan_allocator.cpp -------------------------- --------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file is a part of DataflowSanitizer.
10 //
11 // DataflowSanitizer allocator.
12 //===----------------------------------------------------------------------===//
13 
14 #include "dfsan_allocator.h"
15 
16 #include "dfsan.h"
17 #include "dfsan_flags.h"
18 #include "dfsan_thread.h"
19 #include "sanitizer_common/sanitizer_allocator.h"
20 #include "sanitizer_common/sanitizer_allocator_checks.h"
21 #include "sanitizer_common/sanitizer_allocator_interface.h"
22 #include "sanitizer_common/sanitizer_allocator_report.h"
23 #include "sanitizer_common/sanitizer_errno.h"
24 
25 namespace __dfsan {
26 
27 struct Metadata {
28   uptr requested_size;
29 };
30 
31 struct DFsanMapUnmapCallback {
32   void OnMap(uptr p, uptr size) const { dfsan_set_label(0, (void *)p, size); }
33   void OnUnmap(uptr p, uptr size) const { dfsan_set_label(0, (void *)p, size); }
34 };
35 
36 static const uptr kAllocatorSpace = 0x700000000000ULL;
37 static const uptr kMaxAllowedMallocSize = 8UL << 30;
38 
39 struct AP64 {  // Allocator64 parameters. Deliberately using a short name.
40   static const uptr kSpaceBeg = kAllocatorSpace;
41   static const uptr kSpaceSize = 0x40000000000;  // 4T.
42   static const uptr kMetadataSize = sizeof(Metadata);
43   typedef DefaultSizeClassMap SizeClassMap;
44   typedef DFsanMapUnmapCallback MapUnmapCallback;
45   static const uptr kFlags = 0;
46   using AddressSpaceView = LocalAddressSpaceView;
47 };
48 
49 typedef SizeClassAllocator64<AP64> PrimaryAllocator;
50 
51 typedef CombinedAllocator<PrimaryAllocator> Allocator;
52 typedef Allocator::AllocatorCache AllocatorCache;
53 
54 static Allocator allocator;
55 static AllocatorCache fallback_allocator_cache;
56 static StaticSpinMutex fallback_mutex;
57 
58 static uptr max_malloc_size;
59 
60 void dfsan_allocator_init() {
61   SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
62   allocator.Init(common_flags()->allocator_release_to_os_interval_ms);
63   if (common_flags()->max_allocation_size_mb)
64     max_malloc_size = Min(common_flags()->max_allocation_size_mb << 20,
65                           kMaxAllowedMallocSize);
66   else
67     max_malloc_size = kMaxAllowedMallocSize;
68 }
69 
70 AllocatorCache *GetAllocatorCache(DFsanThreadLocalMallocStorage *ms) {
71   CHECK(ms);
72   CHECK_LE(sizeof(AllocatorCache), sizeof(ms->allocator_cache));
73   return reinterpret_cast<AllocatorCache *>(ms->allocator_cache);
74 }
75 
76 void DFsanThreadLocalMallocStorage::CommitBack() {
77   allocator.SwallowCache(GetAllocatorCache(this));
78 }
79 
80 static void *DFsanAllocate(uptr size, uptr alignment, bool zeroise) {
81   if (size > max_malloc_size) {
82     if (AllocatorMayReturnNull()) {
83       Report("WARNING: DataflowSanitizer failed to allocate 0x%zx bytes\n",
84              size);
85       return nullptr;
86     }
87     BufferedStackTrace stack;
88     ReportAllocationSizeTooBig(size, max_malloc_size, &stack);
89   }
90   if (UNLIKELY(IsRssLimitExceeded())) {
91     if (AllocatorMayReturnNull())
92       return nullptr;
93     BufferedStackTrace stack;
94     ReportRssLimitExceeded(&stack);
95   }
96   DFsanThread *t = GetCurrentThread();
97   void *allocated;
98   if (t) {
99     AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
100     allocated = allocator.Allocate(cache, size, alignment);
101   } else {
102     SpinMutexLock l(&fallback_mutex);
103     AllocatorCache *cache = &fallback_allocator_cache;
104     allocated = allocator.Allocate(cache, size, alignment);
105   }
106   if (UNLIKELY(!allocated)) {
107     SetAllocatorOutOfMemory();
108     if (AllocatorMayReturnNull())
109       return nullptr;
110     BufferedStackTrace stack;
111     ReportOutOfMemory(size, &stack);
112   }
113   Metadata *meta =
114       reinterpret_cast<Metadata *>(allocator.GetMetaData(allocated));
115   meta->requested_size = size;
116   if (zeroise) {
117     internal_memset(allocated, 0, size);
118     dfsan_set_label(0, allocated, size);
119   } else if (flags().zero_in_malloc) {
120     dfsan_set_label(0, allocated, size);
121   }
122   return allocated;
123 }
124 
125 void dfsan_deallocate(void *p) {
126   CHECK(p);
127   Metadata *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(p));
128   uptr size = meta->requested_size;
129   meta->requested_size = 0;
130   if (flags().zero_in_free)
131     dfsan_set_label(0, p, size);
132   DFsanThread *t = GetCurrentThread();
133   if (t) {
134     AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
135     allocator.Deallocate(cache, p);
136   } else {
137     SpinMutexLock l(&fallback_mutex);
138     AllocatorCache *cache = &fallback_allocator_cache;
139     allocator.Deallocate(cache, p);
140   }
141 }
142 
143 void *DFsanReallocate(void *old_p, uptr new_size, uptr alignment) {
144   Metadata *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(old_p));
145   uptr old_size = meta->requested_size;
146   uptr actually_allocated_size = allocator.GetActuallyAllocatedSize(old_p);
147   if (new_size <= actually_allocated_size) {
148     // We are not reallocating here.
149     meta->requested_size = new_size;
150     if (new_size > old_size && flags().zero_in_malloc)
151       dfsan_set_label(0, (char *)old_p + old_size, new_size - old_size);
152     return old_p;
153   }
154   uptr memcpy_size = Min(new_size, old_size);
155   void *new_p = DFsanAllocate(new_size, alignment, false /*zeroise*/);
156   if (new_p) {
157     dfsan_copy_memory(new_p, old_p, memcpy_size);
158     dfsan_deallocate(old_p);
159   }
160   return new_p;
161 }
162 
163 void *DFsanCalloc(uptr nmemb, uptr size) {
164   if (UNLIKELY(CheckForCallocOverflow(size, nmemb))) {
165     if (AllocatorMayReturnNull())
166       return nullptr;
167     BufferedStackTrace stack;
168     ReportCallocOverflow(nmemb, size, &stack);
169   }
170   return DFsanAllocate(nmemb * size, sizeof(u64), true /*zeroise*/);
171 }
172 
173 static uptr AllocationSize(const void *p) {
174   if (!p)
175     return 0;
176   const void *beg = allocator.GetBlockBegin(p);
177   if (beg != p)
178     return 0;
179   Metadata *b = (Metadata *)allocator.GetMetaData(p);
180   return b->requested_size;
181 }
182 
183 void *dfsan_malloc(uptr size) {
184   return SetErrnoOnNull(DFsanAllocate(size, sizeof(u64), false /*zeroise*/));
185 }
186 
187 void *dfsan_calloc(uptr nmemb, uptr size) {
188   return SetErrnoOnNull(DFsanCalloc(nmemb, size));
189 }
190 
191 void *dfsan_realloc(void *ptr, uptr size) {
192   if (!ptr)
193     return SetErrnoOnNull(DFsanAllocate(size, sizeof(u64), false /*zeroise*/));
194   if (size == 0) {
195     dfsan_deallocate(ptr);
196     return nullptr;
197   }
198   return SetErrnoOnNull(DFsanReallocate(ptr, size, sizeof(u64)));
199 }
200 
201 void *dfsan_reallocarray(void *ptr, uptr nmemb, uptr size) {
202   if (UNLIKELY(CheckForCallocOverflow(size, nmemb))) {
203     errno = errno_ENOMEM;
204     if (AllocatorMayReturnNull())
205       return nullptr;
206     BufferedStackTrace stack;
207     ReportReallocArrayOverflow(nmemb, size, &stack);
208   }
209   return dfsan_realloc(ptr, nmemb * size);
210 }
211 
212 void *dfsan_valloc(uptr size) {
213   return SetErrnoOnNull(
214       DFsanAllocate(size, GetPageSizeCached(), false /*zeroise*/));
215 }
216 
217 void *dfsan_pvalloc(uptr size) {
218   uptr PageSize = GetPageSizeCached();
219   if (UNLIKELY(CheckForPvallocOverflow(size, PageSize))) {
220     errno = errno_ENOMEM;
221     if (AllocatorMayReturnNull())
222       return nullptr;
223     BufferedStackTrace stack;
224     ReportPvallocOverflow(size, &stack);
225   }
226   // pvalloc(0) should allocate one page.
227   size = size ? RoundUpTo(size, PageSize) : PageSize;
228   return SetErrnoOnNull(DFsanAllocate(size, PageSize, false /*zeroise*/));
229 }
230 
231 void *dfsan_aligned_alloc(uptr alignment, uptr size) {
232   if (UNLIKELY(!CheckAlignedAllocAlignmentAndSize(alignment, size))) {
233     errno = errno_EINVAL;
234     if (AllocatorMayReturnNull())
235       return nullptr;
236     BufferedStackTrace stack;
237     ReportInvalidAlignedAllocAlignment(size, alignment, &stack);
238   }
239   return SetErrnoOnNull(DFsanAllocate(size, alignment, false /*zeroise*/));
240 }
241 
242 void *dfsan_memalign(uptr alignment, uptr size) {
243   if (UNLIKELY(!IsPowerOfTwo(alignment))) {
244     errno = errno_EINVAL;
245     if (AllocatorMayReturnNull())
246       return nullptr;
247     BufferedStackTrace stack;
248     ReportInvalidAllocationAlignment(alignment, &stack);
249   }
250   return SetErrnoOnNull(DFsanAllocate(size, alignment, false /*zeroise*/));
251 }
252 
253 int dfsan_posix_memalign(void **memptr, uptr alignment, uptr size) {
254   if (UNLIKELY(!CheckPosixMemalignAlignment(alignment))) {
255     if (AllocatorMayReturnNull())
256       return errno_EINVAL;
257     BufferedStackTrace stack;
258     ReportInvalidPosixMemalignAlignment(alignment, &stack);
259   }
260   void *ptr = DFsanAllocate(size, alignment, false /*zeroise*/);
261   if (UNLIKELY(!ptr))
262     // OOM error is already taken care of by DFsanAllocate.
263     return errno_ENOMEM;
264   CHECK(IsAligned((uptr)ptr, alignment));
265   *memptr = ptr;
266   return 0;
267 }
268 
269 }  // namespace __dfsan
270 
271 using namespace __dfsan;
272 
273 uptr __sanitizer_get_current_allocated_bytes() {
274   uptr stats[AllocatorStatCount];
275   allocator.GetStats(stats);
276   return stats[AllocatorStatAllocated];
277 }
278 
279 uptr __sanitizer_get_heap_size() {
280   uptr stats[AllocatorStatCount];
281   allocator.GetStats(stats);
282   return stats[AllocatorStatMapped];
283 }
284 
285 uptr __sanitizer_get_free_bytes() { return 1; }
286 
287 uptr __sanitizer_get_unmapped_bytes() { return 1; }
288 
289 uptr __sanitizer_get_estimated_allocated_size(uptr size) { return size; }
290 
291 int __sanitizer_get_ownership(const void *p) { return AllocationSize(p) != 0; }
292 
293 uptr __sanitizer_get_allocated_size(const void *p) { return AllocationSize(p); }
294