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