[llvm-branch-commits] [nsan] Use sanitizer allocator (PR #102764)

[Fangrui Song via llvm-branch-commits]

================
@@ -0,0 +1,334 @@
+//===- nsan_allocator.cpp -------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// NumericalStabilitySanitizer allocator.
+//
+//===----------------------------------------------------------------------===//
+
+#include "nsan_allocator.h"
+#include "interception/interception.h"
+#include "nsan.h"
+#include "nsan_platform.h"
+#include "nsan_thread.h"
+#include "sanitizer_common/sanitizer_allocator.h"
+#include "sanitizer_common/sanitizer_allocator_checks.h"
+#include "sanitizer_common/sanitizer_allocator_interface.h"
+#include "sanitizer_common/sanitizer_allocator_report.h"
+#include "sanitizer_common/sanitizer_common.h"
+#include "sanitizer_common/sanitizer_errno.h"
+
+DECLARE_REAL(void *, memset, void *dest, int c, uptr n)
+
+using namespace __nsan;
+
+namespace {
+struct Metadata {
+  uptr requested_size;
+};
+
+struct NsanMapUnmapCallback {
+  void OnMap(uptr p, uptr size) const {}
+  void OnMapSecondary(uptr p, uptr size, uptr user_begin,
+                      uptr user_size) const {}
+  void OnUnmap(uptr p, uptr size) const {}
+};
+
+const uptr kMaxAllowedMallocSize = 1ULL << 40;
+
+// Allocator64 parameters. Deliberately using a short name.
+struct AP64 {
+  static const uptr kSpaceBeg = Mapping::kHeapMemBeg;
+  static const uptr kSpaceSize = 0x40000000000; // 4T.
+  static const uptr kMetadataSize = sizeof(Metadata);
+  using SizeClassMap = DefaultSizeClassMap;
+  using MapUnmapCallback = NsanMapUnmapCallback;
+  static const uptr kFlags = 0;
+  using AddressSpaceView = LocalAddressSpaceView;
+};
+} // namespace
+
+using PrimaryAllocator = SizeClassAllocator64<AP64>;
+using Allocator = CombinedAllocator<PrimaryAllocator>;
+using AllocatorCache = Allocator::AllocatorCache;
+
+static Allocator allocator;
+static AllocatorCache fallback_allocator_cache;
+static StaticSpinMutex fallback_mutex;
+
+static uptr max_malloc_size;
+
+void __nsan::NsanAllocatorInit() {
+  SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
+  allocator.Init(common_flags()->allocator_release_to_os_interval_ms);
+  if (common_flags()->max_allocation_size_mb)
+    max_malloc_size = Min(common_flags()->max_allocation_size_mb << 20,
+                          kMaxAllowedMallocSize);
+  else
+    max_malloc_size = kMaxAllowedMallocSize;
+}
+
+static AllocatorCache *GetAllocatorCache(NsanThreadLocalMallocStorage *ms) {
+  CHECK(ms);
+  CHECK_LE(sizeof(AllocatorCache), sizeof(ms->allocator_cache));
+  return reinterpret_cast<AllocatorCache *>(ms->allocator_cache);
+}
+
+void NsanThreadLocalMallocStorage::Init() {
+  allocator.InitCache(GetAllocatorCache(this));
+}
+
+void NsanThreadLocalMallocStorage::CommitBack() {
+  allocator.SwallowCache(GetAllocatorCache(this));
+  allocator.DestroyCache(GetAllocatorCache(this));
+}
+
+static void *NsanAllocate(uptr size, uptr alignment, bool zero) {
+  if (UNLIKELY(size > max_malloc_size)) {
+    if (AllocatorMayReturnNull()) {
+      Report("WARNING: NumericalStabilitySanitizer failed to allocate 0x%zx "
+             "bytes\n",
+             size);
+      return nullptr;
+    }
+    BufferedStackTrace stack;
+    GET_FATAL_STACK_TRACE_IF_EMPTY(&stack);
+    ReportAllocationSizeTooBig(size, max_malloc_size, &stack);
+  }
+  if (UNLIKELY(IsRssLimitExceeded())) {
+    if (AllocatorMayReturnNull())
+      return nullptr;
+    BufferedStackTrace stack;
+    GET_FATAL_STACK_TRACE_IF_EMPTY(&stack);
+    ReportRssLimitExceeded(&stack);
+  }
+  NsanThread *t = GetCurrentThread();
+  void *allocated;
+  if (t) {
+    AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
+    allocated = allocator.Allocate(cache, size, alignment);
+  } else {
+    SpinMutexLock l(&fallback_mutex);
+    AllocatorCache *cache = &fallback_allocator_cache;
+    allocated = allocator.Allocate(cache, size, alignment);
+  }
+  if (UNLIKELY(!allocated)) {
+    SetAllocatorOutOfMemory();
+    if (AllocatorMayReturnNull())
+      return nullptr;
+    BufferedStackTrace stack;
+    GET_FATAL_STACK_TRACE_IF_EMPTY(&stack);
+    ReportOutOfMemory(size, &stack);
+  }
+  auto *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(allocated));
+  meta->requested_size = size;
+  if (zero && allocator.FromPrimary(allocated))
+    REAL(memset)(allocated, 0, size);
+  __nsan_set_value_unknown(allocated, size);
+  RunMallocHooks(allocated, size);
+  return allocated;
+}
+
+void __nsan::NsanDeallocate(void *p) {
+  DCHECK(p);
+  RunFreeHooks(p);
+  auto *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(p));
+  meta->requested_size = 0;
+  if (NsanThread *t = GetCurrentThread()) {
+    AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
+    allocator.Deallocate(cache, p);
+  } else {
+    SpinMutexLock l(&fallback_mutex);
+    AllocatorCache *cache = &fallback_allocator_cache;
+    allocator.Deallocate(cache, p);
+  }
+}
+
+static void *NsanReallocate(void *ptr, uptr new_size, uptr alignment) {
+  Metadata *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(ptr));
+  uptr old_size = meta->requested_size;
+  uptr actually_allocated_size = allocator.GetActuallyAllocatedSize(ptr);
+  if (new_size <= actually_allocated_size) {
+    // We are not reallocating here.
+    meta->requested_size = new_size;
+    if (new_size > old_size)
+      __nsan_set_value_unknown((u8 *)ptr + old_size, new_size - old_size);
+    return ptr;
+  }
+  void *new_p = NsanAllocate(new_size, alignment, false);
+  if (new_p) {
+    __nsan_copy_values(new_p, ptr, Min(new_size, old_size));
+    NsanDeallocate(ptr);
----------------
MaskRay wrote:

`REAL(memcpy)` was missing (sigsegv when testing sqlite-amalgamation) and is now fixed. 

Since we don't check uninitialized memory, whether to call set_value_unknown in NsanDeallocate doesn't matter that much. That said, I added `poison_in_free` to be similar to msan.

https://github.com/llvm/llvm-project/pull/102764