Wallet store interface (#217)

[bytom/vapor.git] / wallet / recovery.go

diff --git a/wallet/recovery.go b/wallet/recovery.go
index ad2ca40..a2aae5a 100644 (file)
--- a/wallet/recovery.go
+++ b/wallet/recovery.go
@@ -1,15 +1,12 @@
 package wallet
 
 import (
-       "encoding/json"
        "fmt"
        "reflect"
        "sync"
        "sync/atomic"
        "time"
 
-       "github.com/tendermint/tmlibs/db"
-
        "github.com/vapor/account"
        "github.com/vapor/blockchain/signers"
        "github.com/vapor/crypto/ed25519/chainkd"
@@ -29,15 +26,14 @@ const (
        addrRecoveryWindow = uint64(128)
 )
 
-//recoveryKey key for db store recovery info.
 var (
-       recoveryKey = []byte("RecoveryInfo")
-
        // ErrRecoveryBusy another recovery in progress, can not get recovery manager lock
        ErrRecoveryBusy = errors.New("another recovery in progress")
 
        // ErrInvalidAcctID can not find account by account id
-       ErrInvalidAcctID = errors.New("invalid account id")
+       ErrInvalidAcctID     = errors.New("invalid account id")
+       ErrGetRecoveryStatus = errors.New("failed to get recovery status.")
+       ErrRecoveryStatus    = errors.New("recovery status is nil.")
 )
 
 // branchRecoveryState maintains the required state in-order to properly
@@ -128,8 +124,8 @@ func newAddressRecoveryState(recoveryWindow uint64, account *account.Account) *a
        }
 }
 
-// recoveryState used to record the status of a recovery process.
-type recoveryState struct {
+// RecoveryState used to record the status of a recovery process.
+type RecoveryState struct {
        // XPubs recovery account xPubs
        XPubs []chainkd.XPub
 
@@ -146,8 +142,8 @@ type recoveryState struct {
        AccountsStatus map[string]*addressRecoveryState
 }
 
-func newRecoveryState() *recoveryState {
-       return &recoveryState{
+func newRecoveryState() *RecoveryState {
+       return &RecoveryState{
                AccountsStatus: make(map[string]*addressRecoveryState),
                StartTime:      time.Now(),
        }
@@ -156,7 +152,7 @@ func newRecoveryState() *recoveryState {
 // stateForScope returns a ScopeRecoveryState for the provided key scope. If one
 // does not already exist, a new one will be generated with the RecoveryState's
 // recoveryWindow.
-func (rs *recoveryState) stateForScope(account *account.Account) {
+func (rs *RecoveryState) stateForScope(account *account.Account) {
        // If the account recovery state already exists, return it.
        if _, ok := rs.AccountsStatus[account.ID]; ok {
                return
@@ -171,7 +167,7 @@ func (rs *recoveryState) stateForScope(account *account.Account) {
 type recoveryManager struct {
        mu sync.Mutex
 
-       db         db.DB
+       store      WalletStore
        accountMgr *account.Manager
 
        locked int32
@@ -180,17 +176,17 @@ type recoveryManager struct {
 
        // state encapsulates and allocates the necessary recovery state for all
        // key scopes and subsidiary derivation paths.
-       state *recoveryState
+       state *RecoveryState
 
        //addresses all addresses derivation lookahead used when
        // attempting to recover the set of used addresses.
        addresses map[bc.Hash]*account.CtrlProgram
 }
 
-// newRecoveryManager create recovery manger.
-func newRecoveryManager(db db.DB, accountMgr *account.Manager) *recoveryManager {
+// NewRecoveryManager create recovery manger.
+func NewRecoveryManager(store WalletStore, accountMgr *account.Manager) *recoveryManager {
        return &recoveryManager{
-               db:         db,
+               store:      store,
                accountMgr: accountMgr,
                addresses:  make(map[bc.Hash]*account.CtrlProgram),
                state:      newRecoveryState(),
@@ -203,16 +199,24 @@ func (m *recoveryManager) AddrResurrect(accts []*account.Account) error {
 
        for _, acct := range accts {
                m.state.stateForScope(acct)
-               if err := m.extendScanAddresses(acct.ID, true); err != nil {
+               if err := m.extendScanAddresses(acct.ID, false); err != nil {
                        return err
                }
 
-               if err := m.extendScanAddresses(acct.ID, false); err != nil {
+               //Bip32 path no change field, no need to create addresses repeatedly.
+               if acct.DeriveRule == signers.BIP0032 {
+                       continue
+               }
+               if err := m.extendScanAddresses(acct.ID, true); err != nil {
                        return err
                }
        }
 
        m.state.StartTime = time.Now()
+       if err := m.commitStatusInfo(); err != nil {
+               return err
+       }
+
        m.started = true
        return nil
 }
@@ -233,18 +237,16 @@ func (m *recoveryManager) AcctResurrect(xPubs []chainkd.XPub) error {
                return err
        }
        m.state.StartTime = time.Now()
+       if err := m.commitStatusInfo(); err != nil {
+               return err
+       }
+
        m.started = true
        return nil
 }
 
 func (m *recoveryManager) commitStatusInfo() error {
-       rawStatus, err := json.Marshal(m.state)
-       if err != nil {
-               return err
-       }
-
-       m.db.Set(recoveryKey, rawStatus)
-       return nil
+       return m.store.SetRecoveryStatus(m.state)
 }
 
 func genAcctAlias(xPubs []chainkd.XPub, index uint64) string {
@@ -321,7 +323,7 @@ func (m *recoveryManager) extendScanAddresses(accountID string, change bool) err
 func (m *recoveryManager) processBlock(b *types.Block) error {
        for _, tx := range b.Transactions {
                for _, output := range tx.Outputs {
-                       if cp, ok := m.addresses[getCPHash(output.ControlProgram)]; ok {
+                       if cp, ok := m.addresses[getCPHash(output.ControlProgram())]; ok {
                                status, ok := m.state.AccountsStatus[cp.AccountID]
                                if !ok {
                                        return ErrInvalidAcctID
@@ -353,7 +355,7 @@ func (m *recoveryManager) FilterRecoveryTxs(b *types.Block) error {
 }
 
 func (m *recoveryManager) finished() {
-       m.db.Delete(recoveryKey)
+       m.store.DeleteRecoveryStatus()
        m.started = false
        m.addresses = make(map[bc.Hash]*account.CtrlProgram)
        m.state = newRecoveryState()
@@ -364,14 +366,17 @@ func (m *recoveryManager) LoadStatusInfo() error {
        m.mu.Lock()
        defer m.mu.Unlock()
 
-       rawStatus := m.db.Get(recoveryKey)
-       if rawStatus == nil {
+       if m.state == nil {
+               return ErrRecoveryStatus
+       }
+       status, err := m.store.GetRecoveryStatus()
+       if err == ErrGetRecoveryStatus {
                return nil
        }
-
-       if err := json.Unmarshal(rawStatus, m.state); err != nil {
+       if err != nil {
                return err
        }
+       m.state = status
 
        if m.state.XPubs != nil && !m.tryStartXPubsRec() {
                return ErrRecoveryBusy