[compiler-rt] FuzzedDataProvider: add ConsumeData and method.
Reviewers: metzman
Subscribers: dberris, #sanitizers, llvm-commits
Tags: #sanitizers, #llvm
Differential Revision: https://reviews.llvm.org/D74359
GitOrigin-RevId: 20a604d3f5bc8cdba161e7e12f12f3f5260aad57
diff --git a/FuzzedDataProvider.h b/FuzzedDataProvider.h
index 3e069eb..b0a3773 100644
--- a/FuzzedDataProvider.h
+++ b/FuzzedDataProvider.h
@@ -237,6 +237,18 @@
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 @@
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 <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 @@
// 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
diff --git a/tests/FuzzedDataProviderUnittest.cpp b/tests/FuzzedDataProviderUnittest.cpp
index 2222834..5eb4653 100644
--- a/tests/FuzzedDataProviderUnittest.cpp
+++ b/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<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();