diff --git a/src/fuzzers/FuzzBuffer.cc b/src/fuzzers/FuzzBuffer.cc index ebae1a1db1..098ce92f5c 100644 --- a/src/fuzzers/FuzzBuffer.cc +++ b/src/fuzzers/FuzzBuffer.cc @@ -17,6 +17,23 @@ bool zeek::detail::FuzzBuffer::Valid() const return true; } +int zeek::detail::FuzzBuffer::ChunkCount() const + { + auto pos = begin; + int chunks = 0; + while (pos < end) + { + pos = (const unsigned char*)memmem(pos, end - pos, + PKT_MAGIC, PKT_MAGIC_LEN); + if ( ! pos ) + break; + pos += PKT_MAGIC_LEN + 1; + chunks++; + } + + return chunks; + } + std::optional zeek::detail::FuzzBuffer::Next() { if ( begin == end ) diff --git a/src/fuzzers/FuzzBuffer.h b/src/fuzzers/FuzzBuffer.h index 1b9cdb58a3..f79f20091e 100644 --- a/src/fuzzers/FuzzBuffer.h +++ b/src/fuzzers/FuzzBuffer.h @@ -43,6 +43,11 @@ public: */ bool Valid() const; + /** + * @return the number of chunks in the fuzz buffer object + */ + int ChunkCount() const; + /** * @return the next chunk to deliver, if one could be extracted */ diff --git a/src/fuzzers/packet-fuzzer.cc b/src/fuzzers/packet-fuzzer.cc index d3bb28b806..59002f6c79 100644 --- a/src/fuzzers/packet-fuzzer.cc +++ b/src/fuzzers/packet-fuzzer.cc @@ -18,6 +18,9 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) if ( ! fb.Valid() ) return 0; + if ( fb.ChunkCount() > 64 ) + return 0; + for ( ; ; ) { auto chunk = fb.Next(); diff --git a/src/fuzzers/pop3-fuzzer.cc b/src/fuzzers/pop3-fuzzer.cc index 5cc4776aad..2c6e01fc5c 100644 --- a/src/fuzzers/pop3-fuzzer.cc +++ b/src/fuzzers/pop3-fuzzer.cc @@ -50,6 +50,9 @@ extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) if ( ! fb.Valid() ) return 0; + if ( fb.ChunkCount() > 64 ) + return 0; + auto conn = add_connection(); auto a = add_analyzer(conn);