From d9b574c312c34ddc88dee7fdb7f55fe748839683 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka@google.com>
Date: Wed, 24 May 2023 16:51:12 -0700
Subject: [PATCH] [hwasan] Fix allocator_interface implementation

__sanitizer_get_current_allocated_bytes had as body, but allocator
caches were not registered to collect stats. It's done by
SizeClassAllocator64LocalCache::Init().

Reviewed By: thurston

Differential Revision: https://reviews.llvm.org/D151389
---
 compiler-rt/lib/hwasan/hwasan_allocator.cpp                       | 8 ++++++--
 compiler-rt/lib/hwasan/hwasan_allocator.h                         | 1 +
 compiler-rt/lib/hwasan/hwasan_thread.cpp                          | 1 +
 .../test/sanitizer_common/TestCases/allocator_interface.cpp       | 3 ---
 4 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/compiler-rt/lib/hwasan/hwasan_allocator.cpp b/compiler-rt/lib/hwasan/hwasan_allocator.cpp
index b81f1e3..d3e8266 100644
--- a/compiler-rt/lib/hwasan/hwasan_allocator.cpp
+++ b/compiler-rt/lib/hwasan/hwasan_allocator.cpp
@@ -149,8 +149,9 @@ void HwasanAllocatorInit() {
   atomic_store_relaxed(&hwasan_allocator_tagging_enabled,
                        !flags()->disable_allocator_tagging);
   SetAllocatorMayReturnNull(common_flags()->allocator_may_return_null);
-  allocator.Init(common_flags()->allocator_release_to_os_interval_ms,
-                 GetAliasRegionStart());
+  allocator.InitLinkerInitialized(
+      common_flags()->allocator_release_to_os_interval_ms,
+      GetAliasRegionStart());
   for (uptr i = 0; i < sizeof(tail_magic); i++)
     tail_magic[i] = GetCurrentThread()->GenerateRandomTag();
   if (common_flags()->max_allocation_size_mb) {
@@ -165,8 +166,11 @@ void HwasanAllocatorLock() { allocator.ForceLock(); }
 
 void HwasanAllocatorUnlock() { allocator.ForceUnlock(); }
 
+void AllocatorThreadStart(AllocatorCache *cache) { allocator.InitCache(cache); }
+
 void AllocatorThreadFinish(AllocatorCache *cache) {
   allocator.SwallowCache(cache);
+  allocator.DestroyCache(cache);
 }
 
 static uptr TaggedSize(uptr size) {
diff --git a/compiler-rt/lib/hwasan/hwasan_allocator.h b/compiler-rt/lib/hwasan/hwasan_allocator.h
index 0168b49..efe253d 100644
--- a/compiler-rt/lib/hwasan/hwasan_allocator.h
+++ b/compiler-rt/lib/hwasan/hwasan_allocator.h
@@ -88,6 +88,7 @@ typedef SizeClassAllocator64<AP64> PrimaryAllocator;
 typedef CombinedAllocator<PrimaryAllocator> Allocator;
 typedef Allocator::AllocatorCache AllocatorCache;
 
+void AllocatorThreadStart(AllocatorCache *cache);
 void AllocatorThreadFinish(AllocatorCache *cache);
 
 class HwasanChunkView {
diff --git a/compiler-rt/lib/hwasan/hwasan_thread.cpp b/compiler-rt/lib/hwasan/hwasan_thread.cpp
index b8a9d28..f5eb79f 100644
--- a/compiler-rt/lib/hwasan/hwasan_thread.cpp
+++ b/compiler-rt/lib/hwasan/hwasan_thread.cpp
@@ -58,6 +58,7 @@ void Thread::Init(uptr stack_buffer_start, uptr stack_buffer_size,
 #endif
   InitStackAndTls(state);
   dtls_ = DTLS_Get();
+  AllocatorThreadStart(allocator_cache());
 }
 
 void Thread::InitStackRingBuffer(uptr stack_buffer_start,
diff --git a/compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp b/compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp
index 9d98fca..c2f3044 100644
--- a/compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp
+++ b/compiler-rt/test/sanitizer_common/TestCases/allocator_interface.cpp
@@ -4,9 +4,6 @@
 // No allocator.
 // UNSUPPORTED: ubsan
 
-// FIXME: implementation is incomplete.
-// XFAIL: hwasan
-
 #include <assert.h>
 #include <sanitizer/allocator_interface.h>
 #include <stdio.h>
-- 
2.7.4