From: yahtoo <yahtoo.ma@gmail.com>
Date: Thu, 1 Feb 2018 07:02:34 +0000 (+0800)
Subject: Add walletdb recover function from keystore (#340)
X-Git-Tag: v1.0.5~323
X-Git-Url: http://git.osdn.net/view?a=commitdiff_plain;h=6fed1ac72b9aacae0526a03a8876205653f8d7f5;p=bytom%2Fbytom.git

Add walletdb recover function from keystore (#340)

* Add walletdb recover function

* Fix walletUpdater goroutine data race bug

* Refactor walletdb recover from keystore code

* Delete useless lock

* Change rescanBlocks signal transfer
---

diff --git a/blockchain/wallet/wallet.go b/blockchain/wallet/wallet.go
index e8fe6fe6..057c2b4e 100755
--- a/blockchain/wallet/wallet.go
+++ b/blockchain/wallet/wallet.go
@@ -2,6 +2,7 @@ package wallet
 
 import (
 	"encoding/json"
+	"fmt"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/tendermint/go-wire/data/base58"
@@ -20,6 +21,9 @@ import (
 //SINGLE single sign
 const SINGLE = 1
 
+//RecoveryIndex walletdb recovery cp number
+const RecoveryIndex = 5000
+
 var walletKey = []byte("walletInfo")
 var privKeyKey = []byte("keysInfo")
 
@@ -52,7 +56,7 @@ type Wallet struct {
 }
 
 //NewWallet return a new wallet instance
-func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, chain *protocol.Chain) (*Wallet, error) {
+func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, chain *protocol.Chain, xpubs []pseudohsm.XPub) (*Wallet, error) {
 	w := &Wallet{
 		DB:             walletDB,
 		AccountMgr:     account,
@@ -62,7 +66,7 @@ func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry,
 		keysInfo:       make([]KeyInfo, 0),
 	}
 
-	if err := w.loadWalletInfo(); err != nil {
+	if err := w.loadWalletInfo(xpubs); err != nil {
 		return nil, err
 	}
 
@@ -71,6 +75,7 @@ func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry,
 	}
 
 	w.ImportPrivKey = w.getImportKeyFlag()
+
 	go w.walletUpdater()
 
 	return w, nil
@@ -78,11 +83,17 @@ func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry,
 
 //GetWalletInfo return stored wallet info and nil,if error,
 //return initial wallet info and err
-func (w *Wallet) loadWalletInfo() error {
+func (w *Wallet) loadWalletInfo(xpubs []pseudohsm.XPub) error {
 	if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
 		return json.Unmarshal(rawWallet, &w.status)
 	}
 
+	for i, v := range xpubs {
+		if err := w.ImportAccountXpubKey(i, v, RecoveryIndex); err != nil {
+			return err
+		}
+	}
+
 	block, err := w.chain.GetBlockByHeight(0)
 	if err != nil {
 		return err
@@ -169,7 +180,6 @@ func (w *Wallet) detachBlock(block *legacy.Block) error {
 //WalletUpdate process every valid block and reverse every invalid block which need to rollback
 func (w *Wallet) walletUpdater() {
 	for {
-		// config.GenesisBlock().hash
 		getRescanNotification(w)
 		checkRescanStatus(w)
 		for !w.chain.InMainChain(w.status.BestHeight, w.status.BestHash) {
@@ -252,6 +262,22 @@ func (w *Wallet) ImportAccountPrivKey(hsm *pseudohsm.HSM, xprv chainkd.XPrv, key
 	return xpub, nil
 }
 
+// ImportAccountXpubKey imports the account key in the Wallet Import Formt.
+func (w *Wallet) ImportAccountXpubKey(xpubIndex int, xpub pseudohsm.XPub, cpIndex uint64) error {
+	accountAlias := fmt.Sprintf("recovery_%d", xpubIndex)
+
+	if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
+		return account.ErrDuplicateAlias
+	}
+
+	newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias, nil)
+	if err != nil {
+		return err
+	}
+
+	return w.recoveryAccountWalletDB(newAccount, &xpub, cpIndex, xpub.Alias)
+}
+
 func (w *Wallet) recoveryAccountWalletDB(account *account.Account, XPub *pseudohsm.XPub, index uint64, keyAlias string) error {
 	if err := w.createProgram(account, XPub, index); err != nil {
 		return err
@@ -279,7 +305,12 @@ func (w *Wallet) createProgram(account *account.Account, XPub *pseudohsm.XPub, i
 }
 
 func (w *Wallet) rescanBlocks() {
-	w.rescanProgress <- struct{}{}
+	select {
+	case <-w.rescanProgress:
+		w.rescanProgress <- struct{}{}
+	default:
+		return
+	}
 }
 
 //GetRescanStatus return key import rescan status
diff --git a/node/node.go b/node/node.go
index 2cf2e45d..681af968 100755
--- a/node/node.go
+++ b/node/node.go
@@ -187,20 +187,22 @@ func NewNode(config *cfg.Config) *Node {
 		return nil
 	}
 
+	hsm, err := pseudohsm.New(config.KeysDir())
+	if err != nil {
+		cmn.Exit(cmn.Fmt("initialize HSM failed: %v", err))
+	}
+
 	if !config.Wallet.Disable {
+		xpubs, _ := hsm.ListKeys()
 		walletDB := dbm.NewDB("wallet", config.DBBackend, config.DBDir())
 		accounts = account.NewManager(walletDB, chain)
 		assets = asset.NewRegistry(walletDB, chain)
-		wallet, err = w.NewWallet(walletDB, accounts, assets, chain)
+		wallet, err = w.NewWallet(walletDB, accounts, assets, chain, xpubs)
 		if err != nil {
 			log.WithField("error", err).Error("init NewWallet")
 		}
 	}
 
-	hsm, err := pseudohsm.New(config.KeysDir())
-	if err != nil {
-		cmn.Exit(cmn.Fmt("initialize HSM failed: %v", err))
-	}
 	bcReactor := bc.NewBlockchainReactor(chain, txPool, accounts, assets, sw, hsm, wallet, txFeed, accessTokens, config.Mining)
 
 	sw.AddReactor("BLOCKCHAIN", bcReactor)