From 21a1a237903a125b928895996ec77f42f173f8d5 Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Wed, 28 Jan 2015 22:39:44 +0000
Subject: [PATCH] [sanitizer] allow to reset the bb/edge coverage data inside
 the process while it is running (single-threaded). Also expose the current
 coverage set to the process.

llvm-svn: 227387
---
 .../include/sanitizer/common_interface_defs.h      |  8 ++++
 .../sanitizer_common/sanitizer_coverage_libcdep.cc | 23 +++++++++-
 .../test/asan/TestCases/Linux/coverage-reset.cc    | 52 ++++++++++++++++++++++
 3 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 compiler-rt/test/asan/TestCases/Linux/coverage-reset.cc

diff --git a/compiler-rt/include/sanitizer/common_interface_defs.h b/compiler-rt/include/sanitizer/common_interface_defs.h
index 4c715c6..cfb06ea 100644
--- a/compiler-rt/include/sanitizer/common_interface_defs.h
+++ b/compiler-rt/include/sanitizer/common_interface_defs.h
@@ -74,6 +74,14 @@ extern "C" {
   // This can be useful for coverage-directed in-process fuzzers.
   uintptr_t __sanitizer_get_total_unique_coverage();
 
+  // Reset the basic-block (edge) coverage to the initial state.
+  // Useful for in-process fuzzing to start collecting coverage from scratch.
+  // Experimental, will likely not work for multi-threaded process.
+  void __sanitizer_reset_coverage();
+  // Set *data to the array of covered PCs and return the size of that array.
+  // Some of the entries in *data will be zero.
+  uintptr_t __sanitizer_get_coverage_guards(uintptr_t **data);
+
   // Annotate the current state of a contiguous container, such as
   // std::vector, std::string or similar.
   // A contiguous container is a container that keeps all of its elements
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_coverage_libcdep.cc b/compiler-rt/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
index 0048f67..71f7bd4 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_coverage_libcdep.cc
@@ -83,6 +83,7 @@ class CoverageData {
 
   void InitializeGuardArray(s32 *guards);
   void InitializeGuards(s32 *guards, uptr n);
+  void ReinitializeGuards();
 
   uptr *data();
   uptr size();
@@ -211,6 +212,13 @@ void CoverageData::Disable() {
   }
 }
 
+void CoverageData::ReinitializeGuards() {
+  // Assuming single thread.
+  atomic_store(&pc_array_index, 0, memory_order_relaxed);
+  for (uptr i = 0; i < guard_array_vec.size(); i++)
+    InitializeGuardArray(guard_array_vec[i]);
+}
+
 void CoverageData::ReInit() {
   Disable();
   if (coverage_enabled) {
@@ -228,8 +236,7 @@ void CoverageData::ReInit() {
   // 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]);
+  ReinitializeGuards();
 }
 
 void CoverageData::BeforeFork() {
@@ -689,4 +696,16 @@ SANITIZER_INTERFACE_ATTRIBUTE
 void __sanitizer_cov_trace_basic_block(s32 *id) {
   coverage_data.TraceBasicBlock(id);
 }
+SANITIZER_INTERFACE_ATTRIBUTE
+void __sanitizer_reset_coverage() {
+  coverage_data.ReinitializeGuards();
+  internal_bzero_aligned16(
+      coverage_data.data(),
+      RoundUpTo(coverage_data.size() * sizeof(coverage_data.data()[0]), 16));
+}
+SANITIZER_INTERFACE_ATTRIBUTE
+uptr __sanitizer_get_coverage_guards(uptr **data) {
+  *data = coverage_data.data();
+  return coverage_data.size();
+}
 }  // extern "C"
diff --git a/compiler-rt/test/asan/TestCases/Linux/coverage-reset.cc b/compiler-rt/test/asan/TestCases/Linux/coverage-reset.cc
new file mode 100644
index 0000000..dedea9e
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/Linux/coverage-reset.cc
@@ -0,0 +1,52 @@
+// Test __sanitizer_reset_coverage().
+
+// RUN: %clangxx_asan -fsanitize-coverage=1 %s -o %t
+// RUN: ASAN_OPTIONS=coverage=1 %run %t
+
+#include <sanitizer/common_interface_defs.h>
+#include <stdio.h>
+#include <assert.h>
+static volatile int sink;
+__attribute__((noinline)) void bar() { sink = 2; }
+__attribute__((noinline)) void foo() { sink = 1; }
+
+#define GET_AND_PRINT_COVERAGE()                                       \
+  bitset = 0;                                                  \
+  for (size_t i = 0; i < n_guards; i++)                        \
+    if (guards[i]) bitset |= 1U << i;                          \
+  printf("line %d: bitset %zd total: %zd\n", __LINE__, bitset, \
+         __sanitizer_get_total_unique_coverage());
+
+#define IS_POWER_OF_TWO(a) ((a & ((a) - 1)) == 0)
+
+int main() {
+  size_t *guards = 0;
+  size_t bitset;
+  size_t n_guards = __sanitizer_get_coverage_guards(&guards);
+
+  GET_AND_PRINT_COVERAGE();
+  size_t main_bit = bitset;
+  assert(IS_POWER_OF_TWO(main_bit));
+
+  foo();
+  GET_AND_PRINT_COVERAGE();
+  size_t foo_bit = bitset & ~main_bit;
+  assert(IS_POWER_OF_TWO(foo_bit));
+
+  bar();
+  GET_AND_PRINT_COVERAGE();
+  size_t bar_bit = bitset & ~(main_bit | foo_bit);
+  assert(IS_POWER_OF_TWO(bar_bit));
+
+  __sanitizer_reset_coverage();
+  GET_AND_PRINT_COVERAGE();
+  assert(bitset == 0);
+
+  foo();
+  GET_AND_PRINT_COVERAGE();
+  assert(bitset == foo_bit);
+
+  bar();
+  GET_AND_PRINT_COVERAGE();
+  assert(bitset == (foo_bit | bar_bit));
+}
-- 
2.7.4