From fd12f9a53355fcc5995097c547969858bfae0752 Mon Sep 17 00:00:00 2001
From: Chengcheng Zhang <943420582@qq.com>
Date: Thu, 4 Jul 2019 17:20:19 +0800
Subject: [PATCH] update

---
 database/wallet_store.go | 13 ++++++++++---
 wallet/recovery.go       | 12 ++++++------
 wallet/store.go          |  6 +++---
 wallet/wallet_test.go    | 12 ++++++++++--
 4 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/database/wallet_store.go b/database/wallet_store.go
index e922ee8a..1442f86a 100644
--- a/database/wallet_store.go
+++ b/database/wallet_store.go
@@ -311,9 +311,16 @@ func (store *WalletStore) GetUnconfirmedTransaction(txID string) (*query.Annotat
 }
 
 // GetRecoveryStatus delete recovery status
-func (store *WalletStore) GetRecoveryStatus(recoveryKey []byte) []byte {
-
-	return store.walletDB.Get(recoveryKey)
+func (store *WalletStore) GetRecoveryStatus() (*wallet.RecoveryState, error) {
+	rawStatus := store.walletDB.Get(dbm.RecoveryKey)
+	if rawStatus == nil {
+		return nil, wallet.ErrGetRecoveryStatus
+	}
+	state := new(wallet.RecoveryState)
+	if err := json.Unmarshal(rawStatus, state); err != nil {
+		return nil, err
+	}
+	return state, nil
 }
 
 // GetWalletInfo get wallet information
diff --git a/wallet/recovery.go b/wallet/recovery.go
index 1828a1ba..91cd7e48 100644
--- a/wallet/recovery.go
+++ b/wallet/recovery.go
@@ -1,7 +1,6 @@
 package wallet
 
 import (
-	"encoding/json"
 	"fmt"
 	"reflect"
 	"sync"
@@ -35,7 +34,8 @@ var (
 	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.")
 )
 
 // branchRecoveryState maintains the required state in-order to properly
@@ -368,14 +368,14 @@ func (m *recoveryManager) LoadStatusInfo() error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	rawStatus := m.store.GetRecoveryStatus(recoveryKey)
-	if rawStatus == nil {
+	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
diff --git a/wallet/store.go b/wallet/store.go
index 05462d2d..8d73adc6 100644
--- a/wallet/store.go
+++ b/wallet/store.go
@@ -23,15 +23,15 @@ type WalletStore interface {
 	GetStandardUTXO(bc.Hash) (*acc.UTXO, error)
 	GetTransaction(string) (*query.AnnotatedTx, error)
 	GetUnconfirmedTransaction(string) (*query.AnnotatedTx, error)
-	GetRecoveryStatus([]byte) []byte
-	GetWalletInfo() (*StatusInfo, error) // need move database.NewWalletStore in wallet package
+	GetRecoveryStatus() (*RecoveryState, error)
+	GetWalletInfo() (*StatusInfo, error)
 	ListAccountUTXOs(string) ([]*acc.UTXO, error)
 	ListTransactions(string, string, uint, bool) ([]*query.AnnotatedTx, error)
 	ListUnconfirmedTransactions() ([]*query.AnnotatedTx, error)
 	SetAssetDefinition(*bc.AssetID, []byte)
 	SetContractUTXO(bc.Hash, *acc.UTXO) error
 	SetGlobalTransactionIndex(string, *bc.Hash, uint64)
-	SetRecoveryStatus(*RecoveryState) error // recoveryManager.state isn't exported outside
+	SetRecoveryStatus(*RecoveryState) error
 	SetTransaction(uint64, *query.AnnotatedTx) error
 	SetUnconfirmedTransaction(string, *query.AnnotatedTx) error
 	SetWalletInfo(*StatusInfo) error
diff --git a/wallet/wallet_test.go b/wallet/wallet_test.go
index ff4710c8..d7bf16c7 100644
--- a/wallet/wallet_test.go
+++ b/wallet/wallet_test.go
@@ -788,8 +788,16 @@ func (store *MockWalletStore) GetUnconfirmedTransaction(txID string) (*query.Ann
 }
 
 // GetRecoveryStatus delete recovery status
-func (store *MockWalletStore) GetRecoveryStatus(recoveryKey []byte) []byte {
-	return store.walletDB.Get(recoveryKey)
+func (store *MockWalletStore) GetRecoveryStatus() (*RecoveryState, error) {
+	rawStatus := store.walletDB.Get(dbm.RecoveryKey)
+	if rawStatus == nil {
+		return nil, ErrGetRecoveryStatus
+	}
+	state := new(RecoveryState)
+	if err := json.Unmarshal(rawStatus, state); err != nil {
+		return nil, err
+	}
+	return state, nil
 }
 
 // GetWalletInfo get wallet information
-- 
2.11.0