1 //===- FuzzerDataFlowTrace.h - Internal header for the Fuzzer ---*- C++ -* ===//
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 // fuzzer::DataFlowTrace; reads and handles a data-flow trace.
9 //
10 // A data flow trace is generated by e.g. dataflow/DataFlow.cpp
11 // and is stored on disk in a separate directory.
12 //
13 // The trace dir contains a file 'functions.txt' which lists function names,
14 // oner per line, e.g.
15 // ==> functions.txt <==
16 // Func2
17 // LLVMFuzzerTestOneInput
18 // Func1
19 //
20 // All other files in the dir are the traces, see dataflow/DataFlow.cpp.
21 // The name of the file is sha1 of the input used to generate the trace.
22 //
23 // Current status:
24 //   the data is parsed and the summary is printed, but the data is not yet
25 //   used in any other way.
26 //===----------------------------------------------------------------------===//
27 
28 #ifndef LLVM_FUZZER_DATA_FLOW_TRACE
29 #define LLVM_FUZZER_DATA_FLOW_TRACE
30 
31 #include "FuzzerDefs.h"
32 #include "FuzzerIO.h"
33 
34 #include <unordered_map>
35 #include <unordered_set>
36 #include <vector>
37 #include <string>
38 
39 namespace fuzzer {
40 
41 int CollectDataFlow(const std::string &DFTBinary, const std::string &DirPath,
42                     const std::vector<SizedFile> &CorporaFiles);
43 
44 class BlockCoverage {
45 public:
46   // These functions guarantee no CoverageVector is longer than UINT32_MAX.
47   bool AppendCoverage(std::istream &IN);
48   bool AppendCoverage(const std::string &S);
49 
50   size_t NumCoveredFunctions() const { return Functions.size(); }
51 
52   uint32_t GetCounter(size_t FunctionId, size_t BasicBlockId) {
53     auto It = Functions.find(FunctionId);
54     if (It == Functions.end())
55       return 0;
56     const auto &Counters = It->second;
57     if (BasicBlockId < Counters.size())
58       return Counters[BasicBlockId];
59     return 0;
60   }
61 
62   uint32_t GetNumberOfBlocks(size_t FunctionId) {
63     auto It = Functions.find(FunctionId);
64     if (It == Functions.end()) return 0;
65     const auto &Counters = It->second;
66     return static_cast<uint32_t>(Counters.size());
67   }
68 
69   uint32_t GetNumberOfCoveredBlocks(size_t FunctionId) {
70     auto It = Functions.find(FunctionId);
71     if (It == Functions.end()) return 0;
72     const auto &Counters = It->second;
73     uint32_t Result = 0;
74     for (auto Cnt: Counters)
75       if (Cnt)
76         Result++;
77     return Result;
78   }
79 
80   std::vector<double> FunctionWeights(size_t NumFunctions) const;
81   void clear() { Functions.clear(); }
82 
83 private:
84   typedef std::vector<uint32_t> CoverageVector;
85 
86   uint32_t NumberOfCoveredBlocks(const CoverageVector &Counters) const {
87     uint32_t Res = 0;
88     for (auto Cnt : Counters)
89       if (Cnt)
90         Res++;
91     return Res;
92   }
93 
94   uint32_t NumberOfUncoveredBlocks(const CoverageVector &Counters) const {
95     return static_cast<uint32_t>(Counters.size()) -
96            NumberOfCoveredBlocks(Counters);
97   }
98 
99   uint32_t SmallestNonZeroCounter(const CoverageVector &Counters) const {
100     assert(!Counters.empty());
101     uint32_t Res = Counters[0];
102     for (auto Cnt : Counters)
103       if (Cnt)
104         Res = Min(Res, Cnt);
105     assert(Res);
106     return Res;
107   }
108 
109   // Function ID => vector of counters.
110   // Each counter represents how many input files trigger the given basic block.
111   std::unordered_map<size_t, CoverageVector> Functions;
112   // Functions that have DFT entry.
113   std::unordered_set<size_t> FunctionsWithDFT;
114 };
115 
116 class DataFlowTrace {
117  public:
118   void ReadCoverage(const std::string &DirPath);
119   bool Init(const std::string &DirPath, std::string *FocusFunction,
120             std::vector<SizedFile> &CorporaFiles, Random &Rand);
121   void Clear() { Traces.clear(); }
122   const std::vector<uint8_t> *Get(const std::string &InputSha1) const {
123     auto It = Traces.find(InputSha1);
124     if (It != Traces.end())
125       return &It->second;
126     return nullptr;
127   }
128 
129  private:
130   // Input's sha1 => DFT for the FocusFunction.
131    std::unordered_map<std::string, std::vector<uint8_t>> Traces;
132    BlockCoverage Coverage;
133    std::unordered_set<std::string> CorporaHashes;
134 };
135 }  // namespace fuzzer
136 
137 #endif // LLVM_FUZZER_DATA_FLOW_TRACE
138