Rollback pr2, rollback 的时候，更新 wallet 的状态， 改内核的下一个版本发 (#493)

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

diff --git a/wallet/wallet.go b/wallet/wallet.go
index 739e949..e22c44e 100644 (file)
--- a/wallet/wallet.go
+++ b/wallet/wallet.go
@@ -80,16 +80,22 @@ func NewWallet(store WalletStore, account *account.Manager, asset *asset.Registr
                return nil, err
        }
 
+       return w, nil
+}
+
+// Run go to run some wallet recorvery and clean tx thread
+func (w *Wallet) Run() error {
        var err error
        w.TxMsgSub, err = w.EventDispatcher.Subscribe(protocol.TxMsgEvent{})
        if err != nil {
-               return nil, err
+               return err
        }
 
        go w.walletUpdater()
        go w.delUnconfirmedTx()
        go w.MemPoolTxQueryLoop()
-       return w, nil
+
+       return nil
 }
 
 // MemPoolTxQueryLoop constantly pass a transaction accepted by mempool to the wallet.
@@ -307,6 +313,22 @@ func (w *Wallet) DeleteAccount(accountID string) (err error) {
        return nil
 }
 
+// Rollback wallet to target height
+func (w *Wallet) Rollback(targetHeight uint64) error {
+       for w.Status.WorkHeight > targetHeight {
+               block, err := w.Chain.GetBlockByHash(&w.Status.WorkHash)
+               if err != nil {
+                       return err
+               }
+
+               if err = w.DetachBlock(block); err != nil {
+                       return err
+               }
+       }
+
+       return nil
+}
+
 func (w *Wallet) UpdateAccountAlias(accountID string, newAlias string) (err error) {
        w.rw.Lock()
        defer w.rw.Unlock()