psiphon-tunnel-core/psiphon/dataStore_bolt.go

// +build !BADGER_DB,!FILES_DB

/*
 * Copyright (c) 2018, Psiphon Inc.
 * All rights reserved.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package psiphon

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"runtime/debug"
	"sync/atomic"
	"time"

	"github.com/Psiphon-Labs/bolt"
	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common"
)

type datastoreDB struct {
	boltDB   *bolt.DB
	isFailed int32
}

type datastoreTx struct {
	db     *datastoreDB
	boltTx *bolt.Tx
}

type datastoreBucket struct {
	db         *datastoreDB
	boltBucket *bolt.Bucket
}

type datastoreCursor struct {
	db         *datastoreDB
	boltCursor *bolt.Cursor
}

func datastoreOpenDB(rootDataDirectory string) (*datastoreDB, error) {

	var db *datastoreDB
	var err error

	for retry := 0; retry < 3; retry++ {

		db, err = tryDatastoreOpenDB(rootDataDirectory, retry > 0)
		if err == nil {
			break
		}

		NoticeAlert("tryDatastoreOpenDB failed: %s", err)

		// The datastore file may be corrupt, so, in subsequent iterations, set the
		// "reset" flag and attempt to delete the file and try again.
	}

	return db, err
}

func tryDatastoreOpenDB(rootDataDirectory string, reset bool) (retdb *datastoreDB, reterr error) {

	// Testing indicates that the bolt Check function can raise SIGSEGV due to
	// invalid mmap buffer accesses in cases such as opening a valid but
	// truncated datastore file.
	//
	// To handle this, we temporarily set SetPanicOnFault in order to treat the
	// fault as a panic, recover any panic, and return an error which will result
	// in a retry with reset.

	// Begin recovery preamble
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)

	defer func() {
		if r := recover(); r != nil {
			retdb = nil
			reterr = common.ContextError(fmt.Errorf("panic: %v", r))
		}
	}()
	// End recovery preamble

	filename := filepath.Join(rootDataDirectory, "psiphon.boltdb")

	if reset {
		NoticeAlert("tryDatastoreOpenDB: reset")
		os.Remove(filename)
	}

	newDB, err := bolt.Open(filename, 0600, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		return nil, common.ContextError(err)
	}

	// Run consistency checks on datastore and emit errors for diagnostics
	// purposes. We assume this will complete quickly for typical size Psiphon
	// datastores and wait for the check to complete before proceeding.
	err = newDB.View(func(tx *bolt.Tx) error {
		return tx.SynchronousCheck()
	})
	if err != nil {
		return nil, common.ContextError(err)
	}

	err = newDB.Update(func(tx *bolt.Tx) error {
		requiredBuckets := [][]byte{
			datastoreServerEntriesBucket,
			datastoreServerEntryTagsBucket,
			datastoreServerEntryTombstoneTagsBucket,
			datastoreSplitTunnelRouteETagsBucket,
			datastoreSplitTunnelRouteDataBucket,
			datastoreUrlETagsBucket,
			datastoreKeyValueBucket,
			datastoreRemoteServerListStatsBucket,
			datastoreFailedTunnelStatsBucket,
			datastoreSLOKsBucket,
			datastoreTacticsBucket,
			datastoreSpeedTestSamplesBucket,
			datastoreDialParametersBucket,
		}
		for _, bucket := range requiredBuckets {
			_, err := tx.CreateBucketIfNotExists(bucket)
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return nil, common.ContextError(err)
	}

	// Cleanup obsolete buckets

	err = newDB.Update(func(tx *bolt.Tx) error {
		obsoleteBuckets := [][]byte{
			[]byte("tunnelStats"),
			[]byte("rankedServerEntries"),
		}
		for _, obsoleteBucket := range obsoleteBuckets {
			if tx.Bucket(obsoleteBucket) != nil {
				err := tx.DeleteBucket(obsoleteBucket)
				if err != nil {
					NoticeAlert("DeleteBucket %s error: %s", obsoleteBucket, err)
					// Continue, since this is not fatal
				}
			}
		}
		return nil
	})
	if err != nil {
		return nil, common.ContextError(err)
	}

	return &datastoreDB{boltDB: newDB}, nil
}

var errDatastoreFailed = errors.New("datastore has failed")

func (db *datastoreDB) isDatastoreFailed() bool {
	return atomic.LoadInt32(&db.isFailed) == 1
}

func (db *datastoreDB) setDatastoreFailed(r interface{}) {
	atomic.StoreInt32(&db.isFailed, 1)
	NoticeAlert("Datastore failed: %s",
		common.ContextError(fmt.Errorf("panic: %v", r)))
}

func (db *datastoreDB) close() error {

	// Limitation: there is no panic recover in this case. We assume boltDB.Close
	// does not make  mmap accesses and prefer to not continue with the datastore
	// file in a locked or open state. We also assume that any locks aquired by
	// boltDB.Close, held by transactions, will be released even if the
	// transaction panics and the database is in the failed state.

	return db.boltDB.Close()
}

func (db *datastoreDB) view(fn func(tx *datastoreTx) error) (reterr error) {

	// Any bolt function that performs mmap buffer accesses can raise SIGBUS  due
	// to underlying storage changes, such as a truncation of the datastore file
	// or removal or network attached storage, etc.
	//
	// To handle this, we temporarily set SetPanicOnFault in order to treat the
	// fault as a panic, recover any panic to avoid crashing the process, and
	// putting this datastoreDB instance into a failed state. All subsequent
	// calls to this datastoreDBinstance or its related datastoreTx and
	// datastoreBucket instances will fail.

	// Begin recovery preamble
	if db.isDatastoreFailed() {
		return errDatastoreFailed
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			db.setDatastoreFailed(r)
			reterr = errDatastoreFailed
		}
	}()
	// End recovery preamble

	return db.boltDB.View(
		func(tx *bolt.Tx) error {
			err := fn(&datastoreTx{db: db, boltTx: tx})
			if err != nil {
				return common.ContextError(err)
			}
			return nil
		})
}

func (db *datastoreDB) update(fn func(tx *datastoreTx) error) (reterr error) {

	// Begin recovery preamble
	if db.isDatastoreFailed() {
		return errDatastoreFailed
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			db.setDatastoreFailed(r)
			reterr = errDatastoreFailed
		}
	}()
	// End recovery preamble

	return db.boltDB.Update(
		func(tx *bolt.Tx) error {
			err := fn(&datastoreTx{db: db, boltTx: tx})
			if err != nil {
				return common.ContextError(err)
			}
			return nil
		})
}

func (tx *datastoreTx) bucket(name []byte) (retbucket *datastoreBucket) {

	// Begin recovery preamble
	if tx.db.isDatastoreFailed() {
		return &datastoreBucket{db: tx.db, boltBucket: nil}
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			tx.db.setDatastoreFailed(r)
			retbucket = &datastoreBucket{db: tx.db, boltBucket: nil}
		}
	}()
	// End recovery preamble

	return &datastoreBucket{db: tx.db, boltBucket: tx.boltTx.Bucket(name)}
}

func (tx *datastoreTx) clearBucket(name []byte) (reterr error) {

	// Begin recovery preamble
	if tx.db.isDatastoreFailed() {
		return errDatastoreFailed
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			tx.db.setDatastoreFailed(r)
			reterr = errDatastoreFailed
		}
	}()
	// End recovery preamble

	err := tx.boltTx.DeleteBucket(name)
	if err != nil {
		return common.ContextError(err)
	}
	_, err = tx.boltTx.CreateBucket(name)
	if err != nil {
		return common.ContextError(err)
	}
	return nil
}

func (b *datastoreBucket) get(key []byte) (retvalue []byte) {

	// Begin recovery preamble
	if b.db.isDatastoreFailed() {
		return nil
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			b.db.setDatastoreFailed(r)
			retvalue = nil
		}
	}()
	// End recovery preamble

	return b.boltBucket.Get(key)
}

func (b *datastoreBucket) put(key, value []byte) (reterr error) {

	// Begin recovery preamble
	if b.db.isDatastoreFailed() {
		return errDatastoreFailed
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			b.db.setDatastoreFailed(r)
			reterr = errDatastoreFailed
		}
	}()
	// End recovery preamble

	err := b.boltBucket.Put(key, value)
	if err != nil {
		return common.ContextError(err)
	}
	return nil
}

func (b *datastoreBucket) delete(key []byte) (reterr error) {

	// Begin recovery preamble
	if b.db.isDatastoreFailed() {
		return errDatastoreFailed
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			b.db.setDatastoreFailed(r)
			reterr = errDatastoreFailed
		}
	}()
	// End recovery preamble

	err := b.boltBucket.Delete(key)
	if err != nil {
		return common.ContextError(err)
	}
	return nil
}

func (b *datastoreBucket) cursor() (retcursor datastoreCursor) {

	// Begin recovery preamble
	if b.db.isDatastoreFailed() {
		return datastoreCursor{db: b.db, boltCursor: nil}
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			b.db.setDatastoreFailed(r)
			retcursor = datastoreCursor{db: b.db, boltCursor: nil}
		}
	}()
	// End recovery preamble

	return datastoreCursor{db: b.db, boltCursor: b.boltBucket.Cursor()}
}

func (c *datastoreCursor) firstKey() (retkey []byte) {

	// Begin recovery preamble
	if c.db.isDatastoreFailed() {
		return nil
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			c.db.setDatastoreFailed(r)
			retkey = nil
		}
	}()
	// End recovery preamble

	key, _ := c.boltCursor.First()
	return key
}

func (c *datastoreCursor) nextKey() (retkey []byte) {

	// Begin recovery preamble
	if c.db.isDatastoreFailed() {
		return nil
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			c.db.setDatastoreFailed(r)
			retkey = nil
		}
	}()
	// End recovery preamble

	key, _ := c.boltCursor.Next()
	return key
}

func (c *datastoreCursor) first() (retkey, retvalue []byte) {

	// Begin recovery preamble
	if c.db.isDatastoreFailed() {
		return nil, nil
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			c.db.setDatastoreFailed(r)
			retkey = nil
			retvalue = nil
		}
	}()
	// End recovery preamble

	return c.boltCursor.First()
}

func (c *datastoreCursor) next() (retkey, retvalue []byte) {

	// Begin recovery preamble
	if c.db.isDatastoreFailed() {
		return nil, nil
	}
	panicOnFault := debug.SetPanicOnFault(true)
	defer debug.SetPanicOnFault(panicOnFault)
	defer func() {
		if r := recover(); r != nil {
			c.db.setDatastoreFailed(r)
			retkey = nil
			retvalue = nil
		}
	}()
	// End recovery preamble

	return c.boltCursor.Next()
}

func (c *datastoreCursor) close() {
	// BoltDB doesn't close cursors.
}