Increase coverage map size dynamically

As we are now instrumenting a known number of edges, we can start with
a small coverage map that grows as needed.

The initial size of the map is 512 and doubles whenever the instrumentor
runs out of new edge IDs.

As the bytecode instrumentation needs to have access to a contiguous
ByteBuffer, we allocate a (hopefully large enough for all practical
purposes) coverage counter buffer in native code once and only adjust
the size of the ByteBuffer dynamically. As libFuzzer is only made aware
of the part of the buffer in active use, this has no performance impact
and only very minor memory usage (< 10 MB).

6 files changed, 110 insertions, 49 deletions

diff --git a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt
index 5694c331..1aa9597c 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt
+++ b/agent/src/main/java/com/code_intelligence/jazzer/instrumentor/EdgeCoverageInstrumentor.kt
@@ -62,7 +62,14 @@ object EdgeCoverageInstrumentor : Instrumentor {
 
 private var nextGlobalEdgeId = 0
 
-private fun nextEdgeId(): Int = nextGlobalEdgeId++
+private fun nextEdgeId(): Int {
+    if (nextGlobalEdgeId >= CoverageMap.mem.capacity()) {
+        if (!EdgeCoverageInstrumentor.isTesting) {
+            CoverageMap.enlargeCoverageMap()
+        }
+    }
+    return nextGlobalEdgeId++
+}
 
 /**
  * The maximal number of stack elements used by [loadCoverageMap].
diff --git a/agent/src/main/java/com/code_intelligence/jazzer/runtime/CoverageMap.java b/agent/src/main/java/com/code_intelligence/jazzer/runtime/CoverageMap.java
index 903f416f..af2424a2 100644
--- a/agent/src/main/java/com/code_intelligence/jazzer/runtime/CoverageMap.java
+++ b/agent/src/main/java/com/code_intelligence/jazzer/runtime/CoverageMap.java
@@ -22,7 +22,12 @@ import java.nio.ByteBuffer;
  * native code.
  */
 final public class CoverageMap {
-  public static int SIZE = 65536;
-  public static ByteBuffer mem = ByteBuffer.allocateDirect(SIZE);
+  public static ByteBuffer mem = ByteBuffer.allocateDirect(0);
 
+  public static void enlargeCoverageMap() {
+    registerNewCoverageCounters();
+    System.out.println("INFO: New number of inline 8-bit counters: " + mem.capacity());
+  }
+
+  private static native void registerNewCoverageCounters();
 }
diff --git a/driver/coverage_tracker.cpp b/driver/coverage_tracker.cpp
index ed1d51eb..213a7b19 100644
--- a/driver/coverage_tracker.cpp
+++ b/driver/coverage_tracker.cpp
@@ -15,8 +15,10 @@
 #include "coverage_tracker.h"
 
 #include <algorithm>
+#include <memory>
+#include <stdexcept>
 
-#include "glog/logging.h"
+#include "absl/strings/str_format.h"
 #include "third_party/jni/jni.h"
 
 extern "C" void __sanitizer_cov_8bit_counters_init(uint8_t *start,
@@ -26,27 +28,35 @@ extern "C" void __sanitizer_cov_pcs_init(const uintptr_t *pcs_beg,
 
 constexpr auto kCoverageMapClass =
     "com/code_intelligence/jazzer/runtime/CoverageMap";
+constexpr auto kByteBufferClass = "java/nio/ByteBuffer";
+
+// The initial size of the Java coverage map (512 counters).
+constexpr std::size_t kInitialCoverageCountersBufferSize = 1u << 9u;
+// The maximum size of the Java coverage map (1,048,576 counters).
+// Since the memory for the coverage map needs to be allocated contiguously,
+// increasing the maximum size incurs additional memory (but not runtime)
+// overhead for all fuzz targets.
+constexpr std::size_t kMaxCoverageCountersBufferSize = 1u << 20u;
+static_assert(kMaxCoverageCountersBufferSize <=
+              std::numeric_limits<jint>::max());
 
 namespace jazzer {
 
-CoverageTracker::CoverageTracker(JVM &jvm) : ExceptionPrinter(jvm) {
-  auto &env = jvm.GetEnv();
-  jclass coverage_map = jvm.FindClass(kCoverageMapClass);
-  jfieldID counters_buffer_id =
-      jvm.GetStaticFieldID(coverage_map, "mem", "Ljava/nio/ByteBuffer;");
-  jobject counters_buffer =
-      env.GetStaticObjectField(coverage_map, counters_buffer_id);
-  counters_ =
-      reinterpret_cast<uint8_t *>(env.GetDirectBufferAddress(counters_buffer));
-  if (counters_ == nullptr) {
-    throw std::runtime_error("Failed to get coverage map address");
-  }
-  counters_size_ = env.GetDirectBufferCapacity(counters_buffer);
-  if (env.ExceptionOccurred()) {
-    LOG(ERROR) << getAndClearException();
-    throw std::runtime_error("failed to retrieve Java coverage buffer");
+uint8_t *CoverageTracker::counters_ = nullptr;
+uint32_t *CoverageTracker::fake_instructions_ = nullptr;
+PCTableEntry *CoverageTracker::pc_entries_ = nullptr;
+
+void CoverageTracker::Setup(JNIEnv &env) {
+  if (counters_ != nullptr) {
+    throw std::runtime_error(
+        "CoverageTracker::Setup must not be called more than once");
   }
-  __sanitizer_cov_8bit_counters_init(counters_, counters_ + counters_size_);
+  JNINativeMethod coverage_tracker_native_methods[]{
+      {(char *)"registerNewCoverageCounters", (char *)"()V",
+       (void *)&RegisterNewCoverageCounters},
+  };
+  jclass coverage_map = env.FindClass(kCoverageMapClass);
+  env.RegisterNatives(coverage_map, coverage_tracker_native_methods, 1);
 
   // libFuzzer requires an array containing the instruction addresses associated
   // with the coverage counters registered above. Given that we are
@@ -56,24 +66,66 @@ CoverageTracker::CoverageTracker(JVM &jvm) : ExceptionPrinter(jvm) {
   // allocated buffer. Note: We intentionally never deallocate the allocations
   // made here as they have static lifetime and we can't guarantee they wouldn't
   // be freed before libFuzzer stops using them.
-  fake_instructions_ = new uint32_t[counters_size_];
-  std::fill(fake_instructions_, fake_instructions_ + counters_size_, 0);
+  constexpr std::size_t counters_size = kMaxCoverageCountersBufferSize;
+  counters_ = new uint8_t[counters_size];
+  Clear();
+
+  // Never deallocated, see above.
+  fake_instructions_ = new uint32_t[counters_size];
+  std::fill(fake_instructions_, fake_instructions_ + counters_size, 0);
+
   // Never deallocated, see above.
-  pc_entries_ = new PCTableEntry[counters_size_];
-  for (std::size_t i = 0; i < counters_size_; ++i) {
+  pc_entries_ = new PCTableEntry[counters_size];
+  for (std::size_t i = 0; i < counters_size; ++i) {
     pc_entries_[i].PC = reinterpret_cast<uintptr_t>(fake_instructions_ + i);
-    // We can't use libFuzzer's value profile tracing of caller-callee
-    // relationships as it relies on compiler built-ins to retrieve the callee
-    // automatically, which does not work with Java methods. We thus don't
-    // report any function PCs.
+    // TODO: Label Java PCs corresponding to functions as such.
     pc_entries_[i].PCFlags = 0;
   }
-  __sanitizer_cov_pcs_init((uintptr_t *)pc_entries_,
-                           (uintptr_t *)(pc_entries_ + counters_size_));
+
+  // Register the first batch of coverage counters.
+  RegisterNewCoverageCounters(env, nullptr);
+}
+
+void JNICALL CoverageTracker::RegisterNewCoverageCounters(JNIEnv &env,
+                                                          jclass cls) {
+  jclass coverage_map = env.FindClass(kCoverageMapClass);
+  jfieldID counters_buffer_id = env.GetStaticFieldID(
+      coverage_map, "mem", absl::StrFormat("L%s;", kByteBufferClass).c_str());
+  jobject counters_buffer =
+      env.GetStaticObjectField(coverage_map, counters_buffer_id);
+
+  jclass byte_buffer = env.FindClass(kByteBufferClass);
+  jmethodID byte_buffer_capacity_id =
+      env.GetMethodID(byte_buffer, "capacity", "()I");
+  jint old_counters_buffer_size =
+      env.CallIntMethod(counters_buffer, byte_buffer_capacity_id);
+
+  jint new_counters_buffer_size;
+  if (old_counters_buffer_size == 0) {
+    new_counters_buffer_size = kInitialCoverageCountersBufferSize;
+  } else {
+    new_counters_buffer_size = 2 * old_counters_buffer_size;
+    if (new_counters_buffer_size > kMaxCoverageCountersBufferSize) {
+      throw std::runtime_error(
+          "Maximal size of the coverage counters buffer exceeded");
+    }
+  }
+
+  jobject new_counters_buffer = env.NewDirectByteBuffer(
+      static_cast<void *>(counters_), new_counters_buffer_size);
+  env.SetStaticObjectField(coverage_map, counters_buffer_id,
+                           new_counters_buffer);
+
+  // Register only the new second half of the counters buffer with libFuzzer.
+  __sanitizer_cov_8bit_counters_init(counters_ + old_counters_buffer_size,
+                                     counters_ + new_counters_buffer_size);
+  __sanitizer_cov_pcs_init(
+      (uintptr_t *)(pc_entries_ + old_counters_buffer_size),
+      (uintptr_t *)(pc_entries_ + new_counters_buffer_size));
 }
 
 void CoverageTracker::Clear() {
-  std::fill(counters_, counters_ + counters_size_, 0);
+  std::fill(counters_, counters_ + kMaxCoverageCountersBufferSize, 0);
 }
 
 uint8_t *CoverageTracker::GetCoverageCounters() { return counters_; }
diff --git a/driver/coverage_tracker.h b/driver/coverage_tracker.h
index eb4b0259..79d8112d 100644
--- a/driver/coverage_tracker.h
+++ b/driver/coverage_tracker.h
@@ -23,7 +23,7 @@ namespace jazzer {
 
 // The members of this struct are only accessed by libFuzzer.
 struct __attribute__((packed)) PCTableEntry {
-  uintptr_t PC, PCFlags;
+  [[maybe_unused]] uintptr_t PC, PCFlags;
 };
 
 // CoverageTracker registers an array of 8-bit coverage counters with
@@ -31,23 +31,21 @@ struct __attribute__((packed)) PCTableEntry {
 // side, where it is populated with the actual coverage information.
 class CoverageTracker : public ExceptionPrinter {
  private:
-  uint8_t *counters_;
-  std::size_t counters_size_;
+  static uint8_t *counters_;
 
-  uint32_t *fake_instructions_;
-  PCTableEntry *pc_entries_;
+  static uint32_t *fake_instructions_;
+  static PCTableEntry *pc_entries_;
 
- public:
-  // Construct the coverage tracker. If the corresponding java class and method
-  // cannot be found it will throw std::runtime_error.
-  explicit CoverageTracker(JVM &jvm);
+  static void JNICALL RegisterNewCoverageCounters(JNIEnv &env, jclass cls);
 
+ public:
+  static void Setup(JNIEnv &env);
   // Clears the coverage counters array manually. It is cleared automatically
   // by libFuzzer prior to running the fuzz target, so this function is only
   // used in tests.
-  void Clear();
+  static void Clear();
 
   // Returns the address of the coverage counters array.
-  uint8_t *GetCoverageCounters();
+  static uint8_t *GetCoverageCounters();
 };
 }  // namespace jazzer
diff --git a/driver/jvm_tooling_test.cpp b/driver/jvm_tooling_test.cpp
index 0558385c..fa0ba747 100644
--- a/driver/jvm_tooling_test.cpp
+++ b/driver/jvm_tooling_test.cpp
@@ -45,6 +45,7 @@ class JvmToolingTest : public ::testing::Test {
     FLAGS_instrumentation_excludes = "**";
 
     jvm_ = std::make_unique<JVM>("test_executable");
+    CoverageTracker::Setup(jvm_->GetEnv());
   }
 
   static void TearDownTestCase() { jvm_.reset(nullptr); }
@@ -166,11 +167,9 @@ TEST_F(JvmToolingTest, FuzzTargetWithInit) {
 }
 
 TEST_F(JvmToolingTest, TestCoverageMap) {
-  CoverageTracker coverage_tracker(*jvm_);
-  coverage_tracker.Clear();
-
+  CoverageTracker::Clear();
   // check that after the initial clear the first coverage counter is 0
-  auto coverage_counters_array = coverage_tracker.GetCoverageCounters();
+  auto coverage_counters_array = CoverageTracker::GetCoverageCounters();
   ASSERT_EQ(0, coverage_counters_array[0]);
 
   FLAGS_target_class = "test/FuzzTargetWithCoverage";
@@ -180,7 +179,7 @@ TEST_F(JvmToolingTest, TestCoverageMap) {
   // increase
   fuzz_target_runner.Run(nullptr, 0);
   ASSERT_EQ(1, coverage_counters_array[0]);
-  coverage_tracker.Clear();
+  CoverageTracker::Clear();
   // back to initial state
   ASSERT_EQ(0, coverage_counters_array[0]);
 
diff --git a/driver/libfuzzer_driver.cpp b/driver/libfuzzer_driver.cpp
index a38991cb..8bdc3030 100644
--- a/driver/libfuzzer_driver.cpp
+++ b/driver/libfuzzer_driver.cpp
@@ -88,7 +88,7 @@ void AbstractLibfuzzerDriver::initJvm(const std::string &executable_path) {
   jvm_ = std::make_unique<jazzer::JVM>(executable_path);
   if (FLAGS_hooks) {
     jazzer::registerFuzzerCallbacks(jvm_->GetEnv());
-    coverage_tracker_ = std::make_unique<jazzer::CoverageTracker>(*jvm_);
+    CoverageTracker::Setup(jvm_->GetEnv());
     // SignalHandler registers its own native methods
     signal_handler_ = std::make_unique<jazzer::SignalHandler>(*jvm_);
     signal_handler_->SetupSignalHandlers();