[libFuzzer] introduce LLVMFuzzerInitialize

llvm-svn: 257980
diff --git a/llvm/lib/Fuzzer/FuzzerInterface.h b/llvm/lib/Fuzzer/FuzzerInterface.h
index 683ce92..b64dff8 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 @@
   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 6078867..44af57d 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 c5af5b05..0d5c113 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 91a61f4..3533d3b 100644
--- a/llvm/lib/Fuzzer/test/CMakeLists.txt
+++ b/llvm/lib/Fuzzer/test/CMakeLists.txt
@@ -18,6 +18,7 @@
   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 0000000..1462747
--- /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 5b98fde..c8a17f5 100644
--- a/llvm/lib/Fuzzer/test/fuzzer.test
+++ b/llvm/lib/Fuzzer/test/fuzzer.test
@@ -37,3 +37,5 @@
 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