diff --git a/compiler-rt/include/fuzzer/FuzzedDataProvider.h b/compiler-rt/include/fuzzer/FuzzedDataProvider.h --- a/compiler-rt/include/fuzzer/FuzzedDataProvider.h +++ b/compiler-rt/include/fuzzer/FuzzedDataProvider.h @@ -237,6 +237,18 @@ return result + range * ConsumeProbability(); } + // 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 @@ 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 @@ } template - std::vector ConsumeBytes(size_t size, size_t num_bytes_to_consume) { + std::vector 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 @@ // To increase the odds even more, we also call |shrink_to_fit| below. std::vector 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 diff --git a/compiler-rt/lib/fuzzer/tests/FuzzedDataProviderUnittest.cpp b/compiler-rt/lib/fuzzer/tests/FuzzedDataProviderUnittest.cpp --- a/compiler-rt/lib/fuzzer/tests/FuzzedDataProviderUnittest.cpp +++ b/compiler-rt/lib/fuzzer/tests/FuzzedDataProviderUnittest.cpp @@ -399,6 +399,25 @@ -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 Expected(Data, Data + sizeof(Buffer)); + EXPECT_EQ(Expected, std::vector(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(Buffer, Buffer + sizeof(Buffer))); + + // Exhaust the buffer. + EXPECT_EQ(std::vector(Data + 12, Data + sizeof(Data)), + DataProv.ConsumeRemainingBytes()); + EXPECT_EQ(size_t(0), DataProv.ConsumeData(Buffer, sizeof(Buffer))); + EXPECT_EQ(Expected, std::vector(Buffer, Buffer + sizeof(Buffer))); +} + int main(int argc, char **argv) { testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS();