From 016a1cc388b4d9abeb901089d4d37e7a692a0699 Mon Sep 17 00:00:00 2001
From: Yongfeng LI <wliyongfeng@gmail.com>
Date: Tue, 17 Apr 2018 11:17:55 +0800
Subject: [PATCH] fix address set marshal bug

---
 wallet/indexer.go     |  2 +-
 wallet/set.go         | 10 +++++-----
 wallet/wallet.go      | 14 +++++++-------
 wallet/wallet_test.go |  7 +++++--
 4 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/wallet/indexer.go b/wallet/indexer.go
index c41d00b5..772ca1eb 100644
--- a/wallet/indexer.go
+++ b/wallet/indexer.go
@@ -390,7 +390,7 @@ transactionLoop:
 			if bytes := w.DB.Get(account.CPKey(hash)); bytes != nil {
 				cp := &account.CtrlProgram{}
 				if err := json.Unmarshal(bytes, cp); err == nil {
-					w.status.selfProgramsOnChain.Add(cp.Address)
+					w.status.OnChainAddresses.Add(cp.Address)
 				}
 
 				annotatedTxs = append(annotatedTxs, w.buildAnnotatedTransaction(tx, b, statusFail, pos))
diff --git a/wallet/set.go b/wallet/set.go
index 627ebd28..5a8c7855 100644
--- a/wallet/set.go
+++ b/wallet/set.go
@@ -1,13 +1,13 @@
 package wallet
 
-type Set map[interface{}]bool
+type AddressSet map[string]bool
 
-func NewSet() Set {
-	return make(Set)
+func NewAddressSet() AddressSet {
+	return make(AddressSet)
 }
 
 // Add Add the specified element to this set if it is not already present (optional operation)
-func (s *Set) Add(i interface{}) bool {
+func (s *AddressSet) Add(i string) bool {
 	_, found := (*s)[i]
 	if found {
 		return false //False if it existed already
@@ -18,7 +18,7 @@ func (s *Set) Add(i interface{}) bool {
 }
 
 // Contains Returns true if this set contains the specified elements
-func (s *Set) Contains(i ...interface{}) bool {
+func (s *AddressSet) Contains(i ...string) bool {
 	for _, val := range i {
 		if _, ok := (*s)[val]; !ok {
 			return false
diff --git a/wallet/wallet.go b/wallet/wallet.go
index 5a061ae5..f5245c5b 100644
--- a/wallet/wallet.go
+++ b/wallet/wallet.go
@@ -29,11 +29,11 @@ var dbKeyForimportingPrivateKey = []byte("importingKeysInfo")
 
 //StatusInfo is base valid block info to handle orphan block rollback
 type StatusInfo struct {
-	WorkHeight uint64
-	WorkHash   bc.Hash
-	BestHeight uint64
-	BestHash   bc.Hash
-	selfProgramsOnChain Set
+	WorkHeight       uint64
+	WorkHash         bc.Hash
+	BestHeight       uint64
+	BestHash         bc.Hash
+	OnChainAddresses AddressSet
 }
 
 //KeyInfo is key import status
@@ -92,7 +92,7 @@ func (w *Wallet) loadWalletInfo() error {
 		return json.Unmarshal(rawWallet, &w.status)
 	}
 
-	w.status.selfProgramsOnChain = NewSet()
+	w.status.OnChainAddresses = NewAddressSet()
 	block, err := w.chain.GetBlockByHeight(0)
 	if err != nil {
 		return err
@@ -364,7 +364,7 @@ func (w *Wallet) updateRescanStatus() {
 
 		if cps, err := w.AccountMgr.ListCtrlProgramsByXpubs(nil, keyInfo.account.XPubs); err == nil {
 			for _, cp := range cps {
-				if !w.status.selfProgramsOnChain.Contains(cp.Address) {
+				if !w.status.OnChainAddresses.Contains(cp.Address) {
 					w.AccountMgr.DeleteAccountControlProgram(cp.ControlProgram)
 				}
 			}
diff --git a/wallet/wallet_test.go b/wallet/wallet_test.go
index 76c80099..4ba4ce7b 100644
--- a/wallet/wallet_test.go
+++ b/wallet/wallet_test.go
@@ -235,14 +235,17 @@ func mockTxData(utxos []*account.UTXO, testAccount *account.Account) (*txbuilder
 }
 
 func mockWallet(walletDB dbm.DB, account *account.Manager, asset *asset.Registry, chain *protocol.Chain) *Wallet {
-	return &Wallet{
+	wallet := &Wallet{
 		DB:                  walletDB,
 		AccountMgr:          account,
 		AssetReg:            asset,
 		chain:               chain,
 		rescanProgress:      make(chan struct{}, 1),
-		selfProgramsOnChain: NewSet(),
 	}
+	wallet.status = StatusInfo{
+		OnChainAddresses: NewAddressSet(),
+	}
+	return wallet
 }
 
 func mockSingleBlock(tx *types.Tx) *types.Block {
-- 
2.11.0