wallet/recovery_test.go

   1 package wallet
   2
   3 import (
   4         "fmt"
   5         "io/ioutil"
   6         "math/rand"
   7         "os"
   8         "reflect"
   9         "testing"
  10         "time"
  11
  12         "github.com/vapor/account"
  13         acc "github.com/vapor/account"
  14         "github.com/vapor/blockchain/pseudohsm"
  15         "github.com/vapor/blockchain/signers"
  16         "github.com/vapor/blockchain/txbuilder"
  17         "github.com/vapor/common"
  18         "github.com/vapor/consensus"
  19         "github.com/vapor/crypto/ed25519/chainkd"
  20         dbm "github.com/vapor/database/leveldb"
  21         "github.com/vapor/errors"
  22         "github.com/vapor/protocol/bc"
  23         "github.com/vapor/protocol/bc/types"
  24 )
  25
  26 // MockBlock mock a block
  27 func MockBlock(txs []*types.Tx) *types.Block {
  28         return &types.Block{
  29                 BlockHeader:  types.BlockHeader{Timestamp: uint64(time.Now().Nanosecond())},
  30                 Transactions: txs,
  31         }
  32 }
  33
  34 func MockSimpleUtxo(index uint64, assetID *bc.AssetID, amount uint64, ctrlProg *account.CtrlProgram) *account.UTXO {
  35         if ctrlProg == nil {
  36                 ctrlProg = &account.CtrlProgram{
  37                         AccountID:      "",
  38                         Address:        "",
  39                         KeyIndex:       uint64(0),
  40                         ControlProgram: []byte{81},
  41                         Change:         false,
  42                 }
  43         }
  44
  45         utxo := &account.UTXO{
  46                 OutputID:            bc.Hash{V0: 1},
  47                 SourceID:            bc.Hash{V0: 1},
  48                 AssetID:             *assetID,
  49                 Amount:              amount,
  50                 SourcePos:           index,
  51                 ControlProgram:      ctrlProg.ControlProgram,
  52                 ControlProgramIndex: ctrlProg.KeyIndex,
  53                 AccountID:           ctrlProg.AccountID,
  54                 Address:             ctrlProg.Address,
  55                 ValidHeight:         0,
  56         }
  57
  58         return utxo
  59 }
  60
  61 func AddInput(sourceID bc.Hash, assetID bc.AssetID, amount uint64, pos uint64, controlProgram []byte) *types.TxInput {
  62         return types.NewSpendInput(nil, sourceID, assetID, amount, pos, controlProgram)
  63 }
  64
  65 func AddTxOutput(assetID bc.AssetID, amount uint64, controlProgram []byte) *types.TxOutput {
  66         return types.NewIntraChainOutput(assetID, amount, controlProgram)
  67 }
  68
  69 func BuildTx(baseUtxo *account.UTXO, signer *signers.Signer) (*txbuilder.Template, error) {
  70         tplBuilder, err := CreateTxBuilder(baseUtxo, signer)
  71         if err != nil {
  72                 return nil, err
  73         }
  74
  75         tpl, _, err := tplBuilder.Build()
  76         if err != nil {
  77                 return nil, err
  78         }
  79
  80         return tpl, nil
  81 }
  82
  83 func CreateTxBuilder(baseUtxo *account.UTXO, signer *signers.Signer) (*txbuilder.TemplateBuilder, error) {
  84         tplBuilder := txbuilder.NewBuilder(time.Now())
  85         txInput := AddInput(bc.Hash{V0: 1}, baseUtxo.AssetID, 10000, uint64(1), baseUtxo.ControlProgram)
  86         txOutput := AddTxOutput(baseUtxo.AssetID, 100, baseUtxo.ControlProgram)
  87         tplBuilder.AddInput(txInput, &txbuilder.SigningInstruction{Position: 0})
  88         tplBuilder.AddOutput(txOutput)
  89         return tplBuilder, nil
  90 }
  91
  92 func MockTxsP2PKH(acctMgr *account.Manager, xPub chainkd.XPub, multiTypeAccount bool) ([]*types.Tx, error) {
  93         txs := []*types.Tx{}
  94         accts := []*account.Account{}
  95         for i := uint32(1); i < 32; i = i + 1 + rand.Uint32()%5 {
  96                 alias := fmt.Sprintf("testAccount%d", i)
  97                 deriveRule := signers.BIP0044
  98                 if multiTypeAccount {
  99                         deriveRule = uint8(rand.Uint32() % 2)
 100                 }
 101                 acct, err := account.CreateAccount([]chainkd.XPub{xPub}, 1, alias, uint64(i), deriveRule)
 102                 if err != nil {
 103                         return nil, err
 104                 }
 105
 106                 if err := acctMgr.SaveAccount(acct); err != nil {
 107                         return nil, err
 108                 }
 109
 110                 accts = append(accts, acct)
 111         }
 112
 113         for _, acct := range accts {
 114                 for i := uint32(1); i < 256; i = i + 1 + rand.Uint32()%16 {
 115                         controlProg, err := account.CreateCtrlProgram(acct, uint64(i), false)
 116                         if err != nil {
 117                                 return nil, err
 118                         }
 119
 120                         if err := acctMgr.SaveControlPrograms(controlProg); err != nil {
 121                                 return nil, err
 122                         }
 123
 124                         utxo := MockSimpleUtxo(0, consensus.BTMAssetID, 1000000000, controlProg)
 125                         tpl, err := BuildTx(utxo, acct.Signer)
 126                         if err != nil {
 127                                 return nil, err
 128                         }
 129
 130                         txs = append(txs, tpl.Transaction)
 131                 }
 132         }
 133
 134         return txs, nil
 135 }
 136
 137 func TestXPubsRecoveryLock(t *testing.T) {
 138         dirPath, err := ioutil.TempDir(".", "")
 139         if err != nil {
 140                 t.Fatal(err)
 141         }
 142         defer os.RemoveAll(dirPath)
 143
 144         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 145         walletStore := NewMockWalletStore(testDB)
 146         hsm, err := pseudohsm.New(dirPath)
 147         if err != nil {
 148                 t.Fatal(err)
 149         }
 150
 151         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 152         if err != nil {
 153                 t.Fatal(err)
 154         }
 155
 156         acctStore := NewMockAccountStore(testDB)
 157         acctMgr := account.NewManager(acctStore, nil)
 158         recoveryMgr := NewRecoveryManager(walletStore, acctMgr)
 159         recoveryMgr.state = newRecoveryState()
 160         recoveryMgr.state.XPubs = []chainkd.XPub{xpub.XPub}
 161         recoveryMgr.state.XPubsStatus = newBranchRecoveryState(acctRecoveryWindow)
 162
 163         recoveryMgr.state.StartTime = time.Now()
 164         recoveryMgr.commitStatusInfo()
 165
 166         if err := recoveryMgr.AcctResurrect([]chainkd.XPub{xpub.XPub}); err != nil {
 167                 t.Fatal("TestXPubsRecoveryLock err:", err)
 168         }
 169
 170         if err := recoveryMgr.AcctResurrect([]chainkd.XPub{xpub.XPub}); err != errors.Root(ErrRecoveryBusy) {
 171                 t.Fatal("TestXPubsRecoveryLock err:", err)
 172         }
 173
 174         if err := recoveryMgr.LoadStatusInfo(); err != errors.Root(ErrRecoveryBusy) {
 175                 t.Fatal("TestXPubsRecoveryLock err:", err)
 176         }
 177
 178         recoveryMgr.stopXPubsRec()
 179         if err := recoveryMgr.LoadStatusInfo(); err != nil {
 180                 t.Fatal("TestXPubsRecoveryLock err:", err)
 181         }
 182         recoveryMgr.finished()
 183         if err := recoveryMgr.AcctResurrect([]chainkd.XPub{xpub.XPub}); err != nil {
 184                 t.Fatal("TestXPubsRecoveryLock err:", err)
 185         }
 186 }
 187
 188 func TestExtendScanAddresses(t *testing.T) {
 189         dirPath, err := ioutil.TempDir(".", "")
 190         if err != nil {
 191                 t.Fatal(err)
 192         }
 193         defer os.RemoveAll(dirPath)
 194
 195         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 196         walletStore := NewMockWalletStore(testDB)
 197         hsm, err := pseudohsm.New(dirPath)
 198         if err != nil {
 199                 t.Fatal(err)
 200         }
 201
 202         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 203         if err != nil {
 204                 t.Fatal(err)
 205         }
 206
 207         acctStore := NewMockAccountStore(testDB)
 208         acctMgr := account.NewManager(acctStore, nil)
 209         recoveryMgr := NewRecoveryManager(walletStore, acctMgr)
 210         acc1 := &account.Account{ID: "testA", Alias: "test1", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub.XPub}, KeyIndex: 1, DeriveRule: signers.BIP0044}}
 211         acc2 := &account.Account{ID: "testB", Alias: "test2"}
 212         acc3 := &account.Account{ID: "testC", Alias: "test3", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub.XPub}, KeyIndex: 2, DeriveRule: 3}}
 213         acc4 := &account.Account{ID: "testD", Alias: "test4", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub.XPub}, KeyIndex: 3, DeriveRule: signers.BIP0032}}
 214
 215         recoveryMgr.state.stateForScope(acc1)
 216         recoveryMgr.state.stateForScope(acc3)
 217         recoveryMgr.state.stateForScope(acc4)
 218
 219         cases := []struct {
 220                 acct       *account.Account
 221                 err        error
 222                 addressLen uint64
 223         }{
 224                 {acc1, nil, addrRecoveryWindow * 2},
 225                 {acc2, ErrInvalidAcctID, addrRecoveryWindow * 2},
 226                 {acc3, signers.ErrDeriveRule, addrRecoveryWindow * 2},
 227                 {acc4, nil, addrRecoveryWindow * 3},
 228         }
 229
 230         for _, c := range cases {
 231                 if err := recoveryMgr.extendScanAddresses(c.acct.ID, true); err != c.err {
 232                         t.Fatal("extend scan addresses err:", err)
 233                 }
 234
 235                 if err := recoveryMgr.extendScanAddresses(c.acct.ID, false); err != c.err {
 236                         t.Fatal("extend scan addresses err:", err)
 237                 }
 238
 239                 if uint64(len(recoveryMgr.addresses)) != c.addressLen {
 240                         t.Fatalf("extend scan addresses err: len:%d,want:%d", len(recoveryMgr.addresses), c.addressLen)
 241                 }
 242         }
 243 }
 244
 245 func TestRecoveryFromXPubs(t *testing.T) {
 246         dirPath, err := ioutil.TempDir(".", "")
 247         if err != nil {
 248                 t.Fatal(err)
 249         }
 250         defer os.RemoveAll(dirPath)
 251
 252         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 253         recoveryDB := dbm.NewDB("recdb", "leveldb", dirPath)
 254         recoveryStore := NewMockWalletStore(recoveryDB)
 255         hsm, err := pseudohsm.New(dirPath)
 256         if err != nil {
 257                 t.Fatal(err)
 258         }
 259
 260         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 261         if err != nil {
 262                 t.Fatal(err)
 263         }
 264
 265         acctStore := NewMockAccountStore(testDB)
 266         acctMgr := account.NewManager(acctStore, nil)
 267         txs, err := MockTxsP2PKH(acctMgr, xpub.XPub, false)
 268         recActStore := NewMockAccountStore(recoveryDB)
 269         recAcctMgr := account.NewManager(recActStore, nil)
 270         recoveryMgr := NewRecoveryManager(recoveryStore, recAcctMgr)
 271
 272         cases := []struct {
 273                 xPubs []chainkd.XPub
 274                 err   error
 275         }{
 276                 {[]chainkd.XPub{xpub.XPub}, nil},
 277                 {[]chainkd.XPub{xpub.XPub, xpub.XPub}, signers.ErrDupeXPub},
 278                 {[]chainkd.XPub{}, signers.ErrNoXPubs},
 279         }
 280
 281         for _, c := range cases {
 282                 if err := recoveryMgr.AcctResurrect(c.xPubs); errors.Root(err) != c.err {
 283                         t.Fatal("recovery from XPubs err:", err)
 284                 }
 285
 286                 if err != nil {
 287                         recoveryMgr.finished()
 288                         continue
 289                 }
 290                 if err := recoveryMgr.FilterRecoveryTxs(MockBlock(txs)); err != nil {
 291                         t.Fatal("recovery from XPubs err:", err)
 292                 }
 293
 294                 Accounts, err := acctMgr.ListAccounts("")
 295                 if err != nil {
 296                         t.Fatal("recovery from XPubs err:", err)
 297                 }
 298
 299                 for _, acct := range Accounts {
 300                         tmp, err := recAcctMgr.GetAccountByXPubsIndex(acct.XPubs, acct.KeyIndex)
 301                         if err != nil && err != acc.ErrFindAccount {
 302                                 t.Fatal("recovery from XPubs err:", err)
 303                         }
 304
 305                         if tmp == nil {
 306                                 t.Fatal("accout recovery from xpubs err:", acct.KeyIndex)
 307                         }
 308
 309                         if acctMgr.GetBip44ContractIndex(acct.ID, true) != recAcctMgr.GetBip44ContractIndex(tmp.ID, true) {
 310                                 t.Fatal("bip44 internal address index recovery from xpubs err")
 311                         }
 312
 313                         if acctMgr.GetBip44ContractIndex(acct.ID, false) != recAcctMgr.GetBip44ContractIndex(tmp.ID, false) {
 314                                 t.Fatal("bip44 external address index recovery from xpubs err")
 315                         }
 316                 }
 317
 318                 recoveryMgr.finished()
 319         }
 320 }
 321
 322 func TestRecoveryByRescanAccount(t *testing.T) {
 323         dirPath, err := ioutil.TempDir(".", "")
 324         if err != nil {
 325                 t.Fatal(err)
 326         }
 327         defer os.RemoveAll(dirPath)
 328
 329         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 330         recoveryDB := dbm.NewDB("recdb", "leveldb", dirPath)
 331         recoveryStore := NewMockWalletStore(recoveryDB)
 332         hsm, err := pseudohsm.New(dirPath)
 333         if err != nil {
 334                 t.Fatal(err)
 335         }
 336
 337         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 338         if err != nil {
 339                 t.Fatal(err)
 340         }
 341
 342         acctStore := NewMockAccountStore(testDB)
 343         acctMgr := account.NewManager(acctStore, nil)
 344         txs, err := MockTxsP2PKH(acctMgr, xpub.XPub, true)
 345         if err != nil {
 346                 t.Fatal("recovery by rescan account err:", err)
 347         }
 348
 349         allAccounts, err := acctMgr.ListAccounts("")
 350         if err != nil {
 351                 t.Fatal("recovery by rescan account err:", err)
 352         }
 353
 354         recActStore := NewMockAccountStore(recoveryDB)
 355         recAcctMgr := account.NewManager(recActStore, nil)
 356         for _, acct := range allAccounts {
 357                 if err := recAcctMgr.SaveAccount(acct); err != nil {
 358                         t.Fatal("recovery by rescan account err:", err)
 359                 }
 360         }
 361
 362         recoveryMgr := NewRecoveryManager(recoveryStore, recAcctMgr)
 363
 364         acct := &account.Account{ID: "testA", Alias: "test1", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub.XPub}, KeyIndex: 1, DeriveRule: 3}}
 365
 366         cases := []struct {
 367                 accounts []*account.Account
 368                 err      error
 369         }{
 370                 {allAccounts, nil},
 371                 {[]*account.Account{acct}, signers.ErrDeriveRule},
 372         }
 373
 374         for _, c := range cases {
 375                 if err := recoveryMgr.AddrResurrect(c.accounts); errors.Root(err) != c.err {
 376                         t.Fatal("recovery by rescan account err:", err)
 377                 }
 378
 379                 if err != nil {
 380                         continue
 381                 }
 382                 recoveryMgr.FilterRecoveryTxs(MockBlock(txs))
 383                 accounts, err := acctMgr.ListAccounts("")
 384                 if err != nil {
 385                         t.Fatal("recovery from XPubs err:", err)
 386                 }
 387
 388                 for _, acct := range accounts {
 389                         tmp, err := recAcctMgr.GetAccountByXPubsIndex(acct.XPubs, acct.KeyIndex)
 390                         if err != nil && err != acc.ErrFindAccount {
 391                                 t.Fatal("recovery from XPubs err:", err)
 392                         }
 393
 394                         if tmp == nil {
 395                                 t.Fatal("accout recovery from xpubs err:", acct.KeyIndex)
 396                         }
 397
 398                         if acctMgr.GetBip44ContractIndex(acct.ID, true) != recAcctMgr.GetBip44ContractIndex(tmp.ID, true) {
 399                                 t.Fatal("bip44 internal address index recovery from xpubs err")
 400                         }
 401
 402                         if acctMgr.GetBip44ContractIndex(acct.ID, false) != recAcctMgr.GetBip44ContractIndex(tmp.ID, false) {
 403                                 t.Fatal("bip44 external address index recovery from xpubs err")
 404                         }
 405
 406                         if acctMgr.GetContractIndex(acct.ID) != recAcctMgr.GetContractIndex(tmp.ID) {
 407                                 t.Fatal("bip32 address index recovery from xpubs err")
 408                         }
 409                 }
 410         }
 411
 412 }
 413
 414 func TestReportFound(t *testing.T) {
 415         dirPath, err := ioutil.TempDir(".", "")
 416         if err != nil {
 417                 t.Fatal(err)
 418         }
 419         defer os.RemoveAll(dirPath)
 420
 421         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 422         testStore := NewMockWalletStore(testDB)
 423         hsm, err := pseudohsm.New(dirPath)
 424         if err != nil {
 425                 t.Fatal(err)
 426         }
 427
 428         xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
 429         if err != nil {
 430                 t.Fatal(err)
 431         }
 432
 433         xpub2, _, err := hsm.XCreate("test_pub2", "password", "en")
 434         if err != nil {
 435                 t.Fatal(err)
 436         }
 437
 438         acctStore := NewMockAccountStore(testDB)
 439         acctMgr := account.NewManager(acctStore, nil)
 440         recoveryMgr := NewRecoveryManager(testStore, acctMgr)
 441         acc1 := &account.Account{ID: "testA", Alias: "test1", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub1.XPub}, KeyIndex: 1, DeriveRule: signers.BIP0044}}
 442         acc2 := &account.Account{ID: "testB", Alias: "test2", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub2.XPub}, KeyIndex: 1, DeriveRule: signers.BIP0032}}
 443         acc3 := &account.Account{ID: "testC", Alias: "test3", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub2.XPub}, KeyIndex: 2, DeriveRule: signers.BIP0044}}
 444
 445         cp1 := &account.CtrlProgram{AccountID: acc1.ID, Address: "address1", KeyIndex: 10, Change: false}
 446         cp2 := &account.CtrlProgram{AccountID: acc1.ID, Address: "address1", KeyIndex: 20, Change: true}
 447         cp3 := &account.CtrlProgram{AccountID: acc2.ID, Address: "address1", KeyIndex: 30, Change: false}
 448         cp4 := &account.CtrlProgram{AccountID: acc2.ID, Address: "address1", KeyIndex: 40, Change: true}
 449         cp5 := &account.CtrlProgram{AccountID: acc3.ID, Address: "address1", KeyIndex: 50, Change: false}
 450         cp6 := &account.CtrlProgram{AccountID: acc3.ID, Address: "address1", KeyIndex: 60, Change: true}
 451
 452         if err := acctMgr.SaveAccount(acc2); err != nil {
 453                 t.Fatal("ReportFound test err:", err)
 454         }
 455
 456         if err := acctMgr.SaveAccount(acc3); err != nil {
 457                 t.Fatal("ReportFound test err:", err)
 458         }
 459
 460         recoveryMgr.state.XPubsStatus = newBranchRecoveryState(acctRecoveryWindow)
 461         recoveryMgr.state.XPubs = []chainkd.XPub{xpub1.XPub}
 462         recoveryMgr.state.stateForScope(acc1)
 463         recoveryMgr.state.stateForScope(acc2)
 464         recoveryMgr.state.stateForScope(acc3)
 465
 466         cases := []struct {
 467                 acct   *account.Account
 468                 cp     *account.CtrlProgram
 469                 err    error
 470                 status *addressRecoveryState
 471         }{
 472                 {acc1, cp1, nil,
 473                         &addressRecoveryState{InternalBranch: &branchRecoveryState{addrRecoveryWindow, 1, 1}, ExternalBranch: &branchRecoveryState{addrRecoveryWindow, 139, 11}, Account: acc1}},
 474                 {acc2, cp3, nil,
 475                         &addressRecoveryState{InternalBranch: &branchRecoveryState{addrRecoveryWindow, 1, 1}, ExternalBranch: &branchRecoveryState{addrRecoveryWindow, 159, 31}, Account: acc2}},
 476                 {acc1, cp2, nil,
 477                         &addressRecoveryState{InternalBranch: &branchRecoveryState{addrRecoveryWindow, 149, 21}, ExternalBranch: &branchRecoveryState{addrRecoveryWindow, 139, 11}, Account: acc1}},
 478                 {acc2, cp4, nil,
 479                         &addressRecoveryState{InternalBranch: &branchRecoveryState{addrRecoveryWindow, 169, 41}, ExternalBranch: &branchRecoveryState{addrRecoveryWindow, 159, 31}, Account: acc2}},
 480                 {acc3, cp5, nil,
 481                         &addressRecoveryState{InternalBranch: &branchRecoveryState{addrRecoveryWindow, 1, 1}, ExternalBranch: &branchRecoveryState{addrRecoveryWindow, 179, 51}, Account: acc3}},
 482                 {acc3, cp6, nil,
 483                         &addressRecoveryState{InternalBranch: &branchRecoveryState{addrRecoveryWindow, 189, 61}, ExternalBranch: &branchRecoveryState{addrRecoveryWindow, 179, 51}, Account: acc3}},
 484         }
 485
 486         for _, c := range cases {
 487                 if err := recoveryMgr.reportFound(c.acct, c.cp); err != c.err {
 488                         t.Fatal("ReportFound test err:", err, c.acct.ID)
 489                 }
 490
 491                 status, ok := recoveryMgr.state.AccountsStatus[c.acct.ID]
 492                 if !ok {
 493                         t.Fatal("ReportFound test err: can not find status")
 494                 }
 495                 if !reflect.DeepEqual(status, c.status) {
 496                         t.Log(c.status.Account, c.status.InternalBranch, c.status.ExternalBranch)
 497                         t.Log(status.Account, status.InternalBranch, status.ExternalBranch)
 498                         t.Fatal("ReportFound test err: recovery status error")
 499                 }
 500         }
 501 }
 502
 503 func TestLoadStatusInfo(t *testing.T) {
 504         dirPath, err := ioutil.TempDir(".", "")
 505         if err != nil {
 506                 t.Fatal(err)
 507         }
 508         defer os.RemoveAll(dirPath)
 509
 510         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 511         testStore := NewMockWalletStore(testDB)
 512         defer os.RemoveAll("temp")
 513
 514         hsm, err := pseudohsm.New(dirPath)
 515         if err != nil {
 516                 t.Fatal(err)
 517         }
 518
 519         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 520         if err != nil {
 521                 t.Fatal(err)
 522         }
 523
 524         acctStore := NewMockAccountStore(testDB)
 525         acctMgr := account.NewManager(acctStore, nil)
 526         recoveryMgr := NewRecoveryManager(testStore, acctMgr)
 527         // StatusInit init recovery status manager.
 528         recoveryMgr.state = newRecoveryState()
 529         recoveryMgr.state.XPubs = []chainkd.XPub{xpub.XPub}
 530         recoveryMgr.state.XPubsStatus = newBranchRecoveryState(acctRecoveryWindow)
 531
 532         recoveryMgr.state.StartTime = time.Now()
 533         if err := recoveryMgr.LoadStatusInfo(); err != nil {
 534                 t.Fatal("TestLoadStatusInfo err:", err)
 535         }
 536
 537         recoveryMgr.commitStatusInfo()
 538
 539         recoveryMgrRestore := NewRecoveryManager(testStore, acctMgr)
 540         if err := recoveryMgrRestore.LoadStatusInfo(); err != nil {
 541                 t.Fatal("TestLoadStatusInfo err:", err)
 542         }
 543
 544         if !reflect.DeepEqual(recoveryMgrRestore.state.XPubsStatus, recoveryMgr.state.XPubsStatus) {
 545                 t.Fatalf("TestLoadStatusInfo XPubsStatus reload err")
 546         }
 547
 548         if !reflect.DeepEqual(recoveryMgrRestore.state.XPubs, recoveryMgr.state.XPubs) {
 549                 t.Fatalf("TestLoadStatusInfo XPubs reload err")
 550         }
 551
 552         if !reflect.DeepEqual(recoveryMgrRestore.state.AccountsStatus, recoveryMgr.state.AccountsStatus) {
 553                 t.Fatalf("TestLoadStatusInfo AccountsStatus reload err")
 554         }
 555
 556         if !recoveryMgrRestore.state.StartTime.Equal(recoveryMgr.state.StartTime) {
 557                 t.Fatalf("TestLoadStatusInfo StartTime reload err")
 558         }
 559
 560         acct := &account.Account{ID: "testA", Alias: "test1", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub.XPub}, KeyIndex: 1, DeriveRule: 3}}
 561         recoveryMgr.state.AccountsStatus[acct.ID] = newAddressRecoveryState(addrRecoveryWindow, acct)
 562         if err := recoveryMgr.commitStatusInfo(); err != nil {
 563                 t.Fatal("TestLoadStatusInfo err:", err)
 564         }
 565         if err := recoveryMgr.LoadStatusInfo(); err == nil {
 566                 t.Fatal("TestLoadStatusInfo err")
 567         }
 568
 569         recoveryMgr.state = nil
 570         if err := recoveryMgr.commitStatusInfo(); err != nil {
 571                 t.Fatal("TestLoadStatusInfo err:", err)
 572         }
 573
 574         if err := recoveryMgr.LoadStatusInfo(); err == nil {
 575                 t.Fatal("TestLoadStatusInfo err")
 576         }
 577 }
 578
 579 func TestLock(t *testing.T) {
 580         dirPath, err := ioutil.TempDir(".", "")
 581         if err != nil {
 582                 t.Fatal(err)
 583         }
 584         defer os.RemoveAll(dirPath)
 585
 586         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 587         testStore := NewMockWalletStore(testDB)
 588         defer os.RemoveAll("temp")
 589
 590         acctStore := NewMockAccountStore(testDB)
 591         acctMgr := account.NewManager(acctStore, nil)
 592         recoveryMgr := NewRecoveryManager(testStore, acctMgr)
 593         if !recoveryMgr.tryStartXPubsRec() {
 594                 t.Fatal("recovery manager try lock test err")
 595         }
 596
 597         if recoveryMgr.tryStartXPubsRec() {
 598                 t.Fatal("recovery manager relock test err")
 599         }
 600
 601         recoveryMgr.stopXPubsRec()
 602
 603         if !recoveryMgr.tryStartXPubsRec() {
 604                 t.Fatal("recovery manager try lock test err")
 605         }
 606 }
 607
 608 func TestStateForScope(t *testing.T) {
 609         state := newRecoveryState()
 610         acc1 := &account.Account{ID: "test1", Alias: "testA"}
 611         state.stateForScope(acc1)
 612         if !reflect.DeepEqual(state.AccountsStatus[acc1.ID].Account, acc1) {
 613                 t.Fatal("state for scope test err")
 614         }
 615
 616         acc2 := &account.Account{ID: "test1", Alias: "testB"}
 617         state.stateForScope(acc2)
 618
 619         if reflect.DeepEqual(state.AccountsStatus[acc2.ID].Account, acc2) {
 620                 t.Fatal("state for scope test err")
 621         }
 622
 623         acc3 := &account.Account{ID: "test2", Alias: "testC"}
 624         state.stateForScope(acc3)
 625         if !reflect.DeepEqual(state.AccountsStatus[acc3.ID].Account, acc3) {
 626                 t.Fatal("state for scope test err")
 627         }
 628 }
 629
 630 func TestContractIndexResidue(t *testing.T) {
 631         dirPath, err := ioutil.TempDir(".", "")
 632         if err != nil {
 633                 t.Fatal(err)
 634         }
 635         defer os.RemoveAll(dirPath)
 636
 637         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 638         testStore := NewMockWalletStore(testDB)
 639         hsm, err := pseudohsm.New(dirPath)
 640         if err != nil {
 641                 t.Fatal(err)
 642         }
 643
 644         xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
 645         if err != nil {
 646                 t.Fatal(err)
 647         }
 648
 649         contractIndexResidue := uint64(5)
 650         acctStore := NewMockAccountStore(testDB)
 651         acctMgr := account.NewManager(acctStore, nil)
 652         recoveryMgr := NewRecoveryManager(testStore, acctMgr)
 653         acct := &account.Account{ID: "testA", Alias: "test1", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub1.XPub}, KeyIndex: 1, DeriveRule: signers.BIP0044}}
 654
 655         cp1 := &account.CtrlProgram{AccountID: acct.ID, Address: "address1", KeyIndex: 10, Change: false}
 656
 657         setContractIndexKey := func(acctMgr *account.Manager, accountID string, change bool) {
 658                 testDB.Set(Bip44ContractIndexKey(accountID, change), common.Unit64ToBytes(contractIndexResidue))
 659         }
 660
 661         delAccount := func(acctMgr *account.Manager, accountID string, change bool) {
 662                 acctMgr.DeleteAccount(accountID)
 663         }
 664
 665         recoveryMgr.state.XPubsStatus = newBranchRecoveryState(acctRecoveryWindow)
 666         recoveryMgr.state.XPubs = []chainkd.XPub{xpub1.XPub}
 667         recoveryMgr.state.stateForScope(acct)
 668
 669         cases := []struct {
 670                 acct       *account.Account
 671                 cp         *account.CtrlProgram
 672                 preProcess func(acctMgr *account.Manager, accountID string, change bool)
 673                 err        error
 674                 wantCPNum  uint64
 675         }{
 676                 {acct, cp1, setContractIndexKey, nil, 5},
 677                 {acct, cp1, delAccount, nil, 10},
 678         }
 679
 680         for _, c := range cases {
 681                 if c.preProcess != nil {
 682                         c.preProcess(acctMgr, c.acct.ID, c.cp.Change)
 683                 }
 684
 685                 if err := acctMgr.SaveAccount(acct); err != nil {
 686                         t.Fatal("ReportFound test err:", err)
 687                 }
 688
 689                 if err := recoveryMgr.reportFound(c.acct, c.cp); err != c.err {
 690                         t.Fatal("ContractIndexResidue test err:", err, c.acct.ID)
 691                 }
 692                 cps, err := acctMgr.ListControlProgram()
 693                 if err != nil {
 694                         t.Fatal("list control program err:", err)
 695                 }
 696
 697                 cpNum := uint64(0)
 698                 for _, cp := range cps {
 699                         if cp.Address == "" || cp.AccountID != c.acct.ID {
 700                                 continue
 701                         }
 702                         cpNum++
 703                 }
 704
 705                 if cpNum != c.wantCPNum {
 706                         t.Fatal("Test contract index residue cp num err want:", c.wantCPNum, " got:", cpNum)
 707                 }
 708         }
 709 }