1//===- FuzzerMerge.cpp - merging corpora ----------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8// Merging corpora.
9//===----------------------------------------------------------------------===//
10
11#include "FuzzerCommand.h"
12#include "FuzzerMerge.h"
13#include "FuzzerIO.h"
14#include "FuzzerInternal.h"
15#include "FuzzerTracePC.h"
16#include "FuzzerUtil.h"
17
18#include <fstream>
19#include <iterator>
20#include <set>
21#include <sstream>
22#include <unordered_set>
23
24namespace fuzzer {
25
26bool Merger::Parse(const std::string &Str, bool ParseCoverage) {
27  std::istringstream SS(Str);
28  return Parse(SS, ParseCoverage);
29}
30
31void Merger::ParseOrExit(std::istream &IS, bool ParseCoverage) {
32  if (!Parse(IS, ParseCoverage)) {
33    Printf("MERGE: failed to parse the control file (unexpected error)\n");
34    exit(1);
35  }
36}
37
38// The control file example:
39//
40// 3 # The number of inputs
41// 1 # The number of inputs in the first corpus, <= the previous number
42// file0
43// file1
44// file2  # One file name per line.
45// STARTED 0 123  # FileID, file size
46// FT 0 1 4 6 8  # FileID COV1 COV2 ...
47// COV 0 7 8 9 # FileID COV1 COV1
48// STARTED 1 456  # If FT is missing, the input crashed while processing.
49// STARTED 2 567
50// FT 2 8 9
51// COV 2 11 12
52bool Merger::Parse(std::istream &IS, bool ParseCoverage) {
53  LastFailure.clear();
54  std::string Line;
55
56  // Parse NumFiles.
57  if (!std::getline(IS, Line, '\n')) return false;
58  std::istringstream L1(Line);
59  size_t NumFiles = 0;
60  L1 >> NumFiles;
61  if (NumFiles == 0 || NumFiles > 10000000) return false;
62
63  // Parse NumFilesInFirstCorpus.
64  if (!std::getline(IS, Line, '\n')) return false;
65  std::istringstream L2(Line);
66  NumFilesInFirstCorpus = NumFiles + 1;
67  L2 >> NumFilesInFirstCorpus;
68  if (NumFilesInFirstCorpus > NumFiles) return false;
69
70  // Parse file names.
71  Files.resize(NumFiles);
72  for (size_t i = 0; i < NumFiles; i++)
73    if (!std::getline(IS, Files[i].Name, '\n'))
74      return false;
75
76  // Parse STARTED, FT, and COV lines.
77  size_t ExpectedStartMarker = 0;
78  const size_t kInvalidStartMarker = -1;
79  size_t LastSeenStartMarker = kInvalidStartMarker;
80  Vector<uint32_t> TmpFeatures;
81  Set<uint32_t> PCs;
82  while (std::getline(IS, Line, '\n')) {
83    std::istringstream ISS1(Line);
84    std::string Marker;
85    size_t N;
86    ISS1 >> Marker;
87    ISS1 >> N;
88    if (Marker == "STARTED") {
89      // STARTED FILE_ID FILE_SIZE
90      if (ExpectedStartMarker != N)
91        return false;
92      ISS1 >> Files[ExpectedStartMarker].Size;
93      LastSeenStartMarker = ExpectedStartMarker;
94      assert(ExpectedStartMarker < Files.size());
95      ExpectedStartMarker++;
96    } else if (Marker == "FT") {
97      // FT FILE_ID COV1 COV2 COV3 ...
98      size_t CurrentFileIdx = N;
99      if (CurrentFileIdx != LastSeenStartMarker)
100        return false;
101      LastSeenStartMarker = kInvalidStartMarker;
102      if (ParseCoverage) {
103        TmpFeatures.clear();  // use a vector from outer scope to avoid resizes.
104        while (ISS1 >> N)
105          TmpFeatures.push_back(N);
106        std::sort(TmpFeatures.begin(), TmpFeatures.end());
107        Files[CurrentFileIdx].Features = TmpFeatures;
108      }
109    } else if (Marker == "COV") {
110      size_t CurrentFileIdx = N;
111      if (ParseCoverage)
112        while (ISS1 >> N)
113          if (PCs.insert(N).second)
114            Files[CurrentFileIdx].Cov.push_back(N);
115    } else {
116      return false;
117    }
118  }
119  if (LastSeenStartMarker != kInvalidStartMarker)
120    LastFailure = Files[LastSeenStartMarker].Name;
121
122  FirstNotProcessedFile = ExpectedStartMarker;
123  return true;
124}
125
126size_t Merger::ApproximateMemoryConsumption() const  {
127  size_t Res = 0;
128  for (const auto &F: Files)
129    Res += sizeof(F) + F.Features.size() * sizeof(F.Features[0]);
130  return Res;
131}
132
133// Decides which files need to be merged (add those to NewFiles).
134// Returns the number of new features added.
135size_t Merger::Merge(const Set<uint32_t> &InitialFeatures,
136                     Set<uint32_t> *NewFeatures,
137                     const Set<uint32_t> &InitialCov, Set<uint32_t> *NewCov,
138                     Vector<std::string> *NewFiles) {
139  NewFiles->clear();
140  assert(NumFilesInFirstCorpus <= Files.size());
141  Set<uint32_t> AllFeatures = InitialFeatures;
142
143  // What features are in the initial corpus?
144  for (size_t i = 0; i < NumFilesInFirstCorpus; i++) {
145    auto &Cur = Files[i].Features;
146    AllFeatures.insert(Cur.begin(), Cur.end());
147  }
148  // Remove all features that we already know from all other inputs.
149  for (size_t i = NumFilesInFirstCorpus; i < Files.size(); i++) {
150    auto &Cur = Files[i].Features;
151    Vector<uint32_t> Tmp;
152    std::set_difference(Cur.begin(), Cur.end(), AllFeatures.begin(),
153                        AllFeatures.end(), std::inserter(Tmp, Tmp.begin()));
154    Cur.swap(Tmp);
155  }
156
157  // Sort. Give preference to
158  //   * smaller files
159  //   * files with more features.
160  std::sort(Files.begin() + NumFilesInFirstCorpus, Files.end(),
161            [&](const MergeFileInfo &a, const MergeFileInfo &b) -> bool {
162              if (a.Size != b.Size)
163                return a.Size < b.Size;
164              return a.Features.size() > b.Features.size();
165            });
166
167  // One greedy pass: add the file's features to AllFeatures.
168  // If new features were added, add this file to NewFiles.
169  for (size_t i = NumFilesInFirstCorpus; i < Files.size(); i++) {
170    auto &Cur = Files[i].Features;
171    // Printf("%s -> sz %zd ft %zd\n", Files[i].Name.c_str(),
172    //       Files[i].Size, Cur.size());
173    bool FoundNewFeatures = false;
174    for (auto Fe: Cur) {
175      if (AllFeatures.insert(Fe).second) {
176        FoundNewFeatures = true;
177        NewFeatures->insert(Fe);
178      }
179    }
180    if (FoundNewFeatures)
181      NewFiles->push_back(Files[i].Name);
182    for (auto Cov : Files[i].Cov)
183      if (InitialCov.find(Cov) == InitialCov.end())
184        NewCov->insert(Cov);
185  }
186  return NewFeatures->size();
187}
188
189Set<uint32_t> Merger::AllFeatures() const {
190  Set<uint32_t> S;
191  for (auto &File : Files)
192    S.insert(File.Features.begin(), File.Features.end());
193  return S;
194}
195
196// Inner process. May crash if the target crashes.
197void Fuzzer::CrashResistantMergeInternalStep(const std::string &CFPath) {
198  Printf("MERGE-INNER: using the control file '%s'\n", CFPath.c_str());
199  Merger M;
200  std::ifstream IF(CFPath);
201  M.ParseOrExit(IF, false);
202  IF.close();
203  if (!M.LastFailure.empty())
204    Printf("MERGE-INNER: '%s' caused a failure at the previous merge step\n",
205           M.LastFailure.c_str());
206
207  Printf("MERGE-INNER: %zd total files;"
208         " %zd processed earlier; will process %zd files now\n",
209         M.Files.size(), M.FirstNotProcessedFile,
210         M.Files.size() - M.FirstNotProcessedFile);
211
212  std::ofstream OF(CFPath, std::ofstream::out | std::ofstream::app);
213  Set<size_t> AllFeatures;
214  auto PrintStatsWrapper = [this, &AllFeatures](const char* Where) {
215    this->PrintStats(Where, "\n", 0, AllFeatures.size());
216  };
217  Set<const TracePC::PCTableEntry *> AllPCs;
218  for (size_t i = M.FirstNotProcessedFile; i < M.Files.size(); i++) {
219    Fuzzer::MaybeExitGracefully();
220    auto U = FileToVector(M.Files[i].Name);
221    if (U.size() > MaxInputLen) {
222      U.resize(MaxInputLen);
223      U.shrink_to_fit();
224    }
225
226    // Write the pre-run marker.
227    OF << "STARTED " << i << " " << U.size() << "\n";
228    OF.flush();  // Flush is important since Command::Execute may crash.
229    // Run.
230    TPC.ResetMaps();
231    ExecuteCallback(U.data(), U.size());
232    // Collect coverage. We are iterating over the files in this order:
233    // * First, files in the initial corpus ordered by size, smallest first.
234    // * Then, all other files, smallest first.
235    // So it makes no sense to record all features for all files, instead we
236    // only record features that were not seen before.
237    Set<size_t> UniqFeatures;
238    TPC.CollectFeatures([&](size_t Feature) {
239      if (AllFeatures.insert(Feature).second)
240        UniqFeatures.insert(Feature);
241    });
242    TPC.UpdateObservedPCs();
243    // Show stats.
244    if (!(TotalNumberOfRuns & (TotalNumberOfRuns - 1)))
245      PrintStatsWrapper("pulse ");
246    if (TotalNumberOfRuns == M.NumFilesInFirstCorpus)
247      PrintStatsWrapper("LOADED");
248    // Write the post-run marker and the coverage.
249    OF << "FT " << i;
250    for (size_t F : UniqFeatures)
251      OF << " " << F;
252    OF << "\n";
253    OF << "COV " << i;
254    TPC.ForEachObservedPC([&](const TracePC::PCTableEntry *TE) {
255      if (AllPCs.insert(TE).second)
256        OF << " " << TPC.PCTableEntryIdx(TE);
257    });
258    OF << "\n";
259    OF.flush();
260  }
261  PrintStatsWrapper("DONE  ");
262}
263
264static size_t WriteNewControlFile(const std::string &CFPath,
265                                  const Vector<SizedFile> &OldCorpus,
266                                  const Vector<SizedFile> &NewCorpus,
267                                  const Vector<MergeFileInfo> &KnownFiles) {
268  std::unordered_set<std::string> FilesToSkip;
269  for (auto &SF: KnownFiles)
270    FilesToSkip.insert(SF.Name);
271
272  Vector<std::string> FilesToUse;
273  auto MaybeUseFile = [=, &FilesToUse](std::string Name) {
274    if (FilesToSkip.find(Name) == FilesToSkip.end())
275      FilesToUse.push_back(Name);
276  };
277  for (auto &SF: OldCorpus)
278    MaybeUseFile(SF.File);
279  auto FilesToUseFromOldCorpus = FilesToUse.size();
280  for (auto &SF: NewCorpus)
281    MaybeUseFile(SF.File);
282
283  RemoveFile(CFPath);
284  std::ofstream ControlFile(CFPath);
285  ControlFile << FilesToUse.size() << "\n";
286  ControlFile << FilesToUseFromOldCorpus << "\n";
287  for (auto &FN: FilesToUse)
288    ControlFile << FN << "\n";
289
290  if (!ControlFile) {
291    Printf("MERGE-OUTER: failed to write to the control file: %s\n",
292           CFPath.c_str());
293    exit(1);
294  }
295
296  return FilesToUse.size();
297}
298
299// Outer process. Does not call the target code and thus should not fail.
300void CrashResistantMerge(const Vector<std::string> &Args,
301                         const Vector<SizedFile> &OldCorpus,
302                         const Vector<SizedFile> &NewCorpus,
303                         Vector<std::string> *NewFiles,
304                         const Set<uint32_t> &InitialFeatures,
305                         Set<uint32_t> *NewFeatures,
306                         const Set<uint32_t> &InitialCov,
307                         Set<uint32_t> *NewCov,
308                         const std::string &CFPath,
309                         bool V /*Verbose*/) {
310  if (NewCorpus.empty() && OldCorpus.empty()) return;  // Nothing to merge.
311  size_t NumAttempts = 0;
312  Vector<MergeFileInfo> KnownFiles;
313  if (FileSize(CFPath)) {
314    VPrintf(V, "MERGE-OUTER: non-empty control file provided: '%s'\n",
315           CFPath.c_str());
316    Merger M;
317    std::ifstream IF(CFPath);
318    if (M.Parse(IF, /*ParseCoverage=*/true)) {
319      VPrintf(V, "MERGE-OUTER: control file ok, %zd files total,"
320             " first not processed file %zd\n",
321             M.Files.size(), M.FirstNotProcessedFile);
322      if (!M.LastFailure.empty())
323        VPrintf(V, "MERGE-OUTER: '%s' will be skipped as unlucky "
324               "(merge has stumbled on it the last time)\n",
325               M.LastFailure.c_str());
326      if (M.FirstNotProcessedFile >= M.Files.size()) {
327        // Merge has already been completed with the given merge control file.
328        if (M.Files.size() == OldCorpus.size() + NewCorpus.size()) {
329          VPrintf(
330              V,
331              "MERGE-OUTER: nothing to do, merge has been completed before\n");
332          exit(0);
333        }
334
335        // Number of input files likely changed, start merge from scratch, but
336        // reuse coverage information from the given merge control file.
337        VPrintf(
338            V,
339            "MERGE-OUTER: starting merge from scratch, but reusing coverage "
340            "information from the given control file\n");
341        KnownFiles = M.Files;
342      } else {
343        // There is a merge in progress, continue.
344        NumAttempts = M.Files.size() - M.FirstNotProcessedFile;
345      }
346    } else {
347      VPrintf(V, "MERGE-OUTER: bad control file, will overwrite it\n");
348    }
349  }
350
351  if (!NumAttempts) {
352    // The supplied control file is empty or bad, create a fresh one.
353    VPrintf(V, "MERGE-OUTER: "
354            "%zd files, %zd in the initial corpus, %zd processed earlier\n",
355            OldCorpus.size() + NewCorpus.size(), OldCorpus.size(),
356            KnownFiles.size());
357    NumAttempts = WriteNewControlFile(CFPath, OldCorpus, NewCorpus, KnownFiles);
358  }
359
360  // Execute the inner process until it passes.
361  // Every inner process should execute at least one input.
362  Command BaseCmd(Args);
363  BaseCmd.removeFlag("merge");
364  BaseCmd.removeFlag("fork");
365  BaseCmd.removeFlag("collect_data_flow");
366  for (size_t Attempt = 1; Attempt <= NumAttempts; Attempt++) {
367    Fuzzer::MaybeExitGracefully();
368    VPrintf(V, "MERGE-OUTER: attempt %zd\n", Attempt);
369    Command Cmd(BaseCmd);
370    Cmd.addFlag("merge_control_file", CFPath);
371    Cmd.addFlag("merge_inner", "1");
372    if (!V) {
373      Cmd.setOutputFile(getDevNull());
374      Cmd.combineOutAndErr();
375    }
376    auto ExitCode = ExecuteCommand(Cmd);
377    if (!ExitCode) {
378      VPrintf(V, "MERGE-OUTER: succesfull in %zd attempt(s)\n", Attempt);
379      break;
380    }
381  }
382  // Read the control file and do the merge.
383  Merger M;
384  std::ifstream IF(CFPath);
385  IF.seekg(0, IF.end);
386  VPrintf(V, "MERGE-OUTER: the control file has %zd bytes\n",
387          (size_t)IF.tellg());
388  IF.seekg(0, IF.beg);
389  M.ParseOrExit(IF, true);
390  IF.close();
391  VPrintf(V,
392          "MERGE-OUTER: consumed %zdMb (%zdMb rss) to parse the control file\n",
393          M.ApproximateMemoryConsumption() >> 20, GetPeakRSSMb());
394
395  M.Files.insert(M.Files.end(), KnownFiles.begin(), KnownFiles.end());
396  M.Merge(InitialFeatures, NewFeatures, InitialCov, NewCov, NewFiles);
397  VPrintf(V, "MERGE-OUTER: %zd new files with %zd new features added; "
398          "%zd new coverage edges\n",
399         NewFiles->size(), NewFeatures->size(), NewCov->size());
400}
401
402} // namespace fuzzer
403