diff --git a/db_whitebox_test.go b/db_whitebox_test.go index 130c3e349..445a964f6 100644 --- a/db_whitebox_test.go +++ b/db_whitebox_test.go @@ -1,6 +1,11 @@ package bbolt import ( + "bytes" + crand "crypto/rand" + "encoding/binary" + "fmt" + "math/rand" "path/filepath" "testing" @@ -8,6 +13,8 @@ import ( "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,189 @@ func TestMethodPage(t *testing.T) { } } +func TestTx_Check_CorruptPage(t *testing.T) { + testCases := []struct { + name string + bucketKey string + expectPanic bool + }{ + { + name: "corrupt page by violating btree invariant", + bucketKey: "testBucket", + expectPanic: false, + }, + { + name: "corrupt page by dumping random bytes", + bucketKey: "testBucket", + expectPanic: true, + }, + } + + 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) + + t.Logf("Corrupting random leaf page in bucket '%v'.", tc.bucketKey) + victimPageId, validPageIds := corruptLeafPage(t, db, tc.expectPanic) + + if tc.expectPanic { + 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) + } + }() + t.Log("Check corrupted page.") + chkConfig.pageId = uint(victimPageId) + tx.check(chkConfig, ch) + require.Equal(t, 1, len(cErrs)) + require.ErrorContains(t, cErrs[0], fmt.Sprintf("leaf page(%d)", victimPageId)) + + t.Log("Check valid pages.") + cErrs = cErrs[:0] + for _, pgId := range validPageIds { + chkConfig.pageId = uint(pgId) + tx.check(chkConfig, ch) + require.Equal(t, 0, len(cErrs)) + } + return nil + }) + require.NoError(t, vErr) + }) + } +} + +// corruptLeafPage write an invalid leafPageElement into the victim page. +func corruptLeafPage(t testing.TB, db *DB, expectPanic bool) (victimPageId common.Pgid, validPageIds []common.Pgid) { + t.Helper() + + victimPageId, validPageIds = findVictimPageId(t, db) + + victimPage, victimBuf, err := guts_cli.ReadPage(db.Path(), uint64(victimPageId)) + require.NoError(t, err) + require.True(t, victimPage.IsLeafPage()) + require.True(t, victimPage.Count() > 0) + + // Dumping random bytes in victim page for corruption. + copy(victimBuf[32:], generateCorruptionBytes(t, expectPanic)) + // Write the corrupt page to db file. + err = guts_cli.WritePage(db.Path(), victimBuf) + require.NoError(t, err) + + return victimPageId, validPageIds +} + +// 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) (victimPageId common.Pgid, validPageIds []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()) + } + + randomIdx := rand.Intn(len(bucketPageIds)) + victimPageId = bucketPageIds[randomIdx] + validPageIds = append(bucketPageIds[:randomIdx], bucketPageIds[randomIdx+1:]...) + return victimPageId, validPageIds +} + +// 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 >= 2200 { + 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, expectPanic bool) []byte { + if expectPanic { + // 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)