[fuzzer] split main() into FuzzerDriver() that takes a callback as a parameter and...
[oota-llvm.git] / lib / Fuzzer / FuzzerDriver.cpp
1 //===- FuzzerDriver.cpp - FuzzerDriver function and flags -----------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 // FuzzerDriver and flag parsing.
10 //===----------------------------------------------------------------------===//
11
12 #include "FuzzerInterface.h"
13 #include "FuzzerInternal.h"
14
15 #include <cstring>
16 #include <unistd.h>
17 #include <iostream>
18 #include <thread>
19 #include <atomic>
20 #include <mutex>
21
22 namespace fuzzer {
23
24 // Program arguments.
25 struct FlagDescription {
26   const char *Name;
27   const char *Description;
28   int   Default;
29   int   *Flag;
30 };
31
32 struct {
33 #define FUZZER_FLAG(Type, Name, Default, Description) Type Name;
34 #include "FuzzerFlags.def"
35 #undef FUZZER_FLAG
36 } Flags;
37
38 static FlagDescription FlagDescriptions [] {
39 #define FUZZER_FLAG(Type, Name, Default, Description) {#Name, Description, Default, &Flags.Name},
40 #include "FuzzerFlags.def"
41 #undef FUZZER_FLAG
42 };
43
44 static const size_t kNumFlags =
45     sizeof(FlagDescriptions) / sizeof(FlagDescriptions[0]);
46
47 static std::vector<std::string> inputs;
48 static const char *ProgName;
49
50 static void PrintHelp() {
51   std::cerr << "Usage: " << ProgName
52             << " [-flag1=val1 [-flag2=val2 ...] ] [dir1 [dir2 ...] ]\n";
53   std::cerr << "\nFlags: (strictly in form -flag=value)\n";
54   size_t MaxFlagLen = 0;
55   for (size_t F = 0; F < kNumFlags; F++)
56     MaxFlagLen = std::max(strlen(FlagDescriptions[F].Name), MaxFlagLen);
57
58   for (size_t F = 0; F < kNumFlags; F++) {
59     const auto &D = FlagDescriptions[F];
60     std::cerr << "  " << D.Name;
61     for (size_t i = 0, n = MaxFlagLen - strlen(D.Name); i < n; i++)
62       std::cerr << " ";
63     std::cerr << "\t";
64     std::cerr << D.Default << "\t" << D.Description << "\n";
65   }
66 }
67
68 static const char *FlagValue(const char *Param, const char *Name) {
69   size_t Len = strlen(Name);
70   if (Param[0] == '-' && strstr(Param + 1, Name) == Param + 1 &&
71       Param[Len + 1] == '=')
72       return &Param[Len + 2];
73   return nullptr;
74 }
75
76 static bool ParseOneFlag(const char *Param) {
77   if (Param[0] != '-') return false;
78   for (size_t F = 0; F < kNumFlags; F++) {
79     const char *Name = FlagDescriptions[F].Name;
80     const char *Str = FlagValue(Param, Name);
81     if (Str)  {
82       int Val = std::stol(Str);
83       *FlagDescriptions[F].Flag = Val;
84       if (Flags.verbosity >= 2)
85         std::cerr << "Flag: " << Name << " " << Val << "\n";
86       return true;
87     }
88   }
89   PrintHelp();
90   exit(1);
91 }
92
93 // We don't use any library to minimize dependencies.
94 static void ParseFlags(int argc, char **argv) {
95   for (size_t F = 0; F < kNumFlags; F++)
96     *FlagDescriptions[F].Flag = FlagDescriptions[F].Default;
97   for (int A = 1; A < argc; A++) {
98     if (ParseOneFlag(argv[A])) continue;
99     inputs.push_back(argv[A]);
100   }
101 }
102
103 static void WorkerThread(const std::string &Cmd, std::atomic<int> *Counter,
104                         int NumJobs, std::atomic<bool> *HasErrors) {
105   static std::mutex CerrMutex;
106   while (true) {
107     int C = (*Counter)++;
108     if (C >= NumJobs) break;
109     std::string Log = "fuzz-" + std::to_string(C) + ".log";
110     std::string ToRun = Cmd + " > " + Log + " 2>&1\n";
111     if (Flags.verbosity)
112       std::cerr << ToRun;
113     int ExitCode = system(ToRun.c_str());
114     if (ExitCode != 0)
115       *HasErrors = true;
116     std::lock_guard<std::mutex> Lock(CerrMutex);
117     std::cerr << "================== Job " << C
118               << " exited with exit code " << ExitCode
119               << " =================\n";
120     fuzzer::CopyFileToErr(Log);
121   }
122 }
123
124 static int RunInMultipleProcesses(int argc, char **argv, int NumWorkers,
125                                   int NumJobs) {
126   std::atomic<int> Counter(0);
127   std::atomic<bool> HasErrors(false);
128   std::string Cmd;
129   for (int i = 0; i < argc; i++) {
130     if (FlagValue(argv[i], "jobs") || FlagValue(argv[i], "workers")) continue;
131     Cmd += argv[i];
132     Cmd += " ";
133   }
134   std::vector<std::thread> V;
135   for (int i = 0; i < NumWorkers; i++)
136     V.push_back(std::thread(WorkerThread, Cmd, &Counter, NumJobs, &HasErrors));
137   for (auto &T : V)
138     T.join();
139   return HasErrors ? 1 : 0;
140 }
141
142 int FuzzerDriver(int argc, char **argv, UserCallback Callback) {
143   using namespace fuzzer;
144
145   ProgName = argv[0];
146   ParseFlags(argc, argv);
147   if (Flags.help) {
148     PrintHelp();
149     return 0;
150   }
151
152   if (Flags.workers > 0 && Flags.jobs > 0)
153     return RunInMultipleProcesses(argc, argv, Flags.workers, Flags.jobs);
154
155   Fuzzer::FuzzingOptions Options;
156   Options.Verbosity = Flags.verbosity;
157   Options.MaxLen = Flags.max_len;
158   Options.DoCrossOver = Flags.cross_over;
159   Options.MutateDepth = Flags.mutate_depth;
160   Options.ExitOnFirst = Flags.exit_on_first;
161   Options.UseFullCoverageSet = Flags.use_full_coverage_set;
162   Options.PreferSmallDuringInitialShuffle =
163       Flags.prefer_small_during_initial_shuffle;
164   if (Flags.runs >= 0)
165     Options.MaxNumberOfRuns = Flags.runs;
166   if (!inputs.empty())
167     Options.OutputCorpus = inputs[0];
168   Fuzzer F(Callback, Options);
169
170   unsigned seed = Flags.seed;
171   // Initialize seed.
172   if (seed == 0)
173     seed = time(0) * 10000 + getpid();
174   if (Flags.verbosity)
175     std::cerr << "Seed: " << seed << "\n";
176   srand(seed);
177
178   // Timer
179   if (Flags.timeout > 0)
180     SetTimer(Flags.timeout);
181
182   for (auto &inp : inputs)
183     F.ReadDir(inp);
184
185   if (F.CorpusSize() == 0)
186     F.AddToCorpus(Unit());  // Can't fuzz empty corpus, so add an empty input.
187   F.ShuffleAndMinimize();
188   if (Flags.save_minimized_corpus)
189     F.SaveCorpus();
190   F.Loop(Flags.iterations < 0 ? INT_MAX : Flags.iterations);
191   if (Flags.verbosity)
192     std::cerr << "Done " << F.getTotalNumberOfRuns()
193               << " runs in " << F.secondsSinceProcessStartUp()
194               << " seconds\n";
195   return 0;
196 }
197
198 }  // namespace fuzzer