wallet/wallet_test.go

   1 package wallet
   2
   3 import (
   4         "encoding/json"
   5         "io/ioutil"
   6         "os"
   7         "reflect"
   8         "testing"
   9         "time"
  10
  11         "github.com/vapor/account"
  12         "github.com/vapor/asset"
  13         "github.com/vapor/blockchain/pseudohsm"
  14         "github.com/vapor/blockchain/signers"
  15         "github.com/vapor/blockchain/txbuilder"
  16         "github.com/vapor/config"
  17         "github.com/vapor/consensus"
  18         "github.com/vapor/crypto/ed25519/chainkd"
  19         "github.com/vapor/database"
  20         dbm "github.com/vapor/database/leveldb"
  21         "github.com/vapor/event"
  22         "github.com/vapor/protocol"
  23         "github.com/vapor/protocol/bc"
  24         "github.com/vapor/protocol/bc/types"
  25 )
  26
  27 func TestEncodeDecodeGlobalTxIndex(t *testing.T) {
  28         want := &struct {
  29                 BlockHash bc.Hash
  30                 Position  uint64
  31         }{
  32                 BlockHash: bc.NewHash([32]byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x20}),
  33                 Position:  1,
  34         }
  35
  36         globalTxIdx := calcGlobalTxIndex(&want.BlockHash, want.Position)
  37         blockHashGot, positionGot := parseGlobalTxIdx(globalTxIdx)
  38         if *blockHashGot != want.BlockHash {
  39                 t.Errorf("blockHash mismatch. Get: %v. Expect: %v", *blockHashGot, want.BlockHash)
  40         }
  41
  42         if positionGot != want.Position {
  43                 t.Errorf("position mismatch. Get: %v. Expect: %v", positionGot, want.Position)
  44         }
  45 }
  46
  47 func TestWalletVersion(t *testing.T) {
  48         // prepare wallet
  49         dirPath, err := ioutil.TempDir(".", "")
  50         if err != nil {
  51                 t.Fatal(err)
  52         }
  53         defer os.RemoveAll(dirPath)
  54
  55         testDB := dbm.NewDB("testdb", "leveldb", "temp")
  56         defer os.RemoveAll("temp")
  57
  58         dispatcher := event.NewDispatcher()
  59         w := mockWallet(testDB, nil, nil, nil, dispatcher, false)
  60
  61         // legacy status test case
  62         type legacyStatusInfo struct {
  63                 WorkHeight uint64
  64                 WorkHash   bc.Hash
  65                 BestHeight uint64
  66                 BestHash   bc.Hash
  67         }
  68         rawWallet, err := json.Marshal(legacyStatusInfo{})
  69         if err != nil {
  70                 t.Fatal("Marshal legacyStatusInfo")
  71         }
  72
  73         w.DB.Set(walletKey, rawWallet)
  74         rawWallet = w.DB.Get(walletKey)
  75         if rawWallet == nil {
  76                 t.Fatal("fail to load wallet StatusInfo")
  77         }
  78
  79         if err := json.Unmarshal(rawWallet, &w.status); err != nil {
  80                 t.Fatal(err)
  81         }
  82
  83         if err := w.checkWalletInfo(); err != errWalletVersionMismatch {
  84                 t.Fatal("fail to detect legacy wallet version")
  85         }
  86
  87         // lower wallet version test case
  88         lowerVersion := StatusInfo{Version: currentVersion - 1}
  89         rawWallet, err = json.Marshal(lowerVersion)
  90         if err != nil {
  91                 t.Fatal("save wallet info")
  92         }
  93
  94         w.DB.Set(walletKey, rawWallet)
  95         rawWallet = w.DB.Get(walletKey)
  96         if rawWallet == nil {
  97                 t.Fatal("fail to load wallet StatusInfo")
  98         }
  99
 100         if err := json.Unmarshal(rawWallet, &w.status); err != nil {
 101                 t.Fatal(err)
 102         }
 103
 104         if err := w.checkWalletInfo(); err != errWalletVersionMismatch {
 105                 t.Fatal("fail to detect expired wallet version")
 106         }
 107 }
 108
 109 func TestWalletUpdate(t *testing.T) {
 110         dirPath, err := ioutil.TempDir(".", "")
 111         if err != nil {
 112                 t.Fatal(err)
 113         }
 114         defer os.RemoveAll(dirPath)
 115
 116         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 117         defer os.RemoveAll("temp")
 118
 119         store := database.NewStore(testDB)
 120         dispatcher := event.NewDispatcher()
 121         txPool := protocol.NewTxPool(store, dispatcher)
 122
 123         chain, err := protocol.NewChain(store, txPool)
 124         if err != nil {
 125                 t.Fatal(err)
 126         }
 127
 128         accountManager := account.NewManager(testDB, chain)
 129         hsm, err := pseudohsm.New(dirPath)
 130         if err != nil {
 131                 t.Fatal(err)
 132         }
 133
 134         xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
 135         if err != nil {
 136                 t.Fatal(err)
 137         }
 138
 139         testAccount, err := accountManager.Create([]chainkd.XPub{xpub1.XPub}, 1, "testAccount", signers.BIP0044)
 140         if err != nil {
 141                 t.Fatal(err)
 142         }
 143
 144         controlProg, err := accountManager.CreateAddress(testAccount.ID, false)
 145         if err != nil {
 146                 t.Fatal(err)
 147         }
 148
 149         controlProg.KeyIndex = 1
 150
 151         reg := asset.NewRegistry(testDB, chain)
 152         asset := bc.AssetID{V0: 5}
 153
 154         utxos := []*account.UTXO{}
 155         btmUtxo := mockUTXO(controlProg, consensus.BTMAssetID)
 156         utxos = append(utxos, btmUtxo)
 157         OtherUtxo := mockUTXO(controlProg, &asset)
 158         utxos = append(utxos, OtherUtxo)
 159
 160         _, txData, err := mockTxData(utxos, testAccount)
 161         if err != nil {
 162                 t.Fatal(err)
 163         }
 164
 165         tx := types.NewTx(*txData)
 166         block := mockSingleBlock(tx)
 167         txStatus := bc.NewTransactionStatus()
 168         txStatus.SetStatus(0, false)
 169         txStatus.SetStatus(1, false)
 170         store.SaveBlock(block, txStatus)
 171
 172         w := mockWallet(testDB, accountManager, reg, chain, dispatcher, true)
 173         err = w.AttachBlock(block)
 174         if err != nil {
 175                 t.Fatal(err)
 176         }
 177
 178         if _, err := w.GetTransactionByTxID(tx.ID.String()); err != nil {
 179                 t.Fatal(err)
 180         }
 181
 182         wants, err := w.GetTransactions("")
 183         if len(wants) != 1 {
 184                 t.Fatal(err)
 185         }
 186
 187         if wants[0].ID != tx.ID {
 188                 t.Fatal("account txID mismatch")
 189         }
 190
 191         for position, tx := range block.Transactions {
 192                 get := w.DB.Get(calcGlobalTxIndexKey(tx.ID.String()))
 193                 bh := block.BlockHeader.Hash()
 194                 expect := calcGlobalTxIndex(&bh, uint64(position))
 195                 if !reflect.DeepEqual(get, expect) {
 196                         t.Fatalf("position#%d: compare retrieved globalTxIdx err", position)
 197                 }
 198         }
 199 }
 200
 201 func TestRescanWallet(t *testing.T) {
 202         // prepare wallet & db
 203         dirPath, err := ioutil.TempDir(".", "")
 204         if err != nil {
 205                 t.Fatal(err)
 206         }
 207         defer os.RemoveAll(dirPath)
 208
 209         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 210         defer os.RemoveAll("temp")
 211
 212         store := database.NewStore(testDB)
 213         dispatcher := event.NewDispatcher()
 214         txPool := protocol.NewTxPool(store, dispatcher)
 215         chain, err := protocol.NewChain(store, txPool)
 216         if err != nil {
 217                 t.Fatal(err)
 218         }
 219
 220         statusInfo := StatusInfo{
 221                 Version:  currentVersion,
 222                 WorkHash: bc.Hash{V0: 0xff},
 223         }
 224         rawWallet, err := json.Marshal(statusInfo)
 225         if err != nil {
 226                 t.Fatal("save wallet info")
 227         }
 228
 229         w := mockWallet(testDB, nil, nil, chain, dispatcher, false)
 230         w.DB.Set(walletKey, rawWallet)
 231         rawWallet = w.DB.Get(walletKey)
 232         if rawWallet == nil {
 233                 t.Fatal("fail to load wallet StatusInfo")
 234         }
 235
 236         if err := json.Unmarshal(rawWallet, &w.status); err != nil {
 237                 t.Fatal(err)
 238         }
 239
 240         // rescan wallet
 241         if err := w.loadWalletInfo(); err != nil {
 242                 t.Fatal(err)
 243         }
 244
 245         block := config.GenesisBlock()
 246         if w.status.WorkHash != block.Hash() {
 247                 t.Fatal("reattach from genesis block")
 248         }
 249 }
 250
 251 func TestMemPoolTxQueryLoop(t *testing.T) {
 252         dirPath, err := ioutil.TempDir(".", "")
 253         if err != nil {
 254                 t.Fatal(err)
 255         }
 256         defer os.RemoveAll(dirPath)
 257
 258         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 259
 260         store := database.NewStore(testDB)
 261         dispatcher := event.NewDispatcher()
 262         txPool := protocol.NewTxPool(store, dispatcher)
 263
 264         chain, err := protocol.NewChain(store, txPool)
 265         if err != nil {
 266                 t.Fatal(err)
 267         }
 268
 269         accountManager := account.NewManager(testDB, chain)
 270         hsm, err := pseudohsm.New(dirPath)
 271         if err != nil {
 272                 t.Fatal(err)
 273         }
 274
 275         xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
 276         if err != nil {
 277                 t.Fatal(err)
 278         }
 279
 280         testAccount, err := accountManager.Create([]chainkd.XPub{xpub1.XPub}, 1, "testAccount", signers.BIP0044)
 281         if err != nil {
 282                 t.Fatal(err)
 283         }
 284
 285         controlProg, err := accountManager.CreateAddress(testAccount.ID, false)
 286         if err != nil {
 287                 t.Fatal(err)
 288         }
 289
 290         controlProg.KeyIndex = 1
 291
 292         reg := asset.NewRegistry(testDB, chain)
 293         asset := bc.AssetID{V0: 5}
 294
 295         utxos := []*account.UTXO{}
 296         btmUtxo := mockUTXO(controlProg, consensus.BTMAssetID)
 297         utxos = append(utxos, btmUtxo)
 298         OtherUtxo := mockUTXO(controlProg, &asset)
 299         utxos = append(utxos, OtherUtxo)
 300
 301         _, txData, err := mockTxData(utxos, testAccount)
 302         if err != nil {
 303                 t.Fatal(err)
 304         }
 305
 306         tx := types.NewTx(*txData)
 307         //block := mockSingleBlock(tx)
 308         txStatus := bc.NewTransactionStatus()
 309         txStatus.SetStatus(0, false)
 310         w, err := NewWallet(testDB, accountManager, reg, hsm, chain, dispatcher, false)
 311         go w.memPoolTxQueryLoop()
 312         w.eventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: protocol.MsgNewTx}})
 313         time.Sleep(time.Millisecond * 10)
 314         if _, err = w.GetUnconfirmedTxByTxID(tx.ID.String()); err != nil {
 315                 t.Fatal("disaptch new tx msg error:", err)
 316         }
 317         w.eventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: protocol.MsgRemoveTx}})
 318         time.Sleep(time.Millisecond * 10)
 319         txs, err := w.GetUnconfirmedTxs(testAccount.ID)
 320         if err != nil {
 321                 t.Fatal("get unconfirmed tx error:", err)
 322         }
 323
 324         if len(txs) != 0 {
 325                 t.Fatal("disaptch remove tx msg error")
 326         }
 327
 328         w.eventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: 2}})
 329 }
 330
 331 func mockUTXO(controlProg *account.CtrlProgram, assetID *bc.AssetID) *account.UTXO {
 332         utxo := &account.UTXO{}
 333         utxo.OutputID = bc.Hash{V0: 1}
 334         utxo.SourceID = bc.Hash{V0: 2}
 335         utxo.AssetID = *assetID
 336         utxo.Amount = 1000000000
 337         utxo.SourcePos = 0
 338         utxo.ControlProgram = controlProg.ControlProgram
 339         utxo.AccountID = controlProg.AccountID
 340         utxo.Address = controlProg.Address
 341         utxo.ControlProgramIndex = controlProg.KeyIndex
 342         return utxo
 343 }
 344
 345 func mockTxData(utxos []*account.UTXO, testAccount *account.Account) (*txbuilder.Template, *types.TxData, error) {
 346         tplBuilder := txbuilder.NewBuilder(time.Now())
 347
 348         for _, utxo := range utxos {
 349                 txInput, sigInst, err := account.UtxoToInputs(testAccount.Signer, utxo)
 350                 if err != nil {
 351                         return nil, nil, err
 352                 }
 353                 tplBuilder.AddInput(txInput, sigInst)
 354
 355                 out := &types.TxOutput{}
 356                 if utxo.AssetID == *consensus.BTMAssetID {
 357                         out = types.NewIntraChainOutput(utxo.AssetID, 100, utxo.ControlProgram)
 358                 } else {
 359                         out = types.NewIntraChainOutput(utxo.AssetID, utxo.Amount, utxo.ControlProgram)
 360                 }
 361                 tplBuilder.AddOutput(out)
 362         }
 363
 364         return tplBuilder.Build()
 365 }
 366
 367 func mockWallet(walletDB dbm.DB, account *account.Manager, asset *asset.Registry, chain *protocol.Chain, dispatcher *event.Dispatcher, txIndexFlag bool) *Wallet {
 368         wallet := &Wallet{
 369                 DB:              walletDB,
 370                 AccountMgr:      account,
 371                 AssetReg:        asset,
 372                 chain:           chain,
 373                 RecoveryMgr:     newRecoveryManager(walletDB, account),
 374                 eventDispatcher: dispatcher,
 375                 TxIndexFlag:     txIndexFlag,
 376         }
 377         wallet.txMsgSub, _ = wallet.eventDispatcher.Subscribe(protocol.TxMsgEvent{})
 378         return wallet
 379 }
 380
 381 func mockSingleBlock(tx *types.Tx) *types.Block {
 382         return &types.Block{
 383                 BlockHeader: types.BlockHeader{
 384                         Version: 1,
 385                         Height:  1,
 386                 },
 387                 Transactions: []*types.Tx{config.GenesisTx(), tx},
 388         }
 389 }