[libFuzzer] handle SIGUSR1/SIGUSR2 and try to exit grafully on these signals

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

7 files changed, 34 insertions, 0 deletions

diff --git a/lib/fuzzer/FuzzerDriver.cpp b/lib/fuzzer/FuzzerDriver.cpp
index abc314420..6480fbe34 100644
--- a/lib/fuzzer/FuzzerDriver.cpp
+++ b/lib/fuzzer/FuzzerDriver.cpp
@@ -639,6 +639,8 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
   Options.HandleSegv = Flags.handle_segv;
   Options.HandleTerm = Flags.handle_term;
   Options.HandleXfsz = Flags.handle_xfsz;
+  Options.HandleUsr1 = Flags.handle_usr1;
+  Options.HandleUsr2 = Flags.handle_usr2;
   SetSignalHandler(Options);
 
   std::atexit(Fuzzer::StaticExitCallback);
diff --git a/lib/fuzzer/FuzzerFlags.def b/lib/fuzzer/FuzzerFlags.def
index 9b2115779..d738a422d 100644
--- a/lib/fuzzer/FuzzerFlags.def
+++ b/lib/fuzzer/FuzzerFlags.def
@@ -114,6 +114,8 @@ FUZZER_FLAG_INT(handle_fpe, 1, "If 1, try to intercept SIGFPE.")
 FUZZER_FLAG_INT(handle_int, 1, "If 1, try to intercept SIGINT.")
 FUZZER_FLAG_INT(handle_term, 1, "If 1, try to intercept SIGTERM.")
 FUZZER_FLAG_INT(handle_xfsz, 1, "If 1, try to intercept SIGXFSZ.")
+FUZZER_FLAG_INT(handle_usr1, 1, "If 1, try to intercept SIGUSR1.")
+FUZZER_FLAG_INT(handle_usr2, 1, "If 1, try to intercept SIGUSR2.")
 FUZZER_FLAG_INT(close_fd_mask, 0, "If 1, close stdout at startup; "
     "if 2, close stderr; if 3, close both. "
     "Be careful, this will also close e.g. asan's stderr/stdout.")
diff --git a/lib/fuzzer/FuzzerInternal.h b/lib/fuzzer/FuzzerInternal.h
index 371624408..cb26fe332 100644
--- a/lib/fuzzer/FuzzerInternal.h
+++ b/lib/fuzzer/FuzzerInternal.h
@@ -63,6 +63,7 @@ public:
   static void StaticExitCallback();
   static void StaticInterruptCallback();
   static void StaticFileSizeExceedCallback();
+  static void StaticGracefulExitCallback();
 
   void ExecuteCallback(const uint8_t *Data, size_t Size);
   bool RunOne(const uint8_t *Data, size_t Size, bool MayDeleteFile = false,
@@ -94,6 +95,7 @@ private:
   void AlarmCallback();
   void CrashCallback();
   void ExitCallback();
+  void MaybeExitGracefully();
   void CrashOnOverwrittenData();
   void InterruptCallback();
   void MutateAndTestOne();
@@ -116,6 +118,8 @@ private:
   uint8_t BaseSha1[kSHA1NumBytes];  // Checksum of the base unit.
   bool RunningCB = false;
 
+  bool GracefulExitRequested = false;
+
   size_t TotalNumberOfRuns = 0;
   size_t NumberOfNewUnitsAdded = 0;
 
diff --git a/lib/fuzzer/FuzzerLoop.cpp b/lib/fuzzer/FuzzerLoop.cpp
index f4771e1df..97bc1c62b 100644
--- a/lib/fuzzer/FuzzerLoop.cpp
+++ b/lib/fuzzer/FuzzerLoop.cpp
@@ -216,6 +216,12 @@ void Fuzzer::StaticInterruptCallback() {
   F->InterruptCallback();
 }
 
+void Fuzzer::StaticGracefulExitCallback() {
+  assert(F);
+  F->GracefulExitRequested = true;
+  Printf("INFO: signal received, trying to exit gracefully\n");
+}
+
 void Fuzzer::StaticFileSizeExceedCallback() {
   Printf("==%lu== ERROR: libFuzzer: file size exceeded\n", GetPid());
   exit(1);
@@ -246,6 +252,13 @@ void Fuzzer::ExitCallback() {
   _Exit(Options.ErrorExitCode);
 }
 
+void Fuzzer::MaybeExitGracefully() {
+  if (!GracefulExitRequested) return;
+  Printf("==%lu== INFO: libFuzzer: exiting as requested\n", GetPid());
+  PrintFinalStats();
+  _Exit(0);
+}
+
 void Fuzzer::InterruptCallback() {
   Printf("==%lu== libFuzzer: run interrupted; exiting\n", GetPid());
   PrintFinalStats();
@@ -621,6 +634,7 @@ void Fuzzer::MutateAndTestOne() {
   for (int i = 0; i < Options.MutateDepth; i++) {
     if (TotalNumberOfRuns >= Options.MaxNumberOfRuns)
       break;
+    MaybeExitGracefully();
     size_t NewSize = 0;
     NewSize = MD.Mutate(CurrentUnitData, Size, CurrentMaxMutationLen);
     assert(NewSize > 0 && "Mutator returned empty unit");
diff --git a/lib/fuzzer/FuzzerMerge.cpp b/lib/fuzzer/FuzzerMerge.cpp
index 89b4821d0..59e34cfa0 100644
--- a/lib/fuzzer/FuzzerMerge.cpp
+++ b/lib/fuzzer/FuzzerMerge.cpp
@@ -223,6 +223,7 @@ void Fuzzer::CrashResistantMergeInternalStep(const std::string &CFPath) {
   std::ofstream OF(CFPath, std::ofstream::out | std::ofstream::app);
   Set<size_t> AllFeatures;
   for (size_t i = M.FirstNotProcessedFile; i < M.Files.size(); i++) {
+    MaybeExitGracefully();
     auto U = FileToVector(M.Files[i].Name);
     if (U.size() > MaxInputLen) {
       U.resize(MaxInputLen);
@@ -334,6 +335,7 @@ void Fuzzer::CrashResistantMerge(const Vector<std::string> &Args,
                              CloneArgsWithoutX(Args, "merge"));
   bool Success = false;
   for (size_t Attempt = 1; Attempt <= NumAttempts; Attempt++) {
+    MaybeExitGracefully();
     Printf("MERGE-OUTER: attempt %zd\n", Attempt);
     auto ExitCode =
         ExecuteCommand(BaseCmd.first + " -merge_control_file=" + CFPath +
diff --git a/lib/fuzzer/FuzzerOptions.h b/lib/fuzzer/FuzzerOptions.h
index 73953e154..0a1fb613c 100644
--- a/lib/fuzzer/FuzzerOptions.h
+++ b/lib/fuzzer/FuzzerOptions.h
@@ -65,6 +65,8 @@ struct FuzzingOptions {
   bool HandleSegv = false;
   bool HandleTerm = false;
   bool HandleXfsz = false;
+  bool HandleUsr1 = false;
+  bool HandleUsr2 = false;
 };
 
 }  // namespace fuzzer
diff --git a/lib/fuzzer/FuzzerUtilPosix.cpp b/lib/fuzzer/FuzzerUtilPosix.cpp
index 24c5ccc35..934b7aa98 100644
--- a/lib/fuzzer/FuzzerUtilPosix.cpp
+++ b/lib/fuzzer/FuzzerUtilPosix.cpp
@@ -40,6 +40,10 @@ static void InterruptHandler(int, siginfo_t *, void *) {
   Fuzzer::StaticInterruptCallback();
 }
 
+static void GracefulExitHandler(int, siginfo_t *, void *) {
+  Fuzzer::StaticGracefulExitCallback();
+}
+
 static void FileSizeExceedHandler(int, siginfo_t *, void *) {
   Fuzzer::StaticFileSizeExceedCallback();
 }
@@ -98,6 +102,10 @@ void SetSignalHandler(const FuzzingOptions& Options) {
     SetSigaction(SIGFPE, CrashHandler);
   if (Options.HandleXfsz)
     SetSigaction(SIGXFSZ, FileSizeExceedHandler);
+  if (Options.HandleUsr1)
+    SetSigaction(SIGUSR1, GracefulExitHandler);
+  if (Options.HandleUsr2)
+    SetSigaction(SIGUSR2, GracefulExitHandler);
 }
 
 void SleepSeconds(int Seconds) {