Merge pull request #734 from bhandras/external_db

wallet: allow using external wallet db

waddrmgr/manager_test.go

@@ -1813,8 +1813,6 @@ func testSync(tc *testContext) bool {
 // It makes use of a test context because the address manager is persistent and
 // much of the testing involves having specific state.
 func TestManager(t *testing.T) {
-	t.Parallel()
-
 	tests := []struct {
 		name                string
 		createdWatchingOnly bool

wallet/loader.go

@@ -6,6 +6,7 @@ package wallet
 
 import (
 	"errors"
+	"fmt"
 	"os"
 	"path/filepath"
 	"sync"
@@ -55,6 +56,9 @@ type Loader struct {
 	timeout        time.Duration
 	recoveryWindow uint32
 	wallet         *Wallet
+	localDB        bool
+	walletExists   func() (bool, error)
+	walletCreated  func(db walletdb.ReadWriteTx) error
 	db             walletdb.DB
 	mu             sync.Mutex
 }
@@ -72,18 +76,42 @@ func NewLoader(chainParams *chaincfg.Params, dbDirPath string,
 		noFreelistSync: noFreelistSync,
 		timeout:        timeout,
 		recoveryWindow: recoveryWindow,
+		localDB:        true,
 	}
 }
 
+// NewLoaderWithDB constructs a Loader with an externally provided DB. This way
+// users are free to use their own walletdb implementation (eg. leveldb, etcd)
+// to store the wallet. Given that the external DB may be shared an additional
+// function is also passed which will override Loader.WalletExists().
+func NewLoaderWithDB(chainParams *chaincfg.Params, recoveryWindow uint32,
+	db walletdb.DB, walletExists func() (bool, error)) (*Loader, error) {
+
+	if db == nil {
+		return nil, fmt.Errorf("no DB provided")
+	}
+
+	if walletExists == nil {
+		return nil, fmt.Errorf("unable to check if wallet exists")
+	}
+
+	return &Loader{
+		chainParams:    chainParams,
+		recoveryWindow: recoveryWindow,
+		localDB:        false,
+		walletExists:   walletExists,
+		db:             db,
+	}, nil
+}
+
 // onLoaded executes each added callback and prevents loader from loading any
 // additional wallets.  Requires mutex to be locked.
-func (l *Loader) onLoaded(w *Wallet, db walletdb.DB) {
+func (l *Loader) onLoaded(w *Wallet) {
 	for _, fn := range l.callbacks {
 		fn(w)
 	}
 
 	l.wallet = w
-	l.db = db
 	l.callbacks = nil // not needed anymore
 }
 
@@ -102,6 +130,15 @@ func (l *Loader) RunAfterLoad(fn func(*Wallet)) {
 	}
 }
 
+// OnWalletCreated adds a function that will be executed the wallet structure
+// is initialized in the wallet database. This is useful if users want to add
+// extra fields in the same transaction (eg. to flag wallet existence).
+func (l *Loader) OnWalletCreated(fn func(walletdb.ReadWriteTx) error) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	l.walletCreated = fn
+}
+
 // CreateNewWallet creates a new wallet using the provided public and private
 // passphrases.  The seed is optional.  If non-nil, addresses are derived from
 // this seed.  If nil, a secure random seed is generated.
@@ -134,8 +171,7 @@ func (l *Loader) createNewWallet(pubPassphrase, privPassphrase,
 		return nil, ErrLoaded
 	}
 
-	dbPath := filepath.Join(l.dbDirPath, WalletDBName)
+	exists, err := l.WalletExists()
-	exists, err := fileExists(dbPath)
 	if err != nil {
 		return nil, err
 	}
@@ -143,25 +179,35 @@ func (l *Loader) createNewWallet(pubPassphrase, privPassphrase,
 		return nil, ErrExists
 	}
 
+	if l.localDB {
+		dbPath := filepath.Join(l.dbDirPath, WalletDBName)
+
 		// Create the wallet database backed by bolt db.
 		err = os.MkdirAll(l.dbDirPath, 0700)
 		if err != nil {
 			return nil, err
 		}
-	db, err := walletdb.Create("bdb", dbPath, l.noFreelistSync, l.timeout)
+		l.db, err = walletdb.Create(
+			"bdb", dbPath, l.noFreelistSync, l.timeout,
+		)
 		if err != nil {
 			return nil, err
 		}
+	}
 
 	// Initialize the newly created database for the wallet before opening.
 	if isWatchingOnly {
-		err = CreateWatchingOnly(db, pubPassphrase, l.chainParams, bday)
+		err := CreateWatchingOnlyWithCallback(
+			l.db, pubPassphrase, l.chainParams, bday,
+			l.walletCreated,
+		)
 		if err != nil {
 			return nil, err
 		}
 	} else {
-		err = Create(
+		err := CreateWithCallback(
-			db, pubPassphrase, privPassphrase, seed, l.chainParams, bday,
+			l.db, pubPassphrase, privPassphrase, seed,
+			l.chainParams, bday, l.walletCreated,
 		)
 		if err != nil {
 			return nil, err
@@ -169,13 +215,13 @@ func (l *Loader) createNewWallet(pubPassphrase, privPassphrase,
 	}
 
 	// Open the newly-created wallet.
-	w, err := Open(db, pubPassphrase, nil, l.chainParams, l.recoveryWindow)
+	w, err := Open(l.db, pubPassphrase, nil, l.chainParams, l.recoveryWindow)
 	if err != nil {
 		return nil, err
 	}
 	w.Start()
 
-	l.onLoaded(w, db)
+	l.onLoaded(w)
 	return w, nil
 }
 
@@ -197,18 +243,23 @@ func (l *Loader) OpenExistingWallet(pubPassphrase []byte, canConsolePrompt bool)
 		return nil, ErrLoaded
 	}
 
+	if l.localDB {
+		var err error
 		// Ensure that the network directory exists.
-	if err := checkCreateDir(l.dbDirPath); err != nil {
+		if err = checkCreateDir(l.dbDirPath); err != nil {
 			return nil, err
 		}
 
 		// Open the database using the boltdb backend.
 		dbPath := filepath.Join(l.dbDirPath, WalletDBName)
-	db, err := walletdb.Open("bdb", dbPath, l.noFreelistSync, l.timeout)
+		l.db, err = walletdb.Open(
+			"bdb", dbPath, l.noFreelistSync, l.timeout,
+		)
 		if err != nil {
 			log.Errorf("Failed to open database: %v", err)
 			return nil, err
 		}
+	}
 
 	var cbs *waddrmgr.OpenCallbacks
 	if canConsolePrompt {
@@ -222,28 +273,35 @@ func (l *Loader) OpenExistingWallet(pubPassphrase []byte, canConsolePrompt bool)
 			ObtainPrivatePass: noConsole,
 		}
 	}
-	w, err := Open(db, pubPassphrase, cbs, l.chainParams, l.recoveryWindow)
+	w, err := Open(l.db, pubPassphrase, cbs, l.chainParams, l.recoveryWindow)
 	if err != nil {
 		// If opening the wallet fails (e.g. because of wrong
 		// passphrase), we must close the backing database to
 		// allow future calls to walletdb.Open().
-		e := db.Close()
+		if l.localDB {
+			e := l.db.Close()
 			if e != nil {
 				log.Warnf("Error closing database: %v", e)
 			}
+		}
+
 		return nil, err
 	}
 	w.Start()
 
-	l.onLoaded(w, db)
+	l.onLoaded(w)
 	return w, nil
 }
 
 // WalletExists returns whether a file exists at the loader's database path.
 // This may return an error for unexpected I/O failures.
 func (l *Loader) WalletExists() (bool, error) {
+	if l.localDB {
 		dbPath := filepath.Join(l.dbDirPath, WalletDBName)
 		return fileExists(dbPath)
+	}
+
+	return l.walletExists()
 }
 
 // LoadedWallet returns the loaded wallet, if any, and a bool for whether the
@@ -270,10 +328,12 @@ func (l *Loader) UnloadWallet() error {
 
 	l.wallet.Stop()
 	l.wallet.WaitForShutdown()
+	if l.localDB {
 		err := l.db.Close()
 		if err != nil {
 			return err
 		}
+	}
 
 	l.wallet = nil
 	l.db = nil

wallet/wallet.go

@@ -3665,6 +3665,29 @@ func (w *Wallet) Database() walletdb.DB {
 	return w.db
 }
 
+// CreateWithCallback is the same as Create with an added callback that will be
+// called in the same transaction the wallet structure is initialized.
+func CreateWithCallback(db walletdb.DB, pubPass, privPass, seed []byte,
+	params *chaincfg.Params, birthday time.Time,
+	cb func(walletdb.ReadWriteTx) error) error {
+
+	return create(
+		db, pubPass, privPass, seed, params, birthday, false, cb,
+	)
+}
+
+// CreateWatchingOnlyWithCallback is the same as CreateWatchingOnly with an
+// added callback that will be called in the same transaction the wallet
+// structure is initialized.
+func CreateWatchingOnlyWithCallback(db walletdb.DB, pubPass []byte,
+	params *chaincfg.Params, birthday time.Time,
+	cb func(walletdb.ReadWriteTx) error) error {
+
+	return create(
+		db, pubPass, nil, nil, params, birthday, true, cb,
+	)
+}
+
 // Create creates an new wallet, writing it to an empty database.  If the passed
 // seed is non-nil, it is used.  Otherwise, a secure random seed of the
 // recommended length is generated.
@@ -3672,7 +3695,7 @@ func Create(db walletdb.DB, pubPass, privPass, seed []byte,
 	params *chaincfg.Params, birthday time.Time) error {
 
 	return create(
-		db, pubPass, privPass, seed, params, birthday, false,
+		db, pubPass, privPass, seed, params, birthday, false, nil,
 	)
 }
 
@@ -3684,12 +3707,13 @@ func CreateWatchingOnly(db walletdb.DB, pubPass []byte,
 	params *chaincfg.Params, birthday time.Time) error {
 
 	return create(
-		db, pubPass, nil, nil, params, birthday, true,
+		db, pubPass, nil, nil, params, birthday, true, nil,
 	)
 }
 
 func create(db walletdb.DB, pubPass, privPass, seed []byte,
-	params *chaincfg.Params, birthday time.Time, isWatchingOnly bool) error {
+	params *chaincfg.Params, birthday time.Time, isWatchingOnly bool,
+	cb func(walletdb.ReadWriteTx) error) error {
 
 	if !isWatchingOnly {
 		// If a seed was provided, ensure that it is of valid length. Otherwise,
@@ -3725,7 +3749,17 @@ func create(db walletdb.DB, pubPass, privPass, seed []byte,
 		if err != nil {
 			return err
 		}
-		return wtxmgr.Create(txmgrNs)
+
+		err = wtxmgr.Create(txmgrNs)
+		if err != nil {
+			return err
+		}
+
+		if cb != nil {
+			return cb(tx)
+		}
+
+		return nil
 	})
 }