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