From e0c5ce72cf7ea6356132dbac65a3e5515bb71a63 Mon Sep 17 00:00:00 2001 From: Olaoluwa Osuntokun Date: Fri, 13 Aug 2021 16:13:06 -0700 Subject: [PATCH] waddrmgr: add new DeriveFromKeyPathCache method for faster key retrieval In this commit, we add a new method `DeriveFromKeyPathCache` that gives callers a way to more quickly obtain a private key they know they'll be using frequently. This method lets a caller avoid the write database transaction as well as the EC operations to derive the key itself (BIP 32). --- waddrmgr/error.go | 5 ++ waddrmgr/manager.go | 21 +++++---- waddrmgr/manager_test.go | 95 ++++++++++++++++++++++++++++++++++++++ waddrmgr/scoped_manager.go | 87 +++++++++++++++++++++++++++++++++- 4 files changed, 198 insertions(+), 10 deletions(-) diff --git a/waddrmgr/error.go b/waddrmgr/error.go index bb3acc2..f836e2d 100644 --- a/waddrmgr/error.go +++ b/waddrmgr/error.go @@ -139,6 +139,10 @@ const ( // ErrBlockNotFound is returned when we attempt to retrieve the hash for // a block that we do not know of. ErrBlockNotFound + + // ErrAccountNotCached is returned when we attempt to perform an + // operation that relies on an account begin cached but it isn't. + ErrAccountNotCached ) // Map of ErrorCode values back to their constant names for pretty printing. @@ -165,6 +169,7 @@ var errorCodeStrings = map[ErrorCode]string{ ErrCallBackBreak: "ErrCallBackBreak", ErrEmptyPassphrase: "ErrEmptyPassphrase", ErrScopeNotFound: "ErrScopeNotFound", + ErrAccountNotCached: "ErrAccountNotCached", } // String returns the ErrorCode as a human-readable name. diff --git a/waddrmgr/manager.go b/waddrmgr/manager.go index 1d36e10..570ef4f 100644 --- a/waddrmgr/manager.go +++ b/waddrmgr/manager.go @@ -17,6 +17,7 @@ import ( "github.com/btcsuite/btcwallet/internal/zero" "github.com/btcsuite/btcwallet/snacl" "github.com/btcsuite/btcwallet/walletdb" + "github.com/lightninglabs/neutrino/cache/lru" ) const ( @@ -602,11 +603,12 @@ func (m *Manager) NewScopedKeyManager(ns walletdb.ReadWriteBucket, // Finally, we'll register this new scoped manager with the root // manager. m.scopedManagers[scope] = &ScopedKeyManager{ - scope: scope, - addrSchema: addrSchema, - rootManager: m, - addrs: make(map[addrKey]ManagedAddress), - acctInfo: make(map[uint32]*accountInfo), + scope: scope, + addrSchema: addrSchema, + rootManager: m, + addrs: make(map[addrKey]ManagedAddress), + acctInfo: make(map[uint32]*accountInfo), + privKeyCache: lru.NewCache(defaultPrivKeyCacheSize), } m.externalAddrSchemas[addrSchema.ExternalAddrType] = append( m.externalAddrSchemas[addrSchema.ExternalAddrType], scope, @@ -1620,10 +1622,11 @@ func loadManager(ns walletdb.ReadBucket, pubPassphrase []byte, } scopedManagers[scope] = &ScopedKeyManager{ - scope: scope, - addrSchema: *scopeSchema, - addrs: make(map[addrKey]ManagedAddress), - acctInfo: make(map[uint32]*accountInfo), + scope: scope, + addrSchema: *scopeSchema, + addrs: make(map[addrKey]ManagedAddress), + acctInfo: make(map[uint32]*accountInfo), + privKeyCache: lru.NewCache(defaultPrivKeyCacheSize), } return nil diff --git a/waddrmgr/manager_test.go b/waddrmgr/manager_test.go index b96367b..3e5e943 100644 --- a/waddrmgr/manager_test.go +++ b/waddrmgr/manager_test.go @@ -14,6 +14,7 @@ import ( "testing" "time" + "github.com/btcsuite/btcd/btcec" "github.com/btcsuite/btcd/chaincfg" "github.com/btcsuite/btcd/chaincfg/chainhash" "github.com/btcsuite/btcutil" @@ -21,6 +22,7 @@ import ( "github.com/btcsuite/btcwallet/snacl" "github.com/btcsuite/btcwallet/walletdb" "github.com/davecgh/go-spew/spew" + "github.com/stretchr/testify/require" ) // failingCryptoKey is an implementation of the EncryptorDecryptor interface @@ -2764,3 +2766,96 @@ func testNewRawAccount(t *testing.T, _ *Manager, db walletdb.DB, accountTargetAddr.AddrHash()) } } + +// TestDeriveFromKeyPathCache tests that the DeriveFromKeyPathCache method will +// properly cache items in the cache, and return corresponding errors if the +// account isn't properly cached. +func TestDeriveFromKeyPathCache(t *testing.T) { + t.Parallel() + + teardown, db := emptyDB(t) + defer teardown() + + // We'll start the test by creating a new root manager that will be + // used for the duration of the test. + var mgr *Manager + err := walletdb.Update(db, func(tx walletdb.ReadWriteTx) error { + ns, err := tx.CreateTopLevelBucket(waddrmgrNamespaceKey) + if err != nil { + return err + } + err = Create( + ns, seed, pubPassphrase, privPassphrase, + &chaincfg.MainNetParams, fastScrypt, time.Time{}, + ) + if err != nil { + return err + } + mgr, err = Open(ns, pubPassphrase, &chaincfg.MainNetParams) + if err != nil { + return err + } + + return mgr.Unlock(ns, privPassphrase) + }) + require.NoError(t, err, "create/open: unexpected error: %v", err) + + defer mgr.Close() + + // Now that we have the manager created, we'll fetch one of the default + // scopes for usage within this test. + scopedMgr, err := mgr.FetchScopedKeyManager(KeyScopeBIP0044) + require.NoError( + t, err, "unable to fetch scope %v: %v", KeyScopeBIP0044, err, + ) + + keyPath := DerivationPath{ + InternalAccount: 0, + Account: hdkeychain.HardenedKeyStart, + Branch: 10, + Index: 1, + } + + // Our test starts here, we'll attempt to derive a new key using the + // cached method. This should fail at first since the account itself + // isn't cached. + _, err = scopedMgr.DeriveFromKeyPathCache(keyPath) + if !IsError(err, ErrAccountNotCached) { + t.Fatalf("didn't get account not cached error: %v", err) + } + + // Now we'll attempt to derive the key using the normal method that + // requires a database transaction. + var derivedKey *btcec.PrivateKey + err = walletdb.Update(db, func(tx walletdb.ReadWriteTx) error { + ns := tx.ReadWriteBucket(waddrmgrNamespaceKey) + + managedAddr, err := scopedMgr.DeriveFromKeyPath( + ns, keyPath, + ) + if err != nil { + return err + } + + derivedKey, err = managedAddr.(ManagedPubKeyAddress).PrivKey() + if err != nil { + return err + } + + return nil + }) + require.NoError(t, err, "unable to derive addr: %v", err) + + // Next attempt to read the key again from the cache, it should succeed + // this time. + cachedKey, err := scopedMgr.DeriveFromKeyPathCache(keyPath) + require.NoError(t, err, "account wasn't cached") + + // We should be able to read the key again. + cachedKey2, err := scopedMgr.DeriveFromKeyPathCache(keyPath) + require.NoError(t, err, "account wasn't cached") + + // All three keys we have now should match exactly. + require.Equal(t, cachedKey.Serialize(), cachedKey2.Serialize()) + require.Equal(t, derivedKey.Serialize(), cachedKey2.Serialize()) +} diff --git a/waddrmgr/scoped_manager.go b/waddrmgr/scoped_manager.go index 7b901db..6d02c31 100644 --- a/waddrmgr/scoped_manager.go +++ b/waddrmgr/scoped_manager.go @@ -14,6 +14,7 @@ import ( "github.com/btcsuite/btcwallet/internal/zero" "github.com/btcsuite/btcwallet/netparams" "github.com/btcsuite/btcwallet/walletdb" + "github.com/lightninglabs/neutrino/cache/lru" ) // HDVersion represents the different supported schemes of hierarchical @@ -51,6 +52,14 @@ const ( HDVersionSimNetBIP0044 HDVersion = 0x0420bd3a // spub ) +const ( + // privKeyCacheSize is the default size of the LRU cache that we'll use + // to cache private keys to avoid DB and EC operations within the + // wallet. With the default sisize, we'll allocate up to 320 KB to + // caching private keys (ignoring pointer overhead, etc). + defaultPrivKeyCacheSize = 10_000 +) + // DerivationPath represents a derivation path from a particular key manager's // scope. Each ScopedKeyManager starts key derivation from the end of their // cointype hardened key: m/purpose'/cointype'. The fields in this struct allow @@ -220,7 +229,7 @@ type ScopedKeyManager struct { rootManager *Manager // addrs is a cached map of all the addresses that we currently - // manager. + // manage. addrs map[addrKey]ManagedAddress // acctInfo houses information about accounts including what is needed @@ -234,6 +243,11 @@ type ScopedKeyManager struct { // order to encrypt it. deriveOnUnlock []*unlockDeriveInfo + // privKeyCache stores the set of private keys that have been marked as + // items to be cached to allow us to avoid the database and EC + // operations each time a key need to be obtained. + privKeyCache *lru.Cache + mtx sync.RWMutex } @@ -579,6 +593,77 @@ func (s *ScopedKeyManager) AccountProperties(ns walletdb.ReadBucket, return props, nil } +// cachedKey is an entry within the LRU map that stores private keys that are +// to be used frequently. We use this wrapper struct to be able too report the +// size of a given element to the cache. +type cachedKey struct { + key *btcec.PrivateKey +} + +// Size returns the size of this element. Rather than have the cache limit +// based on bytes, we simply report that each element is of size 1, meaning we +// can set our cached based on the amount of keys we want to store, rather than +// the total size of all the keys. +func (c *cachedKey) Size() (uint64, error) { + return 1, nil +} + +// DeriveFromKeyPathCache is identical to DeriveFromKeyPath, however it'll fail +// if the account refracted in the DerivationPath isn't already in the +// in-memory cache. Callers looking for faster private key retrieval can opt to +// call this method, which may fail if things aren't in the cache, then fall +// back to the normal variant. The account can information can be drawn into +// the cache if the normal DeriveFromKeyPath method is used, or the account is +// looked up via any other means. +func (s *ScopedKeyManager) DeriveFromKeyPathCache( + kp DerivationPath) (*btcec.PrivateKey, error) { + + s.mtx.Lock() + defer s.mtx.Unlock() + + // First, try to look up the key itself in the proper cache, if the key + // is here, then we don't need to do anything further. + privKeyVal, err := s.privKeyCache.Get(kp) + if err == nil { + return privKeyVal.(*cachedKey).key, nil + } + + // If the key isn't already in the cache, then we'll try to look up the + // account info in the cache, if this fails, then we exit here as we + // can't move forward without creating a DB transaction, and the point + // of this method is to avoid that. + acctInfo, ok := s.acctInfo[kp.InternalAccount] + if !ok { + return nil, managerError( + ErrAccountNotCached, + "", fmt.Errorf("acct %v not cached", kp.InternalAccount), + ) + } + + watchOnly := s.rootManager.WatchOnly() + private := !s.rootManager.IsLocked() && !watchOnly + + // Now that we have the account information, we can derive the key + // directly. + addrKey, err := s.deriveKey(acctInfo, kp.Branch, kp.Index, private) + if err != nil { + return nil, err + } + + // Now that we have the key, we'll attempt to insert it into the cache, + // and return it as is. + privKey, err := addrKey.ECPrivKey() + if err != nil { + return nil, err + } + _, err = s.privKeyCache.Put(kp, &cachedKey{key: privKey}) + if err != nil { + return nil, err + } + + return privKey, nil +} + // DeriveFromKeyPath attempts to derive a maximal child key (under the BIP0044 // scheme) from a given key path. If key derivation isn't possible, then an // error will be returned.