From 23482e1d071ec6d05ddbae8eaffda3221260a35f Mon Sep 17 00:00:00 2001
From: Kostya Serebryany <kcc@google.com>
Date: Thu, 31 Jan 2019 01:40:14 +0000
Subject: [PATCH] [libFuzzer] set libFuzzer's own SEGV handler even one is
 already present, but call that handler from ours (unless we are unprotecting
 lazy counters). Call ProtectLazyCounters later, so that it runs after the
 initialization code in the target.

llvm-svn: 352713
---
 compiler-rt/lib/fuzzer/FuzzerDriver.cpp    |  4 +---
 compiler-rt/lib/fuzzer/FuzzerInternal.h    |  1 -
 compiler-rt/lib/fuzzer/FuzzerLoop.cpp      |  9 ++++-----
 compiler-rt/lib/fuzzer/FuzzerUtilPosix.cpp | 17 +++++++++++++----
 compiler-rt/test/fuzzer/large.test         |  1 +
 5 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/compiler-rt/lib/fuzzer/FuzzerDriver.cpp b/compiler-rt/lib/fuzzer/FuzzerDriver.cpp
index acebd3a7b1e4..0f8389cb3560 100644
--- a/compiler-rt/lib/fuzzer/FuzzerDriver.cpp
+++ b/compiler-rt/lib/fuzzer/FuzzerDriver.cpp
@@ -628,6 +628,7 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
     Options.FocusFunction = Flags.focus_function;
   if (Flags.data_flow_trace)
     Options.DataFlowTrace = Flags.data_flow_trace;
+  Options.LazyCounters = Flags.lazy_counters;
 
   unsigned Seed = Flags.seed;
   // Initialize Seed.
@@ -658,10 +659,7 @@ int FuzzerDriver(int *argc, char ***argv, UserCallback Callback) {
   Options.HandleXfsz = Flags.handle_xfsz;
   Options.HandleUsr1 = Flags.handle_usr1;
   Options.HandleUsr2 = Flags.handle_usr2;
-  Options.LazyCounters = Flags.lazy_counters;
   SetSignalHandler(Options);
-  if (Options.LazyCounters)
-    TPC.ProtectLazyCounters();
 
   std::atexit(Fuzzer::StaticExitCallback);
 
diff --git a/compiler-rt/lib/fuzzer/FuzzerInternal.h b/compiler-rt/lib/fuzzer/FuzzerInternal.h
index 9e3e4bb948c5..9950445bc18f 100644
--- a/compiler-rt/lib/fuzzer/FuzzerInternal.h
+++ b/compiler-rt/lib/fuzzer/FuzzerInternal.h
@@ -59,7 +59,6 @@ public:
   size_t getTotalNumberOfRuns() { return TotalNumberOfRuns; }
 
   static void StaticAlarmCallback();
-  static void StaticSegvSignalCallback(void *Addr);
   static void StaticCrashSignalCallback();
   static void StaticExitCallback();
   static void StaticInterruptCallback();
diff --git a/compiler-rt/lib/fuzzer/FuzzerLoop.cpp b/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
index bf600169c4c6..fb5aa1f11c19 100644
--- a/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
+++ b/compiler-rt/lib/fuzzer/FuzzerLoop.cpp
@@ -205,11 +205,6 @@ void Fuzzer::StaticCrashSignalCallback() {
   F->CrashCallback();
 }
 
-void Fuzzer::StaticSegvSignalCallback(void *Addr) {
-  if (TPC.UnprotectLazyCounters(Addr)) return;
-  StaticCrashSignalCallback();
-}
-
 void Fuzzer::StaticExitCallback() {
   assert(F);
   F->ExitCallback();
@@ -720,6 +715,10 @@ void Fuzzer::ReadAndExecuteSeedCorpora(const Vector<std::string> &CorpusDirs) {
   uint8_t dummy = 0;
   ExecuteCallback(&dummy, 0);
 
+  // Protect lazy counters here, after the once-init code has been executed.
+  if (Options.LazyCounters)
+    TPC.ProtectLazyCounters();
+
   if (SizedFiles.empty()) {
     Printf("INFO: A corpus is not provided, starting from an empty corpus\n");
     Unit U({'\n'}); // Valid ASCII input.
diff --git a/compiler-rt/lib/fuzzer/FuzzerUtilPosix.cpp b/compiler-rt/lib/fuzzer/FuzzerUtilPosix.cpp
index 56b10ffa6685..110785d87413 100644
--- a/compiler-rt/lib/fuzzer/FuzzerUtilPosix.cpp
+++ b/compiler-rt/lib/fuzzer/FuzzerUtilPosix.cpp
@@ -11,6 +11,7 @@
 #if LIBFUZZER_POSIX
 #include "FuzzerIO.h"
 #include "FuzzerInternal.h"
+#include "FuzzerTracePC.h"
 #include <cassert>
 #include <chrono>
 #include <cstring>
@@ -32,9 +33,14 @@ static void AlarmHandler(int, siginfo_t *, void *) {
   Fuzzer::StaticAlarmCallback();
 }
 
-static void SegvHandler(int, siginfo_t *si, void *) {
+static void (*upstream_segv_handler)(int, siginfo_t *, void *);
+
+static void SegvHandler(int sig, siginfo_t *si, void *ucontext) {
   assert(si->si_signo == SIGSEGV);
-  Fuzzer::StaticSegvSignalCallback(si->si_addr);
+  if (TPC.UnprotectLazyCounters(si->si_addr)) return;
+  if (upstream_segv_handler)
+    return upstream_segv_handler(sig, si, ucontext);
+  Fuzzer::StaticCrashSignalCallback();
 }
 
 static void CrashHandler(int, siginfo_t *, void *) {
@@ -61,8 +67,11 @@ static void SetSigaction(int signum,
     exit(1);
   }
   if (sigact.sa_flags & SA_SIGINFO) {
-    if (sigact.sa_sigaction)
-      return;
+    if (sigact.sa_sigaction) {
+      if (signum != SIGSEGV)
+        return;
+      upstream_segv_handler = sigact.sa_sigaction;
+    }
   } else {
     if (sigact.sa_handler != SIG_DFL && sigact.sa_handler != SIG_IGN &&
         sigact.sa_handler != SIG_ERR)
diff --git a/compiler-rt/test/fuzzer/large.test b/compiler-rt/test/fuzzer/large.test
index 91d279bc07d1..99ebbbe4e7da 100644
--- a/compiler-rt/test/fuzzer/large.test
+++ b/compiler-rt/test/fuzzer/large.test
@@ -2,5 +2,6 @@ REQUIRES: linux
 RUN: %cpp_compiler %S/LargeTest.cpp -o %t-LargeTest
 RUN: %run %t-LargeTest -runs=10000
 RUN: %env_asan_opts=handle_segv=0 %run %t-LargeTest -runs=10000 -lazy_counters=1 2>&1 | FileCheck %s
+RUN:                              %run %t-LargeTest -runs=10000 -lazy_counters=1 2>&1 | FileCheck %s
 
 CHECK: pages of counters where protected; libFuzzer's SEGV handler must be installed