7 log "github.com/sirupsen/logrus"
8 "github.com/tendermint/go-wire/data/base58"
9 "github.com/tendermint/tmlibs/db"
11 "github.com/bytom/account"
12 "github.com/bytom/asset"
13 "github.com/bytom/blockchain/pseudohsm"
14 "github.com/bytom/crypto/ed25519/chainkd"
15 "github.com/bytom/crypto/sha3pool"
16 "github.com/bytom/protocol"
17 "github.com/bytom/protocol/bc"
18 "github.com/bytom/protocol/bc/types"
24 //RecoveryIndex walletdb recovery cp number
25 const RecoveryIndex = 5000
27 var walletKey = []byte("walletInfo")
28 var privKeyKey = []byte("keysInfo")
30 //StatusInfo is base valid block info to handle orphan block rollback
31 type StatusInfo struct {
38 //KeyInfo is key import status
40 Alias string `json:"alias"`
41 XPub chainkd.XPub `json:"xpub"`
42 Percent uint8 `json:"percent"`
43 Complete bool `json:"complete"`
46 //Wallet is related to storing account unspent outputs
50 AccountMgr *account.Manager
51 AssetReg *asset.Registry
54 rescanProgress chan struct{}
59 //NewWallet return a new wallet instance
60 func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, hsm *pseudohsm.HSM, chain *protocol.Chain) (*Wallet, error) {
67 rescanProgress: make(chan struct{}, 1),
68 keysInfo: make([]KeyInfo, 0),
71 if err := w.loadWalletInfo(); err != nil {
75 if err := w.loadKeysInfo(); err != nil {
79 w.ImportPrivKey = w.getImportKeyFlag()
86 //GetWalletInfo return stored wallet info and nil,if error,
87 //return initial wallet info and err
88 func (w *Wallet) loadWalletInfo() error {
89 if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
90 return json.Unmarshal(rawWallet, &w.status)
93 block, err := w.chain.GetBlockByHeight(0)
97 return w.AttachBlock(block)
100 func (w *Wallet) commitWalletInfo(batch db.Batch) error {
101 rawWallet, err := json.Marshal(w.status)
103 log.WithField("err", err).Error("save wallet info")
107 batch.Set(walletKey, rawWallet)
112 //GetWalletInfo return stored wallet info and nil,if error,
113 //return initial wallet info and err
114 func (w *Wallet) loadKeysInfo() error {
115 if rawKeyInfo := w.DB.Get(privKeyKey); rawKeyInfo != nil {
116 json.Unmarshal(rawKeyInfo, &w.keysInfo)
122 func (w *Wallet) commitkeysInfo() error {
123 rawKeysInfo, err := json.Marshal(w.keysInfo)
125 log.WithField("err", err).Error("save wallet info")
128 w.DB.Set(privKeyKey, rawKeysInfo)
132 func (w *Wallet) getImportKeyFlag() bool {
133 for _, v := range w.keysInfo {
134 if v.Complete == false {
141 // AttachBlock attach a new block
142 func (w *Wallet) AttachBlock(block *types.Block) error {
143 if block.PreviousBlockHash != w.status.WorkHash {
144 log.Warn("wallet skip attachBlock due to status hash not equal to previous hash")
148 blockHash := block.Hash()
149 txStatus, err := w.chain.GetTransactionStatus(&blockHash)
154 storeBatch := w.DB.NewBatch()
155 w.indexTransactions(storeBatch, block, txStatus)
156 w.buildAccountUTXOs(storeBatch, block, txStatus)
158 w.status.WorkHeight = block.Height
159 w.status.WorkHash = block.Hash()
160 if w.status.WorkHeight >= w.status.BestHeight {
161 w.status.BestHeight = w.status.WorkHeight
162 w.status.BestHash = w.status.WorkHash
164 return w.commitWalletInfo(storeBatch)
167 // DetachBlock detach a block and rollback state
168 func (w *Wallet) DetachBlock(block *types.Block) error {
169 blockHash := block.Hash()
170 txStatus, err := w.chain.GetTransactionStatus(&blockHash)
175 storeBatch := w.DB.NewBatch()
176 w.reverseAccountUTXOs(storeBatch, block, txStatus)
177 w.deleteTransactions(storeBatch, w.status.BestHeight)
179 w.status.BestHeight = block.Height - 1
180 w.status.BestHash = block.PreviousBlockHash
182 if w.status.WorkHeight > w.status.BestHeight {
183 w.status.WorkHeight = w.status.BestHeight
184 w.status.WorkHash = w.status.BestHash
187 return w.commitWalletInfo(storeBatch)
190 //WalletUpdate process every valid block and reverse every invalid block which need to rollback
191 func (w *Wallet) walletUpdater() {
193 getRescanNotification(w)
195 for !w.chain.InMainChain(w.status.BestHash) {
196 block, err := w.chain.GetBlockByHash(&w.status.BestHash)
198 log.WithField("err", err).Error("walletUpdater GetBlockByHash")
202 if err := w.DetachBlock(block); err != nil {
203 log.WithField("err", err).Error("walletUpdater detachBlock")
208 block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight + 1)
210 <-w.chain.BlockWaiter(w.status.WorkHeight + 1)
214 if err := w.AttachBlock(block); err != nil {
215 log.WithField("err", err).Error("walletUpdater stop")
221 func getRescanNotification(w *Wallet) {
223 case <-w.rescanProgress:
224 w.status.WorkHeight = 0
225 block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight)
226 w.status.WorkHash = block.Hash()
232 // ExportAccountPrivKey exports the account private key as a WIF for encoding as a string
233 // in the Wallet Import Formt.
234 func (w *Wallet) ExportAccountPrivKey(xpub chainkd.XPub, auth string) (*string, error) {
235 xprv, err := w.Hsm.LoadChainKDKey(xpub, auth)
240 sha3pool.Sum256(hashed[:], xprv[:])
242 tmp := append(xprv[:], hashed[:4]...)
243 res := base58.Encode(tmp)
247 // ImportAccountPrivKey imports the account key in the Wallet Import Formt.
248 func (w *Wallet) ImportAccountPrivKey(xprv chainkd.XPrv, keyAlias, auth string, index uint64, accountAlias string) (*pseudohsm.XPub, error) {
249 if w.Hsm.HasAlias(keyAlias) {
250 return nil, pseudohsm.ErrDuplicateKeyAlias
252 if w.Hsm.HasKey(xprv) {
253 return nil, pseudohsm.ErrDuplicateKey
256 if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
257 return nil, account.ErrDuplicateAlias
260 xpub, _, err := w.Hsm.ImportXPrvKey(auth, keyAlias, xprv)
265 newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias, nil)
269 if err := w.recoveryAccountWalletDB(newAccount, xpub, index, keyAlias); err != nil {
275 // ImportAccountXpubKey imports the account key in the Wallet Import Formt.
276 func (w *Wallet) ImportAccountXpubKey(xpubIndex int, xpub pseudohsm.XPub, cpIndex uint64) error {
277 accountAlias := fmt.Sprintf("recovery_%d", xpubIndex)
279 if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
280 return account.ErrDuplicateAlias
283 newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias, nil)
288 return w.recoveryAccountWalletDB(newAccount, &xpub, cpIndex, xpub.Alias)
291 func (w *Wallet) recoveryAccountWalletDB(account *account.Account, XPub *pseudohsm.XPub, index uint64, keyAlias string) error {
292 if err := w.createProgram(account, XPub, index); err != nil {
295 w.ImportPrivKey = true
301 w.keysInfo = append(w.keysInfo, tmp)
308 func (w *Wallet) createProgram(account *account.Account, XPub *pseudohsm.XPub, index uint64) error {
309 for i := uint64(0); i < index; i++ {
310 if _, err := w.AccountMgr.CreateAddress(nil, account.ID, false); err != nil {
317 func (w *Wallet) rescanBlocks() {
319 case w.rescanProgress <- struct{}{}:
325 //GetRescanStatus return key import rescan status
326 func (w *Wallet) GetRescanStatus() ([]KeyInfo, error) {
327 keysInfo := make([]KeyInfo, len(w.keysInfo))
329 if rawKeyInfo := w.DB.Get(privKeyKey); rawKeyInfo != nil {
330 if err := json.Unmarshal(rawKeyInfo, &keysInfo); err != nil {
335 var status StatusInfo
336 if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
337 if err := json.Unmarshal(rawWallet, &status); err != nil {
342 for i := range keysInfo {
343 if keysInfo[i].Complete == true || status.BestHeight == 0 {
344 keysInfo[i].Percent = 100
348 keysInfo[i].Percent = uint8(status.WorkHeight * 100 / status.BestHeight)
349 if keysInfo[i].Percent == 100 {
350 keysInfo[i].Complete = true
356 //checkRescanStatus mark private key import process `Complete` if rescan finished
357 func checkRescanStatus(w *Wallet) {
358 if !w.ImportPrivKey || w.status.WorkHeight < w.status.BestHeight {
362 w.ImportPrivKey = false
363 for _, keyInfo := range w.keysInfo {
364 keyInfo.Complete = true