[compiler-rt] FuzzedDataProvider: add ConsumeData and method.

Reviewers: metzman

Subscribers: dberris, #sanitizers, llvm-commits

Tags: #sanitizers, #llvm

Differential Revision: https://reviews.llvm.org/D74359

compiler-rt/include/fuzzer/FuzzedDataProvider.h

@@ -237,6 +237,18 @@ class FuzzedDataProvider {
     return result + range * ConsumeProbability<T>();
   }
 
+  // Writes |num_bytes| of input data to the given destination pointer. If there
+  // is not enough data left, writes all remaining bytes. Return value is the
+  // number of bytes written.
+  // In general, it's better to avoid using this function, but it may be useful
+  // in cases when it's necessary to fill a certain buffer or object with
+  // fuzzing data.
+  size_t ConsumeData(void *destination, size_t num_bytes) {
+    num_bytes = std::min(num_bytes, remaining_bytes_);
+    CopyAndAdvance(destination, num_bytes);
+    return num_bytes;
+  }
+
   // Reports the remaining bytes available for fuzzed input.
   size_t remaining_bytes() { return remaining_bytes_; }
 
@@ -244,6 +256,11 @@ class FuzzedDataProvider {
   FuzzedDataProvider(const FuzzedDataProvider &) = delete;
   FuzzedDataProvider &operator=(const FuzzedDataProvider &) = delete;
 
+  void CopyAndAdvance(void *destination, size_t num_bytes) {
+    std::memcpy(destination, data_ptr_, num_bytes);
+    Advance(num_bytes);
+  }
+
   void Advance(size_t num_bytes) {
     if (num_bytes > remaining_bytes_)
       abort();
@@ -253,7 +270,7 @@ class FuzzedDataProvider {
   }
 
   template <typename T>
-  std::vector<T> ConsumeBytes(size_t size, size_t num_bytes_to_consume) {
+  std::vector<T> ConsumeBytes(size_t size, size_t num_bytes) {
     static_assert(sizeof(T) == sizeof(uint8_t), "Incompatible data type.");
 
     // The point of using the size-based constructor below is to increase the
@@ -264,13 +281,12 @@ class FuzzedDataProvider {
     // To increase the odds even more, we also call |shrink_to_fit| below.
     std::vector<T> result(size);
     if (size == 0) {
-      if (num_bytes_to_consume != 0)
+      if (num_bytes != 0)
         abort();
       return result;
     }
 
-    std::memcpy(result.data(), data_ptr_, num_bytes_to_consume);
-    Advance(num_bytes_to_consume);
+    CopyAndAdvance(result.data(), num_bytes);
 
     // Even though |shrink_to_fit| is also implementation specific, we expect it
     // to provide an additional assurance in case vector's constructor allocated

compiler-rt/lib/fuzzer/tests/FuzzedDataProviderUnittest.cpp

@@ -399,6 +399,25 @@ TEST(FuzzedDataProvider, ConsumeFloatingPoint) {
                                        -13.37, 31.337));
 }
 
+TEST(FuzzedDataProvider, ConsumeData) {
+  FuzzedDataProvider DataProv(Data, sizeof(Data));
+  uint8_t Buffer[10] = {};
+  EXPECT_EQ(sizeof(Buffer), DataProv.ConsumeData(Buffer, sizeof(Buffer)));
+  std::vector<uint8_t> Expected(Data, Data + sizeof(Buffer));
+  EXPECT_EQ(Expected, std::vector<uint8_t>(Buffer, Buffer + sizeof(Buffer)));
+
+  EXPECT_EQ(size_t(2), DataProv.ConsumeData(Buffer, 2));
+  Expected[0] = Data[sizeof(Buffer)];
+  Expected[1] = Data[sizeof(Buffer) + 1];
+  EXPECT_EQ(Expected, std::vector<uint8_t>(Buffer, Buffer + sizeof(Buffer)));
+
+  // Exhaust the buffer.
+  EXPECT_EQ(std::vector<uint8_t>(Data + 12, Data + sizeof(Data)),
+            DataProv.ConsumeRemainingBytes<uint8_t>());
+  EXPECT_EQ(size_t(0), DataProv.ConsumeData(Buffer, sizeof(Buffer)));
+  EXPECT_EQ(Expected, std::vector<uint8_t>(Buffer, Buffer + sizeof(Buffer)));
+}
+
 int main(int argc, char **argv) {
   testing::InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();