[asan] fix coverage between fork() and exec(): reinitialize the guards after fork()

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@225016 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 37 insertions, 13 deletions

diff --git a/lib/sanitizer_common/sanitizer_common.h b/lib/sanitizer_common/sanitizer_common.h
index ebbb00324..a25bacb3f 100644
--- a/lib/sanitizer_common/sanitizer_common.h
+++ b/lib/sanitizer_common/sanitizer_common.h
@@ -380,14 +380,14 @@ INLINE int ToLower(int c) {
 // small vectors.
 // WARNING: The current implementation supports only POD types.
 template<typename T>
-class InternalMmapVector {
+class InternalMmapVectorNoCtor {
  public:
-  explicit InternalMmapVector(uptr initial_capacity) {
+  void Initialize(uptr initial_capacity) {
     capacity_ = Max(initial_capacity, (uptr)1);
     size_ = 0;
-    data_ = (T *)MmapOrDie(capacity_ * sizeof(T), "InternalMmapVector");
+    data_ = (T *)MmapOrDie(capacity_ * sizeof(T), "InternalMmapVectorNoCtor");
   }
-  ~InternalMmapVector() {
+  void Destroy() {
     UnmapOrDie(data_, capacity_ * sizeof(T));
   }
   T &operator[](uptr i) {
@@ -438,15 +438,24 @@ class InternalMmapVector {
     UnmapOrDie(old_data, capacity_ * sizeof(T));
     capacity_ = new_capacity;
   }
-  // Disallow evil constructors.
-  InternalMmapVector(const InternalMmapVector&);
-  void operator=(const InternalMmapVector&);
 
   T *data_;
   uptr capacity_;
   uptr size_;
 };
 
+template<typename T>
+class InternalMmapVector : public InternalMmapVectorNoCtor<T> {
+ public:
+  explicit InternalMmapVector(uptr initial_capacity) {
+    InternalMmapVectorNoCtor<T>::Initialize(initial_capacity);
+  }
+  ~InternalMmapVector() { InternalMmapVectorNoCtor<T>::Destroy(); }
+  // Disallow evil constructors.
+  InternalMmapVector(const InternalMmapVector&);
+  void operator=(const InternalMmapVector&);
+};
+
 // HeapSort for arrays and InternalMmapVector.
 template<class Container, class Compare>
 void InternalSort(Container *v, uptr size, Compare comp) {
diff --git a/lib/sanitizer_common/sanitizer_coverage_libcdep.cc b/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
index eb69172c7..dbb7f9776 100644
--- a/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
+++ b/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
@@ -79,6 +79,7 @@ class CoverageData {
   ALWAYS_INLINE
   void TraceBasicBlock(uptr *cache);
 
+  void InitializeGuardArray(s32 *guards);
   void InitializeGuards(s32 *guards, uptr n);
 
   uptr *data();
@@ -103,6 +104,9 @@ class CoverageData {
   // Descriptor of the file mapped pc array.
   int pc_fd;
 
+  // Vector of coverage guard arrays, protected by mu.
+  InternalMmapVectorNoCtor<s32*> guard_array_vec;
+
   // Caller-Callee (cc) array, size and current index.
   static const uptr kCcArrayMaxSize = FIRST_32_SECOND_64(1 << 18, 1 << 24);
   uptr **cc_array;
@@ -181,6 +185,14 @@ void CoverageData::Init() {
   tr_pc_array_index = 0;
 }
 
+void CoverageData::InitializeGuardArray(s32 *guards) {
+  s32 n = guards[0];
+  for (s32 j = 1; j <= n; j++) {
+    uptr idx = atomic_fetch_add(&pc_array_index, 1, memory_order_relaxed);
+    guards[j] = -static_cast<s32>(idx + 1);
+  }
+}
+
 void CoverageData::ReInit() {
   if (pc_array) {
     internal_munmap(pc_array, sizeof(uptr) * kPcArrayMaxSize);
@@ -199,6 +211,11 @@ void CoverageData::ReInit() {
       Init();
     }
   }
+  // Re-initialize the guards.
+  // We are single-threaded now, no need to grab any lock.
+  CHECK_EQ(atomic_load(&pc_array_index, memory_order_relaxed), 0);
+  for (uptr i = 0; i < guard_array_vec.size(); i++)
+    InitializeGuardArray(guard_array_vec[i]);
 }
 
 void CoverageData::BeforeFork() {
@@ -251,10 +268,9 @@ void CoverageData::InitializeGuards(s32 *guards, uptr n) {
   // to store 'n'.
   CHECK_LT(n, 1 << 30);
   guards[0] = static_cast<s32>(n);
-  for (uptr i = 1; i <= n; i++) {
-    uptr idx = atomic_fetch_add(&pc_array_index, 1, memory_order_relaxed);
-    guards[i] = -static_cast<s32>(idx + 1);
-  }
+  InitializeGuardArray(guards);
+  SpinMutexLock l(&mu);
+  guard_array_vec.push_back(guards);
 }
 
 // If guard is negative, atomically set it to -guard and store the PC in
diff --git a/test/asan/TestCases/Linux/coverage-fork.cc b/test/asan/TestCases/Linux/coverage-fork.cc
index b9e981fcf..38c200942 100644
--- a/test/asan/TestCases/Linux/coverage-fork.cc
+++ b/test/asan/TestCases/Linux/coverage-fork.cc
@@ -33,7 +33,6 @@ int main(int argc, char **argv) {
 }
 
 // CHECK-DAG: Child PID: [[ChildPID:[0-9]+]]
-// Coverage in the child (after fork,before exec) is not expected to work.
-// (disabled): [[ChildPID]].sancov: 1 PCs written
+// CHECK-DAG: [[ChildPID]].sancov: 1 PCs written
 // CHECK-DAG: Parent PID: [[ParentPID:[0-9]+]]
 // CHECK-DAG: [[ParentPID]].sancov: 3 PCs written