From 6aa6eea7bc44993c00ee9d969b328f8347ec41b2 Mon Sep 17 00:00:00 2001
From: Justin Azoff <justin@corelight.com>
Date: Thu, 21 May 2020 12:55:00 -0400
Subject: [PATCH] Speed up ChunkCount validity check

When counting chunks for the purpose of a Valid check, only count up to
chunk_count_limit + 1 chunks.  This speeds up the skipping of the 70,000
chunk test file considerably.

Before:
    Processed 1 inputs in 0.025517s

After:
    Processed 1 inputs in 0.000620s
---
 src/fuzzers/FuzzBuffer.cc | 6 +++---
 src/fuzzers/FuzzBuffer.h  | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/fuzzers/FuzzBuffer.cc b/src/fuzzers/FuzzBuffer.cc
index 0eb166c786..393e830d78 100644
--- a/src/fuzzers/FuzzBuffer.cc
+++ b/src/fuzzers/FuzzBuffer.cc
@@ -14,18 +14,18 @@ bool zeek::detail::FuzzBuffer::Valid(int chunk_count_limit) const
 	if ( memcmp(begin, PKT_MAGIC, PKT_MAGIC_LEN) != 0)
 		return false;
 
-	if ( ChunkCount() > chunk_count_limit )
+	if ( ChunkCount(chunk_count_limit + 1) > chunk_count_limit )
 		return false;
 
 	return true;
 	}
 
-int zeek::detail::FuzzBuffer::ChunkCount() const
+int zeek::detail::FuzzBuffer::ChunkCount(int chunk_count_limit) const
 	{
 	auto pos = begin;
 	int chunks = 0;
 
-	while ( pos < end )
+	while ( pos < end && (chunk_count_limit == 0 || chunks < chunk_count_limit) )
 		{
 		pos = (const unsigned char*)memmem(pos, end - pos,
 		                                   PKT_MAGIC, PKT_MAGIC_LEN);
diff --git a/src/fuzzers/FuzzBuffer.h b/src/fuzzers/FuzzBuffer.h
index 1959e6f18b..830ed65528 100644
--- a/src/fuzzers/FuzzBuffer.h
+++ b/src/fuzzers/FuzzBuffer.h
@@ -49,7 +49,7 @@ public:
 	/**
 	 * @return  the number of chunks in the fuzz buffer object
 	 */
-	int ChunkCount() const;
+	int ChunkCount(int chunk_count_limit = 0) const;
 
 	/**
 	 * @return  the next chunk to deliver, if one could be extracted