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