Add fuzz test for NewIteratorWithStartAndPrefix (#1992)

Co-authored-by: Alberto Benegiamo <alberto.benegiamo@gmail.com>
ava-labs · Nov 8, 2023 · 1329a59 · 1329a59
1 parent 22f3c89
commit 1329a59
Show file tree

Hide file tree

Showing 10 changed files with 148 additions and 46 deletions.
diff --git a/database/corruptabledb/db_test.go b/database/corruptabledb/db_test.go
@@ -26,16 +26,21 @@ func TestInterface(t *testing.T) {
 	}
 }
 
-func FuzzKeyValue(f *testing.F) {
+func newDB() *Database {
 	baseDB := memdb.New()
-	db := New(baseDB)
-	database.FuzzKeyValue(f, db)
+	return New(baseDB)
+}
+
+func FuzzKeyValue(f *testing.F) {
+	database.FuzzKeyValue(f, newDB())
 }
 
 func FuzzNewIteratorWithPrefix(f *testing.F) {
-	baseDB := memdb.New()
-	db := New(baseDB)
-	database.FuzzNewIteratorWithPrefix(f, db)
+	database.FuzzNewIteratorWithPrefix(f, newDB())
+}
+
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	database.FuzzNewIteratorWithStartAndPrefix(f, newDB())
 }
 
 // TestCorruption tests to make sure corruptabledb wrapper works as expected.
@@ -70,9 +75,7 @@ func TestCorruption(t *testing.T) {
 			return err
 		},
 	}
-	baseDB := memdb.New()
-	// wrap this db
-	corruptableDB := New(baseDB)
+	corruptableDB := newDB()
 	_ = corruptableDB.handleError(errTest)
 	for name, testFn := range tests {
 		t.Run(name, func(tt *testing.T) {
@@ -176,9 +179,7 @@ func TestIterator(t *testing.T) {
 			ctrl := gomock.NewController(t)
 
 			// Make a database
-			baseDB := memdb.New()
-			corruptableDB := New(baseDB)
-
+			corruptableDB := newDB()
 			// Put a key-value pair in the database.
 			require.NoError(corruptableDB.Put([]byte{0}, []byte{1}))
 

diff --git a/database/encdb/db_test.go b/database/encdb/db_test.go
@@ -24,28 +24,30 @@ func TestInterface(t *testing.T) {
 	}
 }
 
-func FuzzKeyValue(f *testing.F) {
+func newDB(t testing.TB) database.Database {
 	unencryptedDB := memdb.New()
 	db, err := New([]byte(testPassword), unencryptedDB)
-	require.NoError(f, err)
-	database.FuzzKeyValue(f, db)
+	require.NoError(t, err)
+	return db
+}
+
+func FuzzKeyValue(f *testing.F) {
+	database.FuzzKeyValue(f, newDB(f))
 }
 
 func FuzzNewIteratorWithPrefix(f *testing.F) {
-	unencryptedDB := memdb.New()
-	db, err := New([]byte(testPassword), unencryptedDB)
-	require.NoError(f, err)
-	database.FuzzNewIteratorWithPrefix(f, db)
+	database.FuzzNewIteratorWithPrefix(f, newDB(f))
+}
+
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	database.FuzzNewIteratorWithStartAndPrefix(f, newDB(f))
 }
 
 func BenchmarkInterface(b *testing.B) {
 	for _, size := range database.BenchmarkSizes {
 		keys, values := database.SetupBenchmark(b, size[0], size[1], size[2])
 		for _, bench := range database.Benchmarks {
-			unencryptedDB := memdb.New()
-			db, err := New([]byte(testPassword), unencryptedDB)
-			require.NoError(b, err)
-			bench(b, db, "encdb", keys, values)
+			bench(b, newDB(b), "encdb", keys, values)
 		}
 	}
 }
diff --git a/database/leveldb/db_test.go b/database/leveldb/db_test.go
@@ -26,34 +26,39 @@ func TestInterface(t *testing.T) {
 	}
 }
 
-func FuzzKeyValue(f *testing.F) {
-	folder := f.TempDir()
+func newDB(t testing.TB) database.Database {
+	folder := t.TempDir()
 	db, err := New(folder, nil, logging.NoLog{}, "", prometheus.NewRegistry())
-	require.NoError(f, err)
+	require.NoError(t, err)
+	return db
+}
 
+func FuzzKeyValue(f *testing.F) {
+	db := newDB(f)
 	defer db.Close()
 
 	database.FuzzKeyValue(f, db)
 }
 
 func FuzzNewIteratorWithPrefix(f *testing.F) {
-	folder := f.TempDir()
-	db, err := New(folder, nil, logging.NoLog{}, "", prometheus.NewRegistry())
-	require.NoError(f, err)
-
+	db := newDB(f)
 	defer db.Close()
 
 	database.FuzzNewIteratorWithPrefix(f, db)
 }
 
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	db := newDB(f)
+	defer db.Close()
+
+	database.FuzzNewIteratorWithStartAndPrefix(f, db)
+}
+
 func BenchmarkInterface(b *testing.B) {
 	for _, size := range database.BenchmarkSizes {
 		keys, values := database.SetupBenchmark(b, size[0], size[1], size[2])
 		for _, bench := range database.Benchmarks {
-			folder := b.TempDir()
-
-			db, err := New(folder, nil, logging.NoLog{}, "", prometheus.NewRegistry())
-			require.NoError(b, err)
+			db := newDB(b)
 
 			bench(b, db, "leveldb", keys, values)
 

diff --git a/database/memdb/db_test.go b/database/memdb/db_test.go
@@ -23,6 +23,10 @@ func FuzzNewIteratorWithPrefix(f *testing.F) {
 	database.FuzzNewIteratorWithPrefix(f, New())
 }
 
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	database.FuzzNewIteratorWithStartAndPrefix(f, New())
+}
+
 func BenchmarkInterface(b *testing.B) {
 	for _, size := range database.BenchmarkSizes {
 		keys, values := database.SetupBenchmark(b, size[0], size[1], size[2])

diff --git a/database/meterdb/db_test.go b/database/meterdb/db_test.go
@@ -24,28 +24,30 @@ func TestInterface(t *testing.T) {
 	}
 }
 
-func FuzzKeyValue(f *testing.F) {
+func newDB(t testing.TB) database.Database {
 	baseDB := memdb.New()
 	db, err := New("", prometheus.NewRegistry(), baseDB)
-	require.NoError(f, err)
-	database.FuzzKeyValue(f, db)
+	require.NoError(t, err)
+	return db
+}
+
+func FuzzKeyValue(f *testing.F) {
+	database.FuzzKeyValue(f, newDB(f))
 }
 
 func FuzzNewIteratorWithPrefix(f *testing.F) {
-	baseDB := memdb.New()
-	db, err := New("", prometheus.NewRegistry(), baseDB)
-	require.NoError(f, err)
-	database.FuzzNewIteratorWithPrefix(f, db)
+	database.FuzzNewIteratorWithPrefix(f, newDB(f))
+}
+
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	database.FuzzNewIteratorWithStartAndPrefix(f, newDB(f))
 }
 
 func BenchmarkInterface(b *testing.B) {
 	for _, size := range database.BenchmarkSizes {
 		keys, values := database.SetupBenchmark(b, size[0], size[1], size[2])
 		for _, bench := range database.Benchmarks {
-			baseDB := memdb.New()
-			db, err := New("", prometheus.NewRegistry(), baseDB)
-			require.NoError(b, err)
-			bench(b, db, "meterdb", keys, values)
+			bench(b, newDB(b), "meterdb", keys, values)
 		}
 	}
 }
diff --git a/database/pebble/db_test.go b/database/pebble/db_test.go
@@ -41,6 +41,12 @@ func FuzzNewIteratorWithPrefix(f *testing.F) {
 	_ = db.Close()
 }
 
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	db := newDB(f)
+	database.FuzzNewIteratorWithStartAndPrefix(f, db)
+	_ = db.Close()
+}
+
 func BenchmarkInterface(b *testing.B) {
 	for _, size := range database.BenchmarkSizes {
 		keys, values := database.SetupBenchmark(b, size[0], size[1], size[2])

diff --git a/database/prefixdb/db_test.go b/database/prefixdb/db_test.go
@@ -30,6 +30,10 @@ func FuzzNewIteratorWithPrefix(f *testing.F) {
 	database.FuzzNewIteratorWithPrefix(f, New([]byte(""), memdb.New()))
 }
 
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	database.FuzzNewIteratorWithStartAndPrefix(f, New([]byte(""), memdb.New()))
+}
+
 func BenchmarkInterface(b *testing.B) {
 	for _, size := range database.BenchmarkSizes {
 		keys, values := database.SetupBenchmark(b, size[0], size[1], size[2])

diff --git a/database/rpcdb/db_test.go b/database/rpcdb/db_test.go
@@ -75,6 +75,13 @@ func FuzzNewIteratorWithPrefix(f *testing.F) {
 	db.closeFn()
 }
 
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	db := setupDB(f)
+	database.FuzzNewIteratorWithStartAndPrefix(f, db.client)
+
+	db.closeFn()
+}
+
 func BenchmarkInterface(b *testing.B) {
 	for _, size := range database.BenchmarkSizes {
 		keys, values := database.SetupBenchmark(b, size[0], size[1], size[2])

diff --git a/database/test_database.go b/database/test_database.go
@@ -1266,7 +1266,74 @@ func FuzzNewIteratorWithPrefix(f *testing.F, db Database) {
 			require.Equal(expected[string(iter.Key())], val)
 			numIterElts++
 		}
-		require.Equal(len(expectedList), numIterElts)
+		require.Len(expectedList, numIterElts)
+
+		// Clear the database for the next fuzz iteration.
+		require.NoError(AtomicClear(db, db))
+	})
+}
+
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F, db Database) {
+	const (
+		maxKeyLen   = 32
+		maxValueLen = 32
+	)
+
+	f.Fuzz(func(
+		t *testing.T,
+		randSeed int64,
+		start []byte,
+		prefix []byte,
+		numKeyValues uint,
+	) {
+		require := require.New(t)
+		r := rand.New(rand.NewSource(randSeed)) // #nosec G404
+
+		expected := map[string][]byte{}
+
+		// Put a bunch of key-values
+		for i := 0; i < int(numKeyValues); i++ {
+			key := make([]byte, r.Intn(maxKeyLen))
+			_, _ = r.Read(key) // #nosec G404
+
+			value := make([]byte, r.Intn(maxValueLen))
+			_, _ = r.Read(value) // #nosec G404
+
+			if len(value) == 0 {
+				// Consistently treat zero length values as nil
+				// so that we can compare [expected] and [got] with
+				// require.Equal, which treats nil and empty byte
+				// as being unequal, whereas the database treats
+				// them as being equal.
+				value = nil
+			}
+
+			if bytes.HasPrefix(key, prefix) && bytes.Compare(key, start) >= 0 {
+				expected[string(key)] = value
+			}
+
+			require.NoError(db.Put(key, value))
+		}
+
+		expectedList := maps.Keys(expected)
+		slices.Sort(expectedList)
+
+		iter := db.NewIteratorWithStartAndPrefix(start, prefix)
+		defer iter.Release()
+
+		// Assert the iterator returns the expected key-values.
+		numIterElts := 0
+		for iter.Next() {
+			val := iter.Value()
+			if len(val) == 0 {
+				val = nil
+			}
+			keyStr := string(iter.Key())
+			require.Equal(expectedList[numIterElts], keyStr)
+			require.Equal(expected[keyStr], val)
+			numIterElts++
+		}
+		require.Len(expectedList, numIterElts)
 
 		// Clear the database for the next fuzz iteration.
 		require.NoError(AtomicClear(db, db))

diff --git a/database/versiondb/db_test.go b/database/versiondb/db_test.go
@@ -27,6 +27,10 @@ func FuzzNewIteratorWithPrefix(f *testing.F) {
 	database.FuzzNewIteratorWithPrefix(f, New(memdb.New()))
 }
 
+func FuzzNewIteratorWithStartAndPrefix(f *testing.F) {
+	database.FuzzNewIteratorWithStartAndPrefix(f, New(memdb.New()))
+}
+
 func TestIterate(t *testing.T) {
 	require := require.New(t)