add test check page

Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>

Author: Elbehery

db_whitebox_test.go

@@ -1,13 +1,20 @@
 package bbolt
 
 import (
+	"bytes"
+	crand "crypto/rand"
+	"encoding/binary"
+	"fmt"
+	"math/rand"
 	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"go.etcd.io/bbolt/errors"
+	"go.etcd.io/bbolt/internal/common"
+	"go.etcd.io/bbolt/internal/guts_cli"
 )
 
 func TestOpenWithPreLoadFreelist(t *testing.T) {
@@ -112,6 +119,194 @@ func TestMethodPage(t *testing.T) {
 	}
 }
 
+func TestTx_Check_CorruptPage(t *testing.T) {
+	testCases := []struct {
+		name      string
+		bucketKey string
+		isPanic   bool
+		isErr     bool
+	}{
+		{
+			name:      "no corrupted page",
+			bucketKey: "testBucket",
+			isPanic:   false,
+			isErr:     false,
+		},
+		{
+			name:      "corrupt page by violating btree invariant",
+			bucketKey: "testBucket",
+			isPanic:   false,
+			isErr:     true,
+		},
+		{
+			name:      "corrupt page by dumping random bytes",
+			bucketKey: "testBucket",
+			isPanic:   true,
+			isErr:     false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			fileName, err := prepareData(t)
+			require.NoError(t, err)
+
+			t.Logf("Creating db file '%v'.", fileName)
+			db, err := Open(fileName, 0666, &Options{PageSize: 4096})
+			require.NoError(t, err)
+			defer func() {
+				require.NoError(t, db.Close())
+			}()
+
+			uErr := db.Update(func(tx *Tx) error {
+				t.Logf("Creating bucket '%v'.", tc.bucketKey)
+				b, bErr := tx.CreateBucketIfNotExists([]byte(tc.bucketKey))
+				require.NoError(t, bErr)
+				t.Logf("Generating random data in bucket '%v'.", tc.bucketKey)
+				generateSampleDataInBucket(t, b, 3)
+				return nil
+			})
+			require.NoError(t, uErr)
+
+			var victimPageId common.Pgid
+			if tc.isPanic || tc.isErr {
+				t.Logf("Corrupting random leaf page in bucket '%v'.", tc.bucketKey)
+				victimPageId = corruptLeafPage(t, db, tc.isPanic)
+			}
+
+			if tc.isPanic {
+				defer func() {
+					r := recover()
+					require.NotNil(t, r)
+				}()
+			}
+
+			t.Log("Running consistency check.")
+			vErr := db.View(func(tx *Tx) error {
+				chkConfig := checkConfig{
+					kvStringer: HexKVStringer(),
+				}
+				ch := make(chan error)
+				defer close(ch)
+
+				var cErrs []error
+				go func() {
+					for cErr := range ch {
+						cErrs = append(cErrs, cErr)
+					}
+				}()
+				tx.check(chkConfig, ch)
+
+				if tc.isErr {
+					require.Equal(t, 1, len(cErrs))
+					require.ErrorContains(t, cErrs[0], fmt.Sprintf("leaf page(%d)", victimPageId))
+				}
+				return nil
+			})
+			require.NoError(t, vErr)
+		})
+	}
+}
+
+// corruptLeafPage write an invalid leafPageElement into the victim page.
+func corruptLeafPage(t testing.TB, db *DB, isPanic bool) common.Pgid {
+	t.Helper()
+
+	victimPageId := findVictimPageId(t, db)
+
+	victimPage, victimBuf, err := guts_cli.ReadPage(db.Path(), uint64(victimPageId))
+	require.NoError(t, err)
+	require.NotNil(t, victimPage)
+	require.NotNil(t, victimBuf)
+	require.True(t, victimPage.IsLeafPage())
+	require.True(t, victimPage.Count() > 0)
+
+	// Dumping random bytes in victim page for corruption.
+	copy(victimBuf[16:], generateCorruptionBytes(t, isPanic))
+	// Write the corrupt page to db file.
+	err = guts_cli.WritePage(db.Path(), victimBuf)
+	require.NoError(t, err)
+
+	return victimPageId
+}
+
+// findVictimPageId finds all the leaf pages of a bucket and picks a random leaf page as a victim to be corrupted.
+func findVictimPageId(t testing.TB, db *DB) common.Pgid {
+	t.Helper()
+	// Read DB's RootPage.
+	rootPageId, _, err := guts_cli.GetRootPage(db.Path())
+	require.NoError(t, err)
+	rootPage, _, err := guts_cli.ReadPage(db.Path(), uint64(rootPageId))
+	require.NoError(t, err)
+	require.True(t, rootPage.IsLeafPage())
+	require.Equal(t, 1, len(rootPage.LeafPageElements()))
+	// Find Bucket's RootPage.
+	lpe := rootPage.LeafPageElement(uint16(0))
+	require.Equal(t, uint32(common.BranchPageFlag), lpe.Flags())
+	k := lpe.Key()
+	require.Equal(t, "testBucket", string(k))
+	bucketRootPageId := lpe.Bucket().RootPage()
+	// Read Bucket's RootPage.
+	bucketRootPage, _, err := guts_cli.ReadPage(db.Path(), uint64(bucketRootPageId))
+	require.NoError(t, err)
+	require.Equal(t, uint16(common.BranchPageFlag), bucketRootPage.Flags())
+	// Retrieve Bucket's PageIds
+	var bucketPageIds []common.Pgid
+	for _, bpe := range bucketRootPage.BranchPageElements() {
+		bucketPageIds = append(bucketPageIds, bpe.Pgid())
+	}
+
+	victimPageId := bucketPageIds[rand.Intn(len(bucketPageIds))]
+	return victimPageId
+}
+
+// generateSampleDataInBucket fill in sample data into given bucket to create the given
+// number of leafPages. To control the number of leafPages, sample data are generated in order.
+func generateSampleDataInBucket(t testing.TB, bk *Bucket, lPages int) {
+	t.Helper()
+	currentKey := 1
+	currentVal := 100
+	for i := 0; i < lPages; i++ {
+		currentSize := common.PageHeaderSize
+		for {
+			err := bk.Put(convertIntIntoBytes(t, currentKey), convertIntIntoBytes(t, currentVal))
+			require.NoError(t, err)
+			currentSize += 16 + 4 + 4
+			if currentSize >= 4096 {
+				break
+			}
+			currentKey++
+			currentVal++
+		}
+	}
+}
+
+// generateCorruptionBytes returns random bytes to corrupt a page.
+// It inserts a page element which violates the btree key order if no panic is expected.
+// Otherwise, it dumps random bytes into a page which causes a panic.
+func generateCorruptionBytes(t testing.TB, isPanic bool) []byte {
+	if isPanic {
+		// Generated data size is between pageHeader and pageSize.
+		corruptDataLength := rand.Intn(4096-16) + 16
+		corruptData := make([]byte, corruptDataLength)
+		_, err := crand.Read(corruptData)
+		require.NoError(t, err)
+		return corruptData
+	}
+	// Insert LeafPageElement which violates the BTree range.
+	invalidLPE := common.NewLeafPageElement(0, 0, 0, 0)
+	var buf bytes.Buffer
+	binary.Write(&buf, binary.BigEndian, invalidLPE)
+	return buf.Bytes()
+}
+
+func convertIntIntoBytes(t testing.TB, i int) []byte {
+	t.Helper()
+	buf := make([]byte, 4)
+	binary.BigEndian.PutUint32(buf, uint32(i))
+	return buf
+}
+
 func prepareData(t *testing.T) (string, error) {
 	fileName := filepath.Join(t.TempDir(), "db")
 	db, err := Open(fileName, 0666, nil)