1 // Copyright 2017 Google Inc. All rights reserved.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #include "src/libfuzzer/libfuzzer_macro.h"
16 
17 #include <algorithm>
18 #include <memory>
19 #include <vector>
20 
21 #include "src/binary_format.h"
22 #include "src/libfuzzer/libfuzzer_mutator.h"
23 #include "src/text_format.h"
24 
25 namespace protobuf_mutator {
26 namespace libfuzzer {
27 
28 namespace {
29 
30 class InputReader {
31  public:
InputReader(const uint8_t * data,size_t size)32   InputReader(const uint8_t* data, size_t size) : data_(data), size_(size) {}
33   virtual ~InputReader() = default;
34 
35   virtual bool Read(protobuf::Message* message) const = 0;
36 
data() const37   const uint8_t* data() const { return data_; }
size() const38   size_t size() const { return size_; }
39 
40  private:
41   const uint8_t* data_;
42   size_t size_;
43 };
44 
45 class OutputWriter {
46  public:
OutputWriter(uint8_t * data,size_t size)47   OutputWriter(uint8_t* data, size_t size) : data_(data), size_(size) {}
48   virtual ~OutputWriter() = default;
49 
50   virtual size_t Write(const protobuf::Message& message) = 0;
51 
data() const52   uint8_t* data() const { return data_; }
size() const53   size_t size() const { return size_; }
54 
55  private:
56   uint8_t* data_;
57   size_t size_;
58 };
59 
60 class TextInputReader : public InputReader {
61  public:
62   using InputReader::InputReader;
63 
Read(protobuf::Message * message) const64   bool Read(protobuf::Message* message) const override {
65     return ParseTextMessage(data(), size(), message);
66   }
67 };
68 
69 class TextOutputWriter : public OutputWriter {
70  public:
71   using OutputWriter::OutputWriter;
72 
Write(const protobuf::Message & message)73   size_t Write(const protobuf::Message& message) override {
74     return SaveMessageAsText(message, data(), size());
75   }
76 };
77 
78 class BinaryInputReader : public InputReader {
79  public:
80   using InputReader::InputReader;
81 
Read(protobuf::Message * message) const82   bool Read(protobuf::Message* message) const override {
83     return ParseBinaryMessage(data(), size(), message);
84   }
85 };
86 
87 class BinaryOutputWriter : public OutputWriter {
88  public:
89   using OutputWriter::OutputWriter;
90 
Write(const protobuf::Message & message)91   size_t Write(const protobuf::Message& message) override {
92     return SaveMessageAsBinary(message, data(), size());
93   }
94 };
95 
96 class LastMutationCache {
97  public:
Store(const uint8_t * data,size_t size,protobuf::Message * message)98   void Store(const uint8_t* data, size_t size, protobuf::Message* message) {
99     if (!message_) message_.reset(message->New());
100     message->GetReflection()->Swap(message, message_.get());
101     data_.assign(data, data + size);
102   }
103 
LoadIfSame(const uint8_t * data,size_t size,protobuf::Message * message)104   bool LoadIfSame(const uint8_t* data, size_t size,
105                   protobuf::Message* message) {
106     if (!message_ || size != data_.size() ||
107         !std::equal(data_.begin(), data_.end(), data))
108       return false;
109 
110     message->GetReflection()->Swap(message, message_.get());
111     message_.reset();
112     return true;
113   }
114 
115  private:
116   std::vector<uint8_t> data_;
117   std::unique_ptr<protobuf::Message> message_;
118 };
119 
GetCache()120 LastMutationCache* GetCache() {
121   static LastMutationCache cache;
122   return &cache;
123 }
124 
GetMutator()125 Mutator* GetMutator() {
126   static Mutator mutator;
127   return &mutator;
128 }
129 
GetMaxSize(const InputReader & input,const OutputWriter & output,const protobuf::Message & message)130 size_t GetMaxSize(const InputReader& input, const OutputWriter& output,
131                   const protobuf::Message& message) {
132   size_t max_size = message.ByteSizeLong() + output.size();
133   max_size -= std::min(max_size, input.size());
134   return max_size;
135 }
136 
MutateMessage(unsigned int seed,const InputReader & input,OutputWriter * output,protobuf::Message * message)137 size_t MutateMessage(unsigned int seed, const InputReader& input,
138                      OutputWriter* output, protobuf::Message* message) {
139   GetMutator()->Seed(seed);
140   input.Read(message);
141   size_t max_size = GetMaxSize(input, *output, *message);
142   GetMutator()->Mutate(message, max_size);
143   if (size_t new_size = output->Write(*message)) {
144     assert(new_size <= output->size());
145     GetCache()->Store(output->data(), new_size, message);
146     return new_size;
147   }
148   return 0;
149 }
150 
CrossOverMessages(unsigned int seed,const InputReader & input1,const InputReader & input2,OutputWriter * output,protobuf::Message * message1,protobuf::Message * message2)151 size_t CrossOverMessages(unsigned int seed, const InputReader& input1,
152                          const InputReader& input2, OutputWriter* output,
153                          protobuf::Message* message1,
154                          protobuf::Message* message2) {
155   GetMutator()->Seed(seed);
156   input1.Read(message1);
157   input2.Read(message2);
158   size_t max_size = GetMaxSize(input1, *output, *message1);
159   GetMutator()->CrossOver(*message2, message1, max_size);
160   if (size_t new_size = output->Write(*message1)) {
161     assert(new_size <= output->size());
162     GetCache()->Store(output->data(), new_size, message1);
163     return new_size;
164   }
165   return 0;
166 }
167 
MutateTextMessage(uint8_t * data,size_t size,size_t max_size,unsigned int seed,protobuf::Message * message)168 size_t MutateTextMessage(uint8_t* data, size_t size, size_t max_size,
169                          unsigned int seed, protobuf::Message* message) {
170   TextInputReader input(data, size);
171   TextOutputWriter output(data, max_size);
172   return MutateMessage(seed, input, &output, message);
173 }
174 
CrossOverTextMessages(const uint8_t * data1,size_t size1,const uint8_t * data2,size_t size2,uint8_t * out,size_t max_out_size,unsigned int seed,protobuf::Message * message1,protobuf::Message * message2)175 size_t CrossOverTextMessages(const uint8_t* data1, size_t size1,
176                              const uint8_t* data2, size_t size2, uint8_t* out,
177                              size_t max_out_size, unsigned int seed,
178                              protobuf::Message* message1,
179                              protobuf::Message* message2) {
180   TextInputReader input1(data1, size1);
181   TextInputReader input2(data2, size2);
182   TextOutputWriter output(out, max_out_size);
183   return CrossOverMessages(seed, input1, input2, &output, message1, message2);
184 }
185 
MutateBinaryMessage(uint8_t * data,size_t size,size_t max_size,unsigned int seed,protobuf::Message * message)186 size_t MutateBinaryMessage(uint8_t* data, size_t size, size_t max_size,
187                            unsigned int seed, protobuf::Message* message) {
188   BinaryInputReader input(data, size);
189   BinaryOutputWriter output(data, max_size);
190   return MutateMessage(seed, input, &output, message);
191 }
192 
CrossOverBinaryMessages(const uint8_t * data1,size_t size1,const uint8_t * data2,size_t size2,uint8_t * out,size_t max_out_size,unsigned int seed,protobuf::Message * message1,protobuf::Message * message2)193 size_t CrossOverBinaryMessages(const uint8_t* data1, size_t size1,
194                                const uint8_t* data2, size_t size2, uint8_t* out,
195                                size_t max_out_size, unsigned int seed,
196                                protobuf::Message* message1,
197                                protobuf::Message* message2) {
198   BinaryInputReader input1(data1, size1);
199   BinaryInputReader input2(data2, size2);
200   BinaryOutputWriter output(out, max_out_size);
201   return CrossOverMessages(seed, input1, input2, &output, message1, message2);
202 }
203 
204 }  // namespace
205 
CustomProtoMutator(bool binary,uint8_t * data,size_t size,size_t max_size,unsigned int seed,protobuf::Message * input)206 size_t CustomProtoMutator(bool binary, uint8_t* data, size_t size,
207                           size_t max_size, unsigned int seed,
208                           protobuf::Message* input) {
209   auto mutate = binary ? &MutateBinaryMessage : &MutateTextMessage;
210   return mutate(data, size, max_size, seed, input);
211 }
212 
CustomProtoCrossOver(bool binary,const uint8_t * data1,size_t size1,const uint8_t * data2,size_t size2,uint8_t * out,size_t max_out_size,unsigned int seed,protobuf::Message * input1,protobuf::Message * input2)213 size_t CustomProtoCrossOver(bool binary, const uint8_t* data1, size_t size1,
214                             const uint8_t* data2, size_t size2, uint8_t* out,
215                             size_t max_out_size, unsigned int seed,
216                             protobuf::Message* input1,
217                             protobuf::Message* input2) {
218   auto cross = binary ? &CrossOverBinaryMessages : &CrossOverTextMessages;
219   return cross(data1, size1, data2, size2, out, max_out_size, seed, input1,
220                input2);
221 }
222 
LoadProtoInput(bool binary,const uint8_t * data,size_t size,protobuf::Message * input)223 bool LoadProtoInput(bool binary, const uint8_t* data, size_t size,
224                     protobuf::Message* input) {
225   if (GetCache()->LoadIfSame(data, size, input)) return true;
226   auto result = binary ? ParseBinaryMessage(data, size, input)
227                        : ParseTextMessage(data, size, input);
228   if (!result) return false;
229   GetMutator()->Seed(size);
230   GetMutator()->Fix(input);
231   return true;
232 }
233 
RegisterPostProcessor(const protobuf::Descriptor * desc,std::function<void (protobuf::Message * message,unsigned int seed)> callback)234 void RegisterPostProcessor(
235     const protobuf::Descriptor* desc,
236     std::function<void(protobuf::Message* message, unsigned int seed)>
237         callback) {
238   GetMutator()->RegisterPostProcessor(desc, callback);
239 }
240 
241 }  // namespace libfuzzer
242 }  // namespace protobuf_mutator
243