Register native callbacks on ClassInstrumentor initialization

The native callbacks in the driver should be registered as early as
possible since they may already be needed during agent startup (e.g.,
if a JDK-internal class is to be instrumented).

7 files changed, 58 insertions, 47 deletions

diff --git a/driver/BUILD.bazel b/driver/BUILD.bazel
index cc4559a1..f7b09efe 100644
--- a/driver/BUILD.bazel
+++ b/driver/BUILD.bazel
@@ -44,6 +44,9 @@ cc_library(
     linkopts = [
         "-ldl",
     ],
+    # Needs to be linked statically for JNI_OnLoad_jazzer_initialize to be found
+    # by the JVM.
+    linkstatic = True,
     visibility = ["//visibility:public"],
     deps = [
         ":sanitizer_hooks_with_pc",
@@ -123,6 +126,11 @@ cc_test(
         "//driver/testdata:fuzz_target_mocks_deploy.jar",
     ],
     includes = ["."],
+    linkopts = [
+        # Needs to export symbols dynamically for JNI_OnLoad_jazzer_initialize
+        # to be found by the JVM.
+        "-rdynamic",
+    ],
     deps = [
         ":jvm_tooling_lib",
         ":test_main",
diff --git a/driver/jvm_tooling.cpp b/driver/jvm_tooling.cpp
index 1d27897e..ddfbfd0f 100644
--- a/driver/jvm_tooling.cpp
+++ b/driver/jvm_tooling.cpp
@@ -23,9 +23,11 @@
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/str_split.h"
+#include "coverage_tracker.h"
 #include "gflags/gflags.h"
 #include "glog/logging.h"
 #include "libfuzzer_callbacks.h"
+#include "signal_handler.h"
 #include "tools/cpp/runfiles/runfiles.h"
 #include "utils.h"
 
@@ -76,7 +78,33 @@ DEFINE_string(
     dump_classes_dir, "",
     "path to a directory in which Jazzer should dump the instrumented classes");
 
-DECLARE_bool(hooks);
+DEFINE_bool(hooks, true,
+            "Use JVM hooks to provide coverage information to the fuzzer. The "
+            "fuzzer uses the coverage information to perform smarter input "
+            "selection and mutation. If set to false no "
+            "coverage information will be processed. This can be useful for "
+            "running a regression test on non-instrumented bytecode.");
+
+// Called by the agent when
+// com.code_intelligence.jazzer.instrumentor.ClassInstrumentor is initialized.
+// This only happens when FLAGS_hooks is true.
+extern "C" JNIEXPORT jint JNICALL JNI_OnLoad_jazzer_initialize(JavaVM *vm,
+                                                               void *) {
+  if (!FLAGS_hooks) {
+    LOG(ERROR) << "JNI_OnLoad_jazzer_initialize called with --nohooks";
+    exit(1);
+  }
+  JNIEnv *env = nullptr;
+  jint result = vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION_1_8);
+  if (result != JNI_OK) {
+    LOG(FATAL) << "Failed to get JNI environment";
+    exit(1);
+  }
+  jazzer::registerFuzzerCallbacks(*env);
+  jazzer::CoverageTracker::Setup(*env);
+  jazzer::SignalHandler::Setup(*env);
+  return JNI_VERSION_1_8;
+}
 
 namespace {
 constexpr auto kAgentBazelRunfilesPath = "jazzer/agent/jazzer_agent_deploy.jar";
@@ -203,7 +231,7 @@ JVM::JVM(const std::string &executable_path) {
         .optionString = const_cast<char *>(agent_jvm_arg.c_str())});
   }
 
-  JavaVMInitArgs jvm_init_args = {.version = JNI_VERSION_1_6,
+  JavaVMInitArgs jvm_init_args = {.version = JNI_VERSION_1_8,
                                   .nOptions = (int)options.size(),
                                   .options = options.data(),
                                   .ignoreUnrecognized = JNI_FALSE};
diff --git a/driver/jvm_tooling_test.cpp b/driver/jvm_tooling_test.cpp
index 33bbc5ca..7069acd9 100644
--- a/driver/jvm_tooling_test.cpp
+++ b/driver/jvm_tooling_test.cpp
@@ -45,7 +45,6 @@ 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); }
diff --git a/driver/libfuzzer_driver.cpp b/driver/libfuzzer_driver.cpp
index 4dfb53cd..91929704 100644
--- a/driver/libfuzzer_driver.cpp
+++ b/driver/libfuzzer_driver.cpp
@@ -34,13 +34,6 @@
 
 using namespace std::string_literals;
 
-DEFINE_bool(hooks, true,
-            "Use JVM hooks to provide coverage information to the fuzzer. The "
-            "fuzzer uses the coverage information to perform smarter input "
-            "selection and mutation. If set to false no "
-            "coverage information will be processed. This can be useful for "
-            "running a regression test on non-instrumented bytecode.");
-
 // Defined by glog
 DECLARE_bool(log_prefix);
 
@@ -176,13 +169,6 @@ AbstractLibfuzzerDriver::AbstractLibfuzzerDriver(
 
 void AbstractLibfuzzerDriver::initJvm(const std::string &executable_path) {
   jvm_ = std::make_unique<jazzer::JVM>(executable_path);
-  if (FLAGS_hooks) {
-    jazzer::registerFuzzerCallbacks(jvm_->GetEnv());
-    CoverageTracker::Setup(jvm_->GetEnv());
-    // SignalHandler registers its own native methods
-    signal_handler_ = std::make_unique<jazzer::SignalHandler>(*jvm_);
-    signal_handler_->SetupSignalHandlers();
-  }
 }
 
 LibfuzzerDriver::LibfuzzerDriver(int *argc, char ***argv)
diff --git a/driver/libfuzzer_driver.h b/driver/libfuzzer_driver.h
index f3dd9a49..557277a5 100644
--- a/driver/libfuzzer_driver.h
+++ b/driver/libfuzzer_driver.h
@@ -20,7 +20,6 @@
 #include <string>
 
 #include "absl/strings/match.h"
-#include "coverage_tracker.h"
 #include "fuzz_target_runner.h"
 #include "fuzzed_data_provider.h"
 #include "jvm_tooling.h"
@@ -50,9 +49,6 @@ class AbstractLibfuzzerDriver {
   std::unique_ptr<jazzer::JVM> jvm_;
 
  private:
-  // handles clearing and reading of the coverage map
-  std::unique_ptr<jazzer::CoverageTracker> coverage_tracker_;
-
   // forwards signals caught while the JVM is running
   std::unique_ptr<jazzer::SignalHandler> signal_handler_;
 
diff --git a/driver/signal_handler.cpp b/driver/signal_handler.cpp
index 4dd67d97..05e5953a 100644
--- a/driver/signal_handler.cpp
+++ b/driver/signal_handler.cpp
@@ -36,25 +36,30 @@ void JNICALL handleInterrupt(JNIEnv, jclass) {
 }
 
 namespace jazzer {
-SignalHandler::SignalHandler(JVM &jvm)
-    : jvm_(jvm),
-      jclass_(jvm.FindClass(kSignalHandlerClass)),
-      setup_signal_handlers_method_(
-          jvm.GetStaticMethodID(jclass_, "setupSignalHandlers", "()V")) {}
-
-void SignalHandler::SetupSignalHandlers() {
-  JNINativeMethod signal_handler[]{
+void SignalHandler::Setup(JNIEnv &env) {
+  jclass signal_handler_class = env.FindClass(kSignalHandlerClass);
+  if (env.ExceptionCheck()) {
+    env.ExceptionDescribe();
+    throw std::runtime_error("could not find signal handler class");
+  }
+  JNINativeMethod signal_handler_methods[]{
       {(char *)"handleInterrupt", (char *)"()V", (void *)&handleInterrupt},
   };
-  jvm_.GetEnv().RegisterNatives(jclass_, signal_handler, 1);
-  if (jvm_.GetEnv().ExceptionCheck()) {
-    jvm_.GetEnv().ExceptionDescribe();
+  env.RegisterNatives(signal_handler_class, signal_handler_methods, 1);
+  if (env.ExceptionCheck()) {
+    env.ExceptionDescribe();
     throw std::runtime_error(
         "could not register native callbacks 'handleInterrupt'");
   }
-  jvm_.GetEnv().CallStaticVoidMethod(jclass_, setup_signal_handlers_method_);
-  if (jvm_.GetEnv().ExceptionCheck()) {
-    jvm_.GetEnv().ExceptionDescribe();
+  jmethodID setup_signal_handlers_method_ =
+      env.GetStaticMethodID(signal_handler_class, "setupSignalHandlers", "()V");
+  if (env.ExceptionCheck()) {
+    env.ExceptionDescribe();
+    throw std::runtime_error("could not find setupSignalHandlers method");
+  }
+  env.CallStaticVoidMethod(signal_handler_class, setup_signal_handlers_method_);
+  if (env.ExceptionCheck()) {
+    env.ExceptionDescribe();
     throw std::runtime_error("failed to set up signal handlers");
   }
 }
diff --git a/driver/signal_handler.h b/driver/signal_handler.h
index 84bd380d..d0d17121 100644
--- a/driver/signal_handler.h
+++ b/driver/signal_handler.h
@@ -18,23 +18,12 @@
 
 #include <jni.h>
 
-#include "jvm_tooling.h"
-
 namespace jazzer {
 // SignalHandler registers handlers for signals (e.g. SIGINT) in Java and
 // notifies the driver via native callbacks when the handlers fire.
 class SignalHandler {
- private:
-  const JVM &jvm_;
-  jclass jclass_;
-  jmethodID setup_signal_handlers_method_;
-
  public:
-  // Throws std::runtime_error if the corresponding Java class and method cannot
-  // be found.
-  explicit SignalHandler(JVM &);
-
   // Set up handlers for signal in Java.
-  void SetupSignalHandlers();
+  static void Setup(JNIEnv &env);
 };
 }  // namespace jazzer