From c437fb0172076019c191aaeb137ca8e5fccb77ab Mon Sep 17 00:00:00 2001
From: yahtoo <yahtoo.ma@gmail.com>
Date: Tue, 29 Jan 2019 23:29:30 +0800
Subject: [PATCH] Wallet: fix restore deleted account bug (#1559)

---
 account/accounts.go     |  3 ++
 wallet/recovery.go      |  8 ++++
 wallet/recovery_test.go | 89 +++++++++++++++++++++++++++++++++++++++++
 wallet/wallet.go        |  3 +-
 4 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/account/accounts.go b/account/accounts.go
index ab9bd9e21..48131bf48 100644
--- a/account/accounts.go
+++ b/account/accounts.go
@@ -315,6 +315,9 @@ func (m *Manager) deleteAccountControlPrograms(accountID string) error {
 			m.db.Delete(ContractKey(hash))
 		}
 	}
+	m.db.Delete(bip44ContractIndexKey(accountID, false))
+	m.db.Delete(bip44ContractIndexKey(accountID, true))
+	m.db.Delete(contractIndexKey(accountID))
 	return nil
 }
 
diff --git a/wallet/recovery.go b/wallet/recovery.go
index b3a059019..4b03efe20 100644
--- a/wallet/recovery.go
+++ b/wallet/recovery.go
@@ -216,6 +216,10 @@ func (m *recoveryManager) AddrResurrect(accts []*account.Account) error {
 	}
 
 	m.state.StartTime = time.Now()
+	if err := m.commitStatusInfo(); err != nil {
+		return err
+	}
+
 	m.started = true
 	return nil
 }
@@ -236,6 +240,10 @@ 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
 }
diff --git a/wallet/recovery_test.go b/wallet/recovery_test.go
index dc5d7ea43..a814678b3 100644
--- a/wallet/recovery_test.go
+++ b/wallet/recovery_test.go
@@ -13,6 +13,7 @@ import (
 	"github.com/bytom/blockchain/pseudohsm"
 	"github.com/bytom/blockchain/signers"
 	"github.com/bytom/blockchain/txbuilder"
+	"github.com/bytom/common"
 	"github.com/bytom/consensus"
 	"github.com/bytom/crypto/ed25519/chainkd"
 	"github.com/bytom/errors"
@@ -603,3 +604,91 @@ func TestStateForScope(t *testing.T) {
 		t.Fatal("state for scope test err")
 	}
 }
+
+func bip44ContractIndexKey(accountID string, change bool) []byte {
+	contractIndexPrefix := []byte("ContractIndex")
+	key := append(contractIndexPrefix, accountID...)
+	if change {
+		return append(key, []byte{1}...)
+	}
+	return append(key, []byte{0}...)
+}
+
+func TestContractIndexResidue(t *testing.T) {
+	dirPath, err := ioutil.TempDir(".", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(dirPath)
+
+	testDB := dbm.NewDB("testdb", "leveldb", dirPath)
+	hsm, err := pseudohsm.New(dirPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	contractIndexResidue := uint64(5)
+	acctMgr := account.NewManager(testDB, nil)
+	recoveryMgr := newRecoveryManager(testDB, acctMgr)
+	acct := &account.Account{ID: "testA", Alias: "test1", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub1.XPub}, KeyIndex: 1, DeriveRule: signers.BIP0044}}
+
+	cp1 := &account.CtrlProgram{AccountID: acct.ID, Address: "address1", KeyIndex: 10, Change: false}
+
+	setContractIndexKey := func(acctMgr *account.Manager, accountID string, change bool) {
+		testDB.Set(bip44ContractIndexKey(accountID, change), common.Unit64ToBytes(contractIndexResidue))
+	}
+
+	delAccount := func(acctMgr *account.Manager, accountID string, change bool) {
+		acctMgr.DeleteAccount(accountID)
+	}
+
+	recoveryMgr.state.XPubsStatus = newBranchRecoveryState(acctRecoveryWindow)
+	recoveryMgr.state.XPubs = []chainkd.XPub{xpub1.XPub}
+	recoveryMgr.state.stateForScope(acct)
+
+	cases := []struct {
+		acct       *account.Account
+		cp         *account.CtrlProgram
+		preProcess func(acctMgr *account.Manager, accountID string, change bool)
+		err        error
+		wantCPNum  uint64
+	}{
+		{acct, cp1, setContractIndexKey, nil, 5},
+		{acct, cp1, delAccount, nil, 10},
+	}
+
+	for _, c := range cases {
+		if c.preProcess != nil {
+			c.preProcess(acctMgr, c.acct.ID, c.cp.Change)
+		}
+
+		if err := acctMgr.SaveAccount(acct); err != nil {
+			t.Fatal("ReportFound test err:", err)
+		}
+
+		if err := recoveryMgr.reportFound(c.acct, c.cp); err != c.err {
+			t.Fatal("ContractIndexResidue test err:", err, c.acct.ID)
+		}
+		cps, err := acctMgr.ListControlProgram()
+		if err != nil {
+			t.Fatal("list control program err:", err)
+		}
+
+		cpNum := uint64(0)
+		for _, cp := range cps {
+			if cp.Address == "" || cp.AccountID != c.acct.ID {
+				continue
+			}
+			cpNum++
+		}
+
+		if cpNum != c.wantCPNum {
+			t.Fatal("Test contract index residue cp num err want:", c.wantCPNum, " got:", cpNum)
+		}
+	}
+}
diff --git a/wallet/wallet.go b/wallet/wallet.go
index efdf32aa2..e81259cfe 100644
--- a/wallet/wallet.go
+++ b/wallet/wallet.go
@@ -185,7 +185,8 @@ func (w *Wallet) AttachBlock(block *types.Block) error {
 	}
 
 	if err := w.RecoveryMgr.FilterRecoveryTxs(block); err != nil {
-		return err
+		log.WithField("err", err).Error("filter recovery txs")
+		w.RecoveryMgr.finished()
 	}
 
 	storeBatch := w.DB.NewBatch()