Merge pull request #362 from ahrtr/protect_rollback_20221221

Add protection when `mmap` somehow fails

bolt_windows.go

@@ -107,8 +107,11 @@ func munmap(db *DB) error {
 	}
 
 	addr := (uintptr)(unsafe.Pointer(&db.data[0]))
+	var err1 error
 	if err := syscall.UnmapViewOfFile(addr); err != nil {
-		return os.NewSyscallError("UnmapViewOfFile", err)
+		err1 = os.NewSyscallError("UnmapViewOfFile", err)
 	}
-	return nil
+	db.data = nil
+	db.datasz = 0
+	return err1
 }

db.go

@@ -491,8 +491,19 @@ func (db *DB) mmap(minsz int) error {
 	return nil
 }
 
+func (db *DB) invalidate() {
+	db.dataref = nil
+	db.data = nil
+	db.datasz = 0
+
+	db.meta0 = nil
+	db.meta1 = nil
+}
+
 // munmap unmaps the data file from memory.
 func (db *DB) munmap() error {
+	defer db.invalidate()
+
 	if err := munmap(db); err != nil {
 		return fmt.Errorf("unmap error: " + err.Error())
 	}
@@ -701,6 +712,13 @@ func (db *DB) beginTx() (*Tx, error) {
 		return nil, ErrDatabaseNotOpen
 	}
 
+	// Exit if the database is not correctly mapped.
+	if db.data == nil {
+		db.mmaplock.RUnlock()
+		db.metalock.Unlock()
+		return nil, ErrInvalidMapping
+	}
+
 	// Create a transaction associated with the database.
 	t := &Tx{}
 	t.init(db)
@@ -742,6 +760,12 @@ func (db *DB) beginRWTx() (*Tx, error) {
 		return nil, ErrDatabaseNotOpen
 	}
 
+	// Exit if the database is not correctly mapped.
+	if db.data == nil {
+		db.rwlock.Unlock()
+		return nil, ErrInvalidMapping
+	}
+
 	// Create a transaction associated with the database.
 	t := &Tx{writable: true}
 	t.init(db)
@@ -1016,6 +1040,7 @@ func (db *DB) Stats() Stats {
 // This is for internal access to the raw data bytes from the C cursor, use
 // carefully, or not at all.
 func (db *DB) Info() *Info {
+	_assert(db.data != nil, "database file isn't correctly mapped")
 	return &Info{uintptr(unsafe.Pointer(&db.data[0])), db.pageSize}
 }
 
@@ -1042,7 +1067,7 @@ func (db *DB) meta() *meta {
 		metaB = db.meta0
 	}
 
-	// Use higher meta page if valid. Otherwise fallback to previous, if valid.
+	// Use higher meta page if valid. Otherwise, fallback to previous, if valid.
 	if err := metaA.validate(); err == nil {
 		return metaA
 	} else if err := metaB.validate(); err == nil {

db_test.go

@@ -10,11 +10,13 @@ import (
 	"math/rand"
 	"os"
 	"path/filepath"
+	"reflect"
 	"sync"
 	"testing"
 	"time"
 	"unsafe"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	bolt "go.etcd.io/bbolt"
@@ -1311,6 +1313,28 @@ func TestDB_BatchTime(t *testing.T) {
 	}
 }
 
+// TestDBUnmap verifes that `dataref`, `data` and `datasz` must be reset
+// to zero values respectively after unmapping the db.
+func TestDBUnmap(t *testing.T) {
+	db := btesting.MustCreateDB(t)
+
+	require.NoError(t, db.DB.Close())
+
+	// Ignore the following error:
+	// Error: copylocks: call of reflect.ValueOf copies lock value: go.etcd.io/bbolt.DB contains sync.Once contains sync.Mutex (govet)
+	//nolint:govet
+	v := reflect.ValueOf(*db.DB)
+	dataref := v.FieldByName("dataref")
+	data := v.FieldByName("data")
+	datasz := v.FieldByName("datasz")
+	assert.True(t, dataref.IsNil())
+	assert.True(t, data.IsNil())
+	assert.True(t, datasz.IsZero())
+
+	// Set db.DB to nil to prevent MustCheck from panicking.
+	db.DB = nil
+}
+
 func ExampleDB_Update() {
 	// Open the database.
 	db, err := bolt.Open(tempfile(), 0666, nil)

errors.go

@@ -16,6 +16,9 @@ var (
 	// This typically occurs when a file is not a bolt database.
 	ErrInvalid = errors.New("invalid database")
 
+	// ErrInvalidMapping is returned when the database file fails to get mapped.
+	ErrInvalidMapping = errors.New("database isn't correctly mapped")
+
 	// ErrVersionMismatch is returned when the data file was created with a
 	// different version of Bolt.
 	ErrVersionMismatch = errors.New("version mismatch")

tx.go

@@ -276,13 +276,17 @@ func (tx *Tx) rollback() {
 	}
 	if tx.writable {
 		tx.db.freelist.rollback(tx.meta.txid)
-		if !tx.db.hasSyncedFreelist() {
-			// Reconstruct free page list by scanning the DB to get the whole free page list.
-			// Note: scaning the whole db is heavy if your db size is large in NoSyncFreeList mode.
-			tx.db.freelist.noSyncReload(tx.db.freepages())
-		} else {
-			// Read free page list from freelist page.
-			tx.db.freelist.reload(tx.db.page(tx.db.meta().freelist))
+		// When mmap fails, the `data`, `dataref` and `datasz` may be reset to
+		// zero values, and there is no way to reload free page IDs in this case.
+		if tx.db.data != nil {
+			if !tx.db.hasSyncedFreelist() {
+				// Reconstruct free page list by scanning the DB to get the whole free page list.
+				// Note: scaning the whole db is heavy if your db size is large in NoSyncFreeList mode.
+				tx.db.freelist.noSyncReload(tx.db.freepages())
+			} else {
+				// Read free page list from freelist page.
+				tx.db.freelist.reload(tx.db.page(tx.db.meta().freelist))
+			}
 		}
 	}
 	tx.close()