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         dbm "github.com/tendermint/tmlibs/db"
  13
  14         "github.com/vapor/account"
  15         "github.com/vapor/blockchain/pseudohsm"
  16         "github.com/vapor/blockchain/signers"
  17         "github.com/vapor/blockchain/txbuilder"
  18         "github.com/vapor/consensus"
  19         "github.com/vapor/crypto/ed25519/chainkd"
  20         "github.com/vapor/protocol/bc"
  21         "github.com/vapor/protocol/bc/types"
  22 )
  23
  24 // MockBlock mock a block
  25 func MockBlock(txs []*types.Tx) *types.Block {
  26         return &types.Block{
  27                 BlockHeader:  types.BlockHeader{Timestamp: uint64(time.Now().Nanosecond())},
  28                 Transactions: txs,
  29         }
  30 }
  31
  32 func MockSimpleUtxo(index uint64, assetID *bc.AssetID, amount uint64, ctrlProg *account.CtrlProgram) *account.UTXO {
  33         if ctrlProg == nil {
  34                 ctrlProg = &account.CtrlProgram{
  35                         AccountID:      "",
  36                         Address:        "",
  37                         KeyIndex:       uint64(0),
  38                         ControlProgram: []byte{81},
  39                         Change:         false,
  40                 }
  41         }
  42
  43         utxo := &account.UTXO{
  44                 OutputID:            bc.Hash{V0: 1},
  45                 SourceID:            bc.Hash{V0: 1},
  46                 AssetID:             *assetID,
  47                 Amount:              amount,
  48                 SourcePos:           index,
  49                 ControlProgram:      ctrlProg.ControlProgram,
  50                 ControlProgramIndex: ctrlProg.KeyIndex,
  51                 AccountID:           ctrlProg.AccountID,
  52                 Address:             ctrlProg.Address,
  53                 ValidHeight:         0,
  54         }
  55
  56         return utxo
  57 }
  58
  59 func AddTxOutput(assetID bc.AssetID, amount uint64, controlProgram []byte) *types.TxOutput {
  60         out := types.NewTxOutput(assetID, amount, controlProgram)
  61         return out
  62 }
  63
  64 func BuildTx(baseUtxo *account.UTXO, signer *signers.Signer) (*txbuilder.Template, error) {
  65         tplBuilder, err := CreateTxBuilder(baseUtxo, signer)
  66         if err != nil {
  67                 return nil, err
  68         }
  69
  70         tpl, _, err := tplBuilder.Build()
  71         if err != nil {
  72                 return nil, err
  73         }
  74
  75         return tpl, nil
  76 }
  77
  78 func CreateTxBuilder(baseUtxo *account.UTXO, signer *signers.Signer) (*txbuilder.TemplateBuilder, error) {
  79         tplBuilder := txbuilder.NewBuilder(time.Now())
  80         txOutput := AddTxOutput(baseUtxo.AssetID, 100, baseUtxo.ControlProgram)
  81         tplBuilder.AddOutput(txOutput)
  82         return tplBuilder, nil
  83 }
  84
  85 func MockTxsP2PKH(acctMgr *account.Manager, xPub chainkd.XPub, multiTypeAccount bool) ([]*types.Tx, error) {
  86         txs := []*types.Tx{}
  87         accts := []*account.Account{}
  88         for i := uint32(1); i < 32; i = i + 1 + rand.Uint32()%5 {
  89                 alias := fmt.Sprintf("testAccount%d", i)
  90                 deriveRule := signers.BIP0044
  91                 if multiTypeAccount {
  92                         deriveRule = uint8(rand.Uint32() % 2)
  93                 }
  94                 acct, err := account.CreateAccount([]chainkd.XPub{xPub}, 1, alias, uint64(i), deriveRule)
  95                 if err != nil {
  96                         return nil, err
  97                 }
  98
  99                 if err := acctMgr.SaveAccount(acct); err != nil {
 100                         return nil, err
 101                 }
 102
 103                 accts = append(accts, acct)
 104         }
 105
 106         for _, acct := range accts {
 107                 for i := uint32(1); i < 256; i = i + 1 + rand.Uint32()%16 {
 108                         controlProg, err := account.CreateCtrlProgram(acct, uint64(i), false)
 109                         if err != nil {
 110                                 return nil, err
 111                         }
 112
 113                         if err := acctMgr.SaveControlPrograms(controlProg); err != nil {
 114                                 return nil, err
 115                         }
 116
 117                         utxo := MockSimpleUtxo(0, consensus.BTMAssetID, 1000000000, controlProg)
 118                         tpl, err := BuildTx(utxo, acct.Signer)
 119                         if err != nil {
 120                                 return nil, err
 121                         }
 122
 123                         txs = append(txs, tpl.Transaction)
 124                 }
 125         }
 126
 127         return txs, nil
 128 }
 129
 130 func TestExtendScanAddresses(t *testing.T) {
 131         dirPath, err := ioutil.TempDir(".", "")
 132         if err != nil {
 133                 t.Fatal(err)
 134         }
 135         defer os.RemoveAll(dirPath)
 136
 137         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 138         hsm, err := pseudohsm.New(dirPath)
 139         if err != nil {
 140                 t.Fatal(err)
 141         }
 142
 143         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 144         if err != nil {
 145                 t.Fatal(err)
 146         }
 147
 148         acctMgr := account.NewManager(testDB, nil)
 149         recoveryMgr := newRecoveryManager(testDB, acctMgr)
 150         acc1 := &account.Account{ID: "testA", Alias: "test1", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub.XPub}, KeyIndex: 1, DeriveRule: signers.BIP0044}}
 151         acc2 := &account.Account{ID: "testB", Alias: "test2"}
 152         acc3 := &account.Account{ID: "testC", Alias: "test3", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub.XPub}, KeyIndex: 2, DeriveRule: 3}}
 153         acc4 := &account.Account{ID: "testD", Alias: "test4", Signer: &signers.Signer{XPubs: []chainkd.XPub{xpub.XPub}, KeyIndex: 3, DeriveRule: signers.BIP0032}}
 154
 155         recoveryMgr.state.stateForScope(acc1)
 156         recoveryMgr.state.stateForScope(acc3)
 157         recoveryMgr.state.stateForScope(acc4)
 158
 159         cases := []struct {
 160                 acct       *account.Account
 161                 err        error
 162                 addressLen uint64
 163         }{
 164                 {acc1, nil, addrRecoveryWindow * 2},
 165                 {acc2, ErrInvalidAcctID, addrRecoveryWindow * 2},
 166                 {acc3, signers.ErrDeriveRule, addrRecoveryWindow * 2},
 167                 {acc4, nil, addrRecoveryWindow * 3},
 168         }
 169
 170         for _, c := range cases {
 171                 if err := recoveryMgr.extendScanAddresses(c.acct.ID, true); err != c.err {
 172                         t.Fatal("extend scan addresses err:", err)
 173                 }
 174
 175                 if err := recoveryMgr.extendScanAddresses(c.acct.ID, false); err != c.err {
 176                         t.Fatal("extend scan addresses err:", err)
 177                 }
 178
 179                 if uint64(len(recoveryMgr.addresses)) != c.addressLen {
 180                         t.Fatalf("extend scan addresses err: len:%d,want:%d", len(recoveryMgr.addresses), c.addressLen)
 181                 }
 182         }
 183 }
 184
 185 func TestRecoveryFromXPubs(t *testing.T) {
 186         dirPath, err := ioutil.TempDir(".", "")
 187         if err != nil {
 188                 t.Fatal(err)
 189         }
 190         defer os.RemoveAll(dirPath)
 191
 192         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 193         recoveryDB := dbm.NewDB("recdb", "leveldb", dirPath)
 194         hsm, err := pseudohsm.New(dirPath)
 195         if err != nil {
 196                 t.Fatal(err)
 197         }
 198
 199         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 200         if err != nil {
 201                 t.Fatal(err)
 202         }
 203
 204         acctMgr := account.NewManager(testDB, nil)
 205         txs, err := MockTxsP2PKH(acctMgr, xpub.XPub, false)
 206         recAcctMgr := account.NewManager(recoveryDB, nil)
 207         recoveryMgr := newRecoveryManager(recoveryDB, recAcctMgr)
 208         if err := recoveryMgr.AcctResurrect([]chainkd.XPub{xpub.XPub}); err != nil {
 209                 t.Fatal("recovery from XPubs err:", err)
 210         }
 211
 212         if err := recoveryMgr.FilterRecoveryTxs(MockBlock(txs)); err != nil {
 213                 t.Fatal("recovery from XPubs err:", err)
 214         }
 215
 216         Accounts, err := acctMgr.ListAccounts("")
 217         if err != nil {
 218                 t.Fatal("recovery from XPubs err:", err)
 219         }
 220
 221         for _, acct := range Accounts {
 222                 tmp, err := recAcctMgr.GetAccountByXPubsIndex(acct.XPubs, acct.KeyIndex)
 223                 if err != nil {
 224                         t.Fatal("recovery from XPubs err:", err)
 225                 }
 226
 227                 if tmp == nil {
 228                         t.Fatal("accout recovery from xpubs err:", acct.KeyIndex)
 229                 }
 230
 231                 if acctMgr.GetBip44ContractIndex(acct.ID, true) != recAcctMgr.GetBip44ContractIndex(tmp.ID, true) {
 232                         t.Fatal("bip44 internal address index recovery from xpubs err")
 233                 }
 234
 235                 if acctMgr.GetBip44ContractIndex(acct.ID, false) != recAcctMgr.GetBip44ContractIndex(tmp.ID, false) {
 236                         t.Fatal("bip44 external address index recovery from xpubs err")
 237                 }
 238         }
 239 }
 240
 241 func TestRecoveryByRescanAccount(t *testing.T) {
 242         dirPath, err := ioutil.TempDir(".", "")
 243         if err != nil {
 244                 t.Fatal(err)
 245         }
 246         defer os.RemoveAll(dirPath)
 247
 248         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 249         recoveryDB := dbm.NewDB("recdb", "leveldb", dirPath)
 250         hsm, err := pseudohsm.New(dirPath)
 251         if err != nil {
 252                 t.Fatal(err)
 253         }
 254
 255         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 256         if err != nil {
 257                 t.Fatal(err)
 258         }
 259
 260         acctMgr := account.NewManager(testDB, nil)
 261         txs, err := MockTxsP2PKH(acctMgr, xpub.XPub, true)
 262         if err != nil {
 263                 t.Fatal("recovery by rescan account err:", err)
 264         }
 265
 266         allAccounts, err := acctMgr.ListAccounts("")
 267         if err != nil {
 268                 t.Fatal("recovery by rescan account err:", err)
 269         }
 270
 271         recAcctMgr := account.NewManager(recoveryDB, nil)
 272         for _, acct := range allAccounts {
 273                 if err := recAcctMgr.SaveAccount(acct); err != nil {
 274                         t.Fatal("recovery by rescan account err:", err)
 275                 }
 276         }
 277
 278         recoveryMgr := newRecoveryManager(recoveryDB, recAcctMgr)
 279         if err := recoveryMgr.AddrResurrect(allAccounts); err != nil {
 280                 t.Fatal("recovery by rescan account err:", err)
 281         }
 282
 283         recoveryMgr.FilterRecoveryTxs(MockBlock(txs))
 284         Accounts, err := acctMgr.ListAccounts("")
 285         for _, acct := range Accounts {
 286                 tmp, err := recAcctMgr.GetAccountByXPubsIndex(acct.XPubs, acct.KeyIndex)
 287                 if err != nil {
 288                         t.Fatal("recovery from XPubs err:", err)
 289                 }
 290
 291                 if tmp == nil {
 292                         t.Fatal("accout recovery from xpubs err:", acct.KeyIndex)
 293                 }
 294
 295                 if acctMgr.GetBip44ContractIndex(acct.ID, true) != recAcctMgr.GetBip44ContractIndex(tmp.ID, true) {
 296                         t.Fatal("bip44 internal address index recovery from xpubs err")
 297                 }
 298
 299                 if acctMgr.GetBip44ContractIndex(acct.ID, false) != recAcctMgr.GetBip44ContractIndex(tmp.ID, false) {
 300                         t.Fatal("bip44 external address index recovery from xpubs err")
 301                 }
 302
 303                 if acctMgr.GetContractIndex(acct.ID) != recAcctMgr.GetContractIndex(tmp.ID) {
 304                         t.Fatal("bip32 address index recovery from xpubs err")
 305                 }
 306         }
 307 }
 308
 309 func TestLoadStatusInfo(t *testing.T) {
 310         dirPath, err := ioutil.TempDir(".", "")
 311         if err != nil {
 312                 t.Fatal(err)
 313         }
 314         defer os.RemoveAll(dirPath)
 315
 316         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 317         defer os.RemoveAll("temp")
 318
 319         hsm, err := pseudohsm.New(dirPath)
 320         if err != nil {
 321                 t.Fatal(err)
 322         }
 323
 324         xpub, _, err := hsm.XCreate("test_pub", "password", "en")
 325         if err != nil {
 326                 t.Fatal(err)
 327         }
 328
 329         acctMgr := account.NewManager(testDB, nil)
 330         recoveryMgr := newRecoveryManager(testDB, acctMgr)
 331         // StatusInit init recovery status manager.
 332         recoveryMgr.state = newRecoveryState()
 333         recoveryMgr.state.XPubs = []chainkd.XPub{xpub.XPub}
 334         recoveryMgr.state.XPubsStatus = newBranchRecoveryState(acctRecoveryWindow)
 335
 336         recoveryMgr.state.StartTime = time.Now()
 337         recoveryMgr.commitStatusInfo()
 338
 339         recoveryMgrRestore := newRecoveryManager(testDB, acctMgr)
 340         recoveryMgrRestore.LoadStatusInfo()
 341
 342         if !reflect.DeepEqual(recoveryMgrRestore.state.XPubsStatus, recoveryMgr.state.XPubsStatus) {
 343                 t.Fatalf("testLoadStatusInfo XPubsStatus reload err")
 344         }
 345
 346         if !reflect.DeepEqual(recoveryMgrRestore.state.XPubs, recoveryMgr.state.XPubs) {
 347                 t.Fatalf("testLoadStatusInfo XPubs reload err")
 348         }
 349
 350         if !reflect.DeepEqual(recoveryMgrRestore.state.AccountsStatus, recoveryMgr.state.AccountsStatus) {
 351                 t.Fatalf("testLoadStatusInfo AccountsStatus reload err")
 352         }
 353
 354         if !recoveryMgrRestore.state.StartTime.Equal(recoveryMgr.state.StartTime) {
 355                 t.Fatalf("testLoadStatusInfo StartTime reload err")
 356         }
 357 }
 358
 359 func TestLock(t *testing.T) {
 360         dirPath, err := ioutil.TempDir(".", "")
 361         if err != nil {
 362                 t.Fatal(err)
 363         }
 364         defer os.RemoveAll(dirPath)
 365
 366         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 367         defer os.RemoveAll("temp")
 368
 369         acctMgr := account.NewManager(testDB, nil)
 370         recoveryMgr := newRecoveryManager(testDB, acctMgr)
 371         if !recoveryMgr.tryStartXPubsRec() {
 372                 t.Fatal("recovery manager try lock test err")
 373         }
 374
 375         if recoveryMgr.tryStartXPubsRec() {
 376                 t.Fatal("recovery manager relock test err")
 377         }
 378
 379         recoveryMgr.stopXPubsRec()
 380
 381         if !recoveryMgr.tryStartXPubsRec() {
 382                 t.Fatal("recovery manager try lock test err")
 383         }
 384 }
 385
 386 func TestStateForScope(t *testing.T) {
 387         state := newRecoveryState()
 388         acc1 := &account.Account{ID: "test1", Alias: "testA"}
 389         state.stateForScope(acc1)
 390         if !reflect.DeepEqual(state.AccountsStatus[acc1.ID].Account, acc1) {
 391                 t.Fatal("state for scope test err")
 392         }
 393
 394         acc2 := &account.Account{ID: "test1", Alias: "testB"}
 395         state.stateForScope(acc2)
 396
 397         if reflect.DeepEqual(state.AccountsStatus[acc2.ID].Account, acc2) {
 398                 t.Fatal("state for scope test err")
 399         }
 400
 401         acc3 := &account.Account{ID: "test2", Alias: "testC"}
 402         state.stateForScope(acc3)
 403         if !reflect.DeepEqual(state.AccountsStatus[acc3.ID].Account, acc3) {
 404                 t.Fatal("state for scope test err")
 405         }
 406 }