From a795db6b12ffab7daa34f726d3cf11c059d1e9a7 Mon Sep 17 00:00:00 2001 From: Andras Banki-Horvath Date: Tue, 27 Apr 2021 22:18:50 +0200 Subject: [PATCH 1/3] wallet: support for external wallet DB --- wallet/loader.go | 123 +++++++++++++++++++++++++++++++++-------------- 1 file changed, 86 insertions(+), 37 deletions(-) diff --git a/wallet/loader.go b/wallet/loader.go index 029114c..32f3bd5 100644 --- a/wallet/loader.go +++ b/wallet/loader.go @@ -6,6 +6,7 @@ package wallet import ( "errors" + "fmt" "os" "path/filepath" "sync" @@ -55,6 +56,8 @@ type Loader struct { timeout time.Duration recoveryWindow uint32 wallet *Wallet + localDB bool + walletExists func() (bool, error) db walletdb.DB mu sync.Mutex } @@ -72,18 +75,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 } @@ -134,8 +161,7 @@ func (l *Loader) createNewWallet(pubPassphrase, privPassphrase, return nil, ErrLoaded } - dbPath := filepath.Join(l.dbDirPath, WalletDBName) - exists, err := fileExists(dbPath) + exists, err := l.WalletExists() if err != nil { return nil, err } @@ -143,25 +169,34 @@ func (l *Loader) createNewWallet(pubPassphrase, privPassphrase, return nil, ErrExists } - // 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) - if err != nil { - return nil, err + 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 + } + 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 := CreateWatchingOnly( + l.db, pubPassphrase, l.chainParams, bday, + ) if err != nil { return nil, err } } else { - err = Create( - db, pubPassphrase, privPassphrase, seed, l.chainParams, bday, + err := Create( + l.db, pubPassphrase, privPassphrase, seed, + l.chainParams, bday, ) if err != nil { return nil, err @@ -169,13 +204,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,17 +232,22 @@ func (l *Loader) OpenExistingWallet(pubPassphrase []byte, canConsolePrompt bool) return nil, ErrLoaded } - // Ensure that the network directory exists. - if err := checkCreateDir(l.dbDirPath); err != nil { - return nil, err - } + if l.localDB { + var err error + // Ensure that the network directory exists. + 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) - if err != nil { - log.Errorf("Failed to open database: %v", err) - return nil, err + // Open the database using the boltdb backend. + dbPath := filepath.Join(l.dbDirPath, WalletDBName) + 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 @@ -222,28 +262,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 e != nil { - log.Warnf("Error closing database: %v", e) + 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) { - dbPath := filepath.Join(l.dbDirPath, WalletDBName) - return fileExists(dbPath) + 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,9 +317,11 @@ func (l *Loader) UnloadWallet() error { l.wallet.Stop() l.wallet.WaitForShutdown() - err := l.db.Close() - if err != nil { - return err + if l.localDB { + err := l.db.Close() + if err != nil { + return err + } } l.wallet = nil From 98ba16748e6ca077d53dab638decd1cca07409d2 Mon Sep 17 00:00:00 2001 From: Andras Banki-Horvath Date: Wed, 28 Apr 2021 16:43:59 +0200 Subject: [PATCH 2/3] loader: add txn callback when wallet is created --- wallet/loader.go | 17 ++++++++++++++--- wallet/wallet.go | 42 ++++++++++++++++++++++++++++++++++++++---- 2 files changed, 52 insertions(+), 7 deletions(-) diff --git a/wallet/loader.go b/wallet/loader.go index 32f3bd5..89812a1 100644 --- a/wallet/loader.go +++ b/wallet/loader.go @@ -58,6 +58,7 @@ type Loader struct { wallet *Wallet localDB bool walletExists func() (bool, error) + walletCreated func(db walletdb.ReadWriteTx) error db walletdb.DB mu sync.Mutex } @@ -129,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. @@ -187,16 +197,17 @@ func (l *Loader) createNewWallet(pubPassphrase, privPassphrase, // Initialize the newly created database for the wallet before opening. if isWatchingOnly { - err := CreateWatchingOnly( + err := CreateWatchingOnlyWithCallback( l.db, pubPassphrase, l.chainParams, bday, + l.walletCreated, ) if err != nil { return nil, err } } else { - err := Create( + err := CreateWithCallback( l.db, pubPassphrase, privPassphrase, seed, - l.chainParams, bday, + l.chainParams, bday, l.walletCreated, ) if err != nil { return nil, err diff --git a/wallet/wallet.go b/wallet/wallet.go index 549199f..7e85a89 100644 --- a/wallet/wallet.go +++ b/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 }) } From 13966db554a60dccb503cff58ecea4f650cbc5ef Mon Sep 17 00:00:00 2001 From: Andras Banki-Horvath Date: Thu, 29 Apr 2021 11:31:04 +0200 Subject: [PATCH 3/3] waddrmgr: test flake fix --- waddrmgr/manager_test.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/waddrmgr/manager_test.go b/waddrmgr/manager_test.go index 557bae4..fa46775 100644 --- a/waddrmgr/manager_test.go +++ b/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