7 files changed, 72 insertions, 2 deletions
diff --git a/llvm/docs/LibFuzzer.rst b/llvm/docs/LibFuzzer.rst
index 84adff3616f..8f4163bd895 100644
--- a/llvm/docs/LibFuzzer.rst
+++ b/llvm/docs/LibFuzzer.rst
@@ -336,6 +336,35 @@ User-supplied mutators
 LibFuzzer allows to use custom (user-supplied) mutators,
 see FuzzerInterface.h_
 
+Startup initialization
+----------------------
+If the library being tested needs to be initialized, there are several options.
+
+The simplest way is to have a statically initialized global object::
+
+   static bool Initialized = DoInitialization();
+
+Alternatively, you may define an optional init function and it will receive
+the program arguments that you can read and modify::
+
+   extern "C" int LLVMFuzzerInitialize(int *argc, char ***argv) {
+    ReadAndMaybeModify(argc, argv);
+    return 0;
+   }
+
+Finally, you may use your own ``main()`` and call ``FuzzerDriver``
+from there, see FuzzerInterface.h_.
+
+Try to avoid initialization inside the target function itself as
+it will skew the coverage data. Don't do this::
+
+    extern "C" int LLVMFuzzerTestOneInput(...) {
+      static bool initialized = false;
+      if (!initialized) { 
+         ...
+      }
+    }
+
 Fuzzing components of LLVM
 ==========================
 
diff --git a/llvm/lib/Fuzzer/FuzzerInterface.h b/llvm/lib/Fuzzer/FuzzerInterface.h
index 683ce9206d1..b64dff8fe43 100644
--- a/llvm/lib/Fuzzer/FuzzerInterface.h
+++ b/llvm/lib/Fuzzer/FuzzerInterface.h
@@ -23,7 +23,6 @@
 #include <string>
 
 namespace fuzzer {
-typedef std::vector<uint8_t> Unit;
 
 /// Returns an int 0. Values other than zero are reserved for future.
 typedef int (*UserCallback)(const uint8_t *Data, size_t Size);
@@ -39,6 +38,15 @@ int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
   return 0;
 }
 
+// Optional.
+// Define this only if you need to read/modify argc/argv at startup
+// and you are using libFuzzer's main().
+// Must return 0.
+int LLVMFuzzerInitialize(int *argc, char ***argv) {
+  ReadAndMaybeModify(argc, argv);
+  return 0;
+}
+
 // Implement your own main() or use the one from FuzzerMain.cpp.
 int main(int argc, char **argv) {
   InitializeMeIfNeeded();
diff --git a/llvm/lib/Fuzzer/FuzzerInternal.h b/llvm/lib/Fuzzer/FuzzerInternal.h
index 60788676611..44af57d5312 100644
--- a/llvm/lib/Fuzzer/FuzzerInternal.h
+++ b/llvm/lib/Fuzzer/FuzzerInternal.h
@@ -25,6 +25,7 @@
 
 namespace fuzzer {
 using namespace std::chrono;
+typedef std::vector<uint8_t> Unit;
 
 std::string FileToString(const std::string &Path);
 Unit FileToVector(const std::string &Path);
diff --git a/llvm/lib/Fuzzer/FuzzerMain.cpp b/llvm/lib/Fuzzer/FuzzerMain.cpp
index c5af5b05909..0d5c1132ea5 100644
--- a/llvm/lib/Fuzzer/FuzzerMain.cpp
+++ b/llvm/lib/Fuzzer/FuzzerMain.cpp
@@ -12,9 +12,15 @@
 #include "FuzzerInterface.h"
 #include "FuzzerInternal.h"
 
+extern "C" {
 // This function should be defined by the user.
-extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size);
+int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size);
+// This function may optionally be defined by the user.
+__attribute__((weak)) int LLVMFuzzerInitialize(int *argc, char ***argv);
+}  // extern "C"
 
 int main(int argc, char **argv) {
+  if (LLVMFuzzerInitialize)
+    LLVMFuzzerInitialize(&argc, &argv);
   return fuzzer::FuzzerDriver(argc, argv, LLVMFuzzerTestOneInput);
 }
diff --git a/llvm/lib/Fuzzer/test/CMakeLists.txt b/llvm/lib/Fuzzer/test/CMakeLists.txt
index 91a61f43375..3533d3b2a41 100644
--- a/llvm/lib/Fuzzer/test/CMakeLists.txt
+++ b/llvm/lib/Fuzzer/test/CMakeLists.txt
@@ -18,6 +18,7 @@ set(Tests
   CounterTest
   FourIndependentBranchesTest
   FullCoverageSetTest
+  InitializeTest
   MemcmpTest
   NullDerefTest
   RepeatedMemcmp
diff --git a/llvm/lib/Fuzzer/test/InitializeTest.cpp b/llvm/lib/Fuzzer/test/InitializeTest.cpp
new file mode 100644
index 00000000000..14627479708
--- /dev/null
+++ b/llvm/lib/Fuzzer/test/InitializeTest.cpp
@@ -0,0 +1,23 @@
+// Make sure LLVMFuzzerInitialize is called.
+#include <assert.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+static char *argv0;
+
+extern "C" int LLVMFuzzerInitialize(int *argc, char ***argv) {
+  assert(argc > 0);
+  argv0 = **argv;
+  return 0;
+}
+
+extern "C" int LLVMFuzzerTestOneInput(const uint8_t *Data, size_t Size) {
+  if (strncmp(reinterpret_cast<const char*>(Data), argv0, Size)) {
+    fprintf(stderr, "BINGO\n");
+    exit(1);
+  }
+  return 0;
+}
diff --git a/llvm/lib/Fuzzer/test/fuzzer.test b/llvm/lib/Fuzzer/test/fuzzer.test
index 5b98fde6815..c8a17f53880 100644
--- a/llvm/lib/Fuzzer/test/fuzzer.test
+++ b/llvm/lib/Fuzzer/test/fuzzer.test
@@ -37,3 +37,5 @@ PCS:BINGO
 RUN: not LLVMFuzzer-BufferOverflowOnInput 2>&1 | FileCheck %s --check-prefix=OOB
 OOB: AddressSanitizer: heap-buffer-overflow
 OOB: is located 0 bytes to the right of 3-byte region
+
+RUN: not LLVMFuzzer-InitializeTest 2>&1 | FileCheck %s