wallet/wallet_test.go

   1 package wallet
   2
   3 import (
   4         "encoding/binary"
   5         "encoding/hex"
   6         "encoding/json"
   7         "fmt"
   8         "io/ioutil"
   9         "os"
  10         "sort"
  11         "strings"
  12         "testing"
  13         "time"
  14
  15         "github.com/vapor/account"
  16         acc "github.com/vapor/account"
  17         "github.com/vapor/asset"
  18         "github.com/vapor/blockchain/query"
  19         "github.com/vapor/blockchain/signers"
  20         "github.com/vapor/blockchain/txbuilder"
  21         "github.com/vapor/common"
  22         "github.com/vapor/config"
  23         "github.com/vapor/consensus"
  24         "github.com/vapor/crypto/ed25519/chainkd"
  25         "github.com/vapor/crypto/sha3pool"
  26         dbm "github.com/vapor/database/leveldb"
  27         "github.com/vapor/errors"
  28         "github.com/vapor/event"
  29         "github.com/vapor/protocol"
  30         "github.com/vapor/protocol/bc"
  31         "github.com/vapor/protocol/bc/types"
  32 )
  33
  34 func TestEncodeDecodeGlobalTxIndex(t *testing.T) {
  35         want := &struct {
  36                 BlockHash bc.Hash
  37                 Position  uint64
  38         }{
  39                 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}),
  40                 Position:  1,
  41         }
  42
  43         globalTxIdx := CalcGlobalTxIndex(&want.BlockHash, want.Position)
  44         blockHashGot, positionGot := parseGlobalTxIdx(globalTxIdx)
  45         if *blockHashGot != want.BlockHash {
  46                 t.Errorf("blockHash mismatch. Get: %v. Expect: %v", *blockHashGot, want.BlockHash)
  47         }
  48
  49         if positionGot != want.Position {
  50                 t.Errorf("position mismatch. Get: %v. Expect: %v", positionGot, want.Position)
  51         }
  52 }
  53
  54 func TestWalletVersion(t *testing.T) {
  55         // prepare wallet
  56         dirPath, err := ioutil.TempDir(".", "")
  57         if err != nil {
  58                 t.Fatal(err)
  59         }
  60         defer os.RemoveAll(dirPath)
  61
  62         testDB := dbm.NewDB("testdb", "leveldb", "temp")
  63         walletStore := NewMockWalletStore(testDB)
  64         defer func() {
  65                 testDB.Close()
  66                 os.RemoveAll("temp")
  67         }()
  68
  69         dispatcher := event.NewDispatcher()
  70         w := mockWallet(walletStore, nil, nil, nil, dispatcher, false)
  71
  72         walletStatus := new(StatusInfo)
  73         if err := w.Store.SetWalletInfo(walletStatus); err != nil {
  74                 t.Fatal(err)
  75         }
  76
  77         status, err := w.Store.GetWalletInfo()
  78         if err != nil {
  79                 t.Fatal(err)
  80         }
  81         w.Status = *status
  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         if err := w.Store.SetWalletInfo(&lowerVersion); err != nil {
  90                 t.Fatal(err)
  91         }
  92
  93         status, err = w.Store.GetWalletInfo()
  94         if err != nil {
  95                 t.Fatal(err)
  96         }
  97         w.Status = *status
  98
  99         if err := w.checkWalletInfo(); err != errWalletVersionMismatch {
 100                 t.Fatal("fail to detect expired wallet version")
 101         }
 102 }
 103
 104 func mockUTXO(controlProg *account.CtrlProgram, assetID *bc.AssetID) *account.UTXO {
 105         utxo := &account.UTXO{}
 106         utxo.OutputID = bc.Hash{V0: 1}
 107         utxo.SourceID = bc.Hash{V0: 2}
 108         utxo.AssetID = *assetID
 109         utxo.Amount = 1000000000
 110         utxo.SourcePos = 0
 111         utxo.ControlProgram = controlProg.ControlProgram
 112         utxo.AccountID = controlProg.AccountID
 113         utxo.Address = controlProg.Address
 114         utxo.ControlProgramIndex = controlProg.KeyIndex
 115         return utxo
 116 }
 117
 118 func mockTxData(utxos []*account.UTXO, testAccount *account.Account) (*txbuilder.Template, *types.TxData, error) {
 119         tplBuilder := txbuilder.NewBuilder(time.Now())
 120
 121         for _, utxo := range utxos {
 122                 txInput, sigInst, err := account.UtxoToInputs(testAccount.Signer, utxo)
 123                 if err != nil {
 124                         return nil, nil, err
 125                 }
 126                 tplBuilder.AddInput(txInput, sigInst)
 127
 128                 out := &types.TxOutput{}
 129                 if utxo.AssetID == *consensus.BTMAssetID {
 130                         out = types.NewIntraChainOutput(utxo.AssetID, 100, utxo.ControlProgram)
 131                 } else {
 132                         out = types.NewIntraChainOutput(utxo.AssetID, utxo.Amount, utxo.ControlProgram)
 133                 }
 134                 tplBuilder.AddOutput(out)
 135         }
 136
 137         return tplBuilder.Build()
 138 }
 139
 140 func mockWallet(store WalletStore, account *account.Manager, asset *asset.Registry, chain *protocol.Chain, dispatcher *event.Dispatcher, txIndexFlag bool) *Wallet {
 141         wallet := &Wallet{
 142                 Store:           store,
 143                 AccountMgr:      account,
 144                 AssetReg:        asset,
 145                 Chain:           chain,
 146                 RecoveryMgr:     NewRecoveryManager(store, account),
 147                 EventDispatcher: dispatcher,
 148                 TxIndexFlag:     txIndexFlag,
 149         }
 150         wallet.TxMsgSub, _ = wallet.EventDispatcher.Subscribe(protocol.TxMsgEvent{})
 151         return wallet
 152 }
 153
 154 func mockSingleBlock(tx *types.Tx) *types.Block {
 155         return &types.Block{
 156                 BlockHeader: types.BlockHeader{
 157                         Version: 1,
 158                         Height:  1,
 159                 },
 160                 Transactions: []*types.Tx{config.GenesisTx(), tx},
 161         }
 162 }
 163
 164 // errors
 165 var (
 166         errAccntTxIDNotFound = errors.New("account TXID not found")
 167         errGetAsset          = errors.New("Failed to find asset definition")
 168 )
 169
 170 const (
 171         utxoPrefix byte = iota //UTXOPrefix is StandardUTXOKey prefix
 172         contractPrefix
 173         contractIndexPrefix
 174         accountPrefix // AccountPrefix is account ID prefix
 175         accountIndexPrefix
 176 )
 177
 178 // leveldb key prefix
 179 var (
 180         colon               byte = 0x3a
 181         accountStore             = []byte("AS:")
 182         UTXOPrefix               = append(accountStore, utxoPrefix, colon)
 183         ContractPrefix           = append(accountStore, contractPrefix, colon)
 184         ContractIndexPrefix      = append(accountStore, contractIndexPrefix, colon)
 185         AccountPrefix            = append(accountStore, accountPrefix, colon) // AccountPrefix is account ID prefix
 186         AccountIndexPrefix       = append(accountStore, accountIndexPrefix, colon)
 187 )
 188
 189 const (
 190         sutxoPrefix byte = iota //SUTXOPrefix is ContractUTXOKey prefix
 191         accountAliasPrefix
 192         txPrefix            //TxPrefix is wallet database transactions prefix
 193         txIndexPrefix       //TxIndexPrefix is wallet database tx index prefix
 194         unconfirmedTxPrefix //UnconfirmedTxPrefix is txpool unconfirmed transactions prefix
 195         globalTxIndexPrefix //GlobalTxIndexPrefix is wallet database global tx index prefix
 196         walletKey
 197         miningAddressKey
 198         coinbaseAbKey
 199         recoveryKey //recoveryKey key for db store recovery info.
 200 )
 201
 202 var (
 203         walletStore         = []byte("WS:")
 204         SUTXOPrefix         = append(walletStore, sutxoPrefix, colon)
 205         AccountAliasPrefix  = append(walletStore, accountAliasPrefix, colon)
 206         TxPrefix            = append(walletStore, txPrefix, colon)            //TxPrefix is wallet database transactions prefix
 207         TxIndexPrefix       = append(walletStore, txIndexPrefix, colon)       //TxIndexPrefix is wallet database tx index prefix
 208         UnconfirmedTxPrefix = append(walletStore, unconfirmedTxPrefix, colon) //UnconfirmedTxPrefix is txpool unconfirmed transactions prefix
 209         GlobalTxIndexPrefix = append(walletStore, globalTxIndexPrefix, colon) //GlobalTxIndexPrefix is wallet database global tx index prefix
 210         WalletKey           = append(walletStore, walletKey)
 211         MiningAddressKey    = append(walletStore, miningAddressKey)
 212         CoinbaseAbKey       = append(walletStore, coinbaseAbKey)
 213         RecoveryKey         = append(walletStore, recoveryKey)
 214 )
 215
 216 func accountIndexKey(xpubs []chainkd.XPub) []byte {
 217         var hash [32]byte
 218         var xPubs []byte
 219         cpy := append([]chainkd.XPub{}, xpubs[:]...)
 220         sort.Sort(signers.SortKeys(cpy))
 221         for _, xpub := range cpy {
 222                 xPubs = append(xPubs, xpub[:]...)
 223         }
 224         sha3pool.Sum256(hash[:], xPubs)
 225         return append([]byte(AccountIndexPrefix), hash[:]...)
 226 }
 227
 228 func Bip44ContractIndexKey(accountID string, change bool) []byte {
 229         key := append([]byte(ContractIndexPrefix), accountID...)
 230         if change {
 231                 return append(key, []byte{1}...)
 232         }
 233         return append(key, []byte{0}...)
 234 }
 235
 236 // ContractKey account control promgram store prefix
 237 func ContractKey(hash bc.Hash) []byte {
 238         return append([]byte(ContractPrefix), hash.Bytes()...)
 239 }
 240
 241 // AccountIDKey account id store prefix
 242 func AccountIDKey(accountID string) []byte {
 243         return append([]byte(AccountPrefix), []byte(accountID)...)
 244 }
 245
 246 // StandardUTXOKey makes an account unspent outputs key to store
 247 func StandardUTXOKey(id bc.Hash) []byte {
 248         return append(UTXOPrefix, id.Bytes()...)
 249 }
 250
 251 // ContractUTXOKey makes a smart contract unspent outputs key to store
 252 func ContractUTXOKey(id bc.Hash) []byte {
 253         return append(SUTXOPrefix, id.Bytes()...)
 254 }
 255
 256 func calcDeleteKey(blockHeight uint64) []byte {
 257         return []byte(fmt.Sprintf("%s%016x", TxPrefix, blockHeight))
 258 }
 259
 260 func calcTxIndexKey(txID string) []byte {
 261         return append(TxIndexPrefix, []byte(txID)...)
 262 }
 263
 264 func calcAnnotatedKey(formatKey string) []byte {
 265         return append(TxPrefix, []byte(formatKey)...)
 266 }
 267
 268 func calcUnconfirmedTxKey(formatKey string) []byte {
 269         return append(UnconfirmedTxPrefix, []byte(formatKey)...)
 270 }
 271
 272 func CalcGlobalTxIndexKey(txID string) []byte {
 273         return append(GlobalTxIndexPrefix, []byte(txID)...)
 274 }
 275
 276 func CalcGlobalTxIndex(blockHash *bc.Hash, position uint64) []byte {
 277         txIdx := make([]byte, 40)
 278         copy(txIdx[:32], blockHash.Bytes())
 279         binary.BigEndian.PutUint64(txIdx[32:], position)
 280         return txIdx
 281 }
 282
 283 func formatKey(blockHeight uint64, position uint32) string {
 284         return fmt.Sprintf("%016x%08x", blockHeight, position)
 285 }
 286
 287 func contractIndexKey(accountID string) []byte {
 288         return append([]byte(ContractIndexPrefix), []byte(accountID)...)
 289 }
 290
 291 func accountAliasKey(name string) []byte {
 292         return append([]byte(AccountAliasPrefix), []byte(name)...)
 293 }
 294
 295 // MockWalletStore store wallet using leveldb
 296 type MockWalletStore struct {
 297         db    dbm.DB
 298         batch dbm.Batch
 299 }
 300
 301 // NewMockWalletStore create new MockWalletStore struct
 302 func NewMockWalletStore(db dbm.DB) *MockWalletStore {
 303         return &MockWalletStore{
 304                 db:    db,
 305                 batch: nil,
 306         }
 307 }
 308
 309 // InitBatch initial new wallet store
 310 func (store *MockWalletStore) InitBatch() WalletStore {
 311         newStore := NewMockWalletStore(store.db)
 312         newStore.batch = newStore.db.NewBatch()
 313         return newStore
 314 }
 315
 316 // CommitBatch commit batch
 317 func (store *MockWalletStore) CommitBatch() error {
 318         if store.batch == nil {
 319                 return errors.New("MockWalletStore commit fail, store batch is nil.")
 320         }
 321         store.batch.Write()
 322         store.batch = nil
 323         return nil
 324 }
 325
 326 // DeleteContractUTXO delete contract utxo by outputID
 327 func (store *MockWalletStore) DeleteContractUTXO(outputID bc.Hash) {
 328         if store.batch == nil {
 329                 store.db.Delete(ContractUTXOKey(outputID))
 330         } else {
 331                 store.batch.Delete(ContractUTXOKey(outputID))
 332         }
 333 }
 334
 335 // DeleteRecoveryStatus delete recovery status
 336 func (store *MockWalletStore) DeleteRecoveryStatus() {
 337         if store.batch == nil {
 338                 store.db.Delete(RecoveryKey)
 339         } else {
 340                 store.batch.Delete(RecoveryKey)
 341         }
 342 }
 343
 344 // DeleteTransactions delete transactions when orphan block rollback
 345 func (store *MockWalletStore) DeleteTransactions(height uint64) {
 346         batch := store.db.NewBatch()
 347         if store.batch != nil {
 348                 batch = store.batch
 349         }
 350         txIter := store.db.IteratorPrefix(calcDeleteKey(height))
 351         defer txIter.Release()
 352
 353         tmpTx := query.AnnotatedTx{}
 354         for txIter.Next() {
 355                 if err := json.Unmarshal(txIter.Value(), &tmpTx); err == nil {
 356                         batch.Delete(calcTxIndexKey(tmpTx.ID.String()))
 357                 }
 358                 batch.Delete(txIter.Key())
 359         }
 360         if store.batch == nil {
 361                 batch.Write()
 362         }
 363 }
 364
 365 // DeleteUnconfirmedTransaction delete unconfirmed tx by txID
 366 func (store *MockWalletStore) DeleteUnconfirmedTransaction(txID string) {
 367         if store.batch == nil {
 368                 store.db.Delete(calcUnconfirmedTxKey(txID))
 369         } else {
 370                 store.batch.Delete(calcUnconfirmedTxKey(txID))
 371         }
 372 }
 373
 374 // DeleteWalletTransactions delete all txs in wallet
 375 func (store *MockWalletStore) DeleteWalletTransactions() {
 376         batch := store.db.NewBatch()
 377         if store.batch != nil {
 378                 batch = store.batch
 379         }
 380         txIter := store.db.IteratorPrefix([]byte(TxPrefix))
 381         defer txIter.Release()
 382
 383         for txIter.Next() {
 384                 batch.Delete(txIter.Key())
 385         }
 386
 387         txIndexIter := store.db.IteratorPrefix([]byte(TxIndexPrefix))
 388         defer txIndexIter.Release()
 389
 390         for txIndexIter.Next() {
 391                 batch.Delete(txIndexIter.Key())
 392         }
 393         if store.batch == nil {
 394                 batch.Write()
 395         }
 396 }
 397
 398 // DeleteWalletUTXOs delete all txs in wallet
 399 func (store *MockWalletStore) DeleteWalletUTXOs() {
 400         batch := store.db.NewBatch()
 401         if store.batch != nil {
 402                 batch = store.batch
 403         }
 404         ruIter := store.db.IteratorPrefix([]byte(UTXOPrefix))
 405         defer ruIter.Release()
 406         for ruIter.Next() {
 407                 batch.Delete(ruIter.Key())
 408         }
 409
 410         suIter := store.db.IteratorPrefix([]byte(SUTXOPrefix))
 411         defer suIter.Release()
 412         for suIter.Next() {
 413                 batch.Delete(suIter.Key())
 414         }
 415         if store.batch == nil {
 416                 batch.Write()
 417         }
 418 }
 419
 420 // GetAsset get asset by assetID
 421 func (store *MockWalletStore) GetAsset(assetID *bc.AssetID) (*asset.Asset, error) {
 422         definitionByte := store.db.Get(asset.ExtAssetKey(assetID))
 423         if definitionByte == nil {
 424                 return nil, errGetAsset
 425         }
 426         definitionMap := make(map[string]interface{})
 427         if err := json.Unmarshal(definitionByte, &definitionMap); err != nil {
 428                 return nil, err
 429         }
 430         alias := assetID.String()
 431         externalAsset := &asset.Asset{
 432                 AssetID:           *assetID,
 433                 Alias:             &alias,
 434                 DefinitionMap:     definitionMap,
 435                 RawDefinitionByte: definitionByte,
 436         }
 437         return externalAsset, nil
 438 }
 439
 440 // GetGlobalTransactionIndex get global tx by txID
 441 func (store *MockWalletStore) GetGlobalTransactionIndex(txID string) []byte {
 442         return store.db.Get(CalcGlobalTxIndexKey(txID))
 443 }
 444
 445 // GetStandardUTXO get standard utxo by id
 446 func (store *MockWalletStore) GetStandardUTXO(outid bc.Hash) (*acc.UTXO, error) {
 447         rawUTXO := store.db.Get(StandardUTXOKey(outid))
 448         if rawUTXO == nil {
 449                 return nil, fmt.Errorf("failed get standard UTXO, outputID: %s ", outid.String())
 450         }
 451         UTXO := new(acc.UTXO)
 452         if err := json.Unmarshal(rawUTXO, UTXO); err != nil {
 453                 return nil, err
 454         }
 455         return UTXO, nil
 456 }
 457
 458 // GetTransaction get tx by txid
 459 func (store *MockWalletStore) GetTransaction(txID string) (*query.AnnotatedTx, error) {
 460         formatKey := store.db.Get(calcTxIndexKey(txID))
 461         if formatKey == nil {
 462                 return nil, errAccntTxIDNotFound
 463         }
 464         rawTx := store.db.Get(calcAnnotatedKey(string(formatKey)))
 465         tx := new(query.AnnotatedTx)
 466         if err := json.Unmarshal(rawTx, tx); err != nil {
 467                 return nil, err
 468         }
 469         return tx, nil
 470 }
 471
 472 // GetUnconfirmedTransaction get unconfirmed tx by txID
 473 func (store *MockWalletStore) GetUnconfirmedTransaction(txID string) (*query.AnnotatedTx, error) {
 474         rawUnconfirmedTx := store.db.Get(calcUnconfirmedTxKey(txID))
 475         if rawUnconfirmedTx == nil {
 476                 return nil, fmt.Errorf("failed get unconfirmed tx, txID: %s ", txID)
 477         }
 478         tx := new(query.AnnotatedTx)
 479         if err := json.Unmarshal(rawUnconfirmedTx, tx); err != nil {
 480                 return nil, err
 481         }
 482         return tx, nil
 483 }
 484
 485 // GetRecoveryStatus delete recovery status
 486 func (store *MockWalletStore) GetRecoveryStatus() (*RecoveryState, error) {
 487         rawStatus := store.db.Get(RecoveryKey)
 488         if rawStatus == nil {
 489                 return nil, ErrGetRecoveryStatus
 490         }
 491         state := new(RecoveryState)
 492         if err := json.Unmarshal(rawStatus, state); err != nil {
 493                 return nil, err
 494         }
 495         return state, nil
 496 }
 497
 498 // GetWalletInfo get wallet information
 499 func (store *MockWalletStore) GetWalletInfo() (*StatusInfo, error) {
 500         rawStatus := store.db.Get([]byte(WalletKey))
 501         if rawStatus == nil {
 502                 return nil, fmt.Errorf("failed get wallet info")
 503         }
 504         status := new(StatusInfo)
 505         if err := json.Unmarshal(rawStatus, status); err != nil {
 506                 return nil, err
 507         }
 508         return status, nil
 509 }
 510
 511 // ListAccountUTXOs get all account unspent outputs
 512 func (store *MockWalletStore) ListAccountUTXOs(id string, isSmartContract bool) ([]*acc.UTXO, error) {
 513         prefix := UTXOPrefix
 514         if isSmartContract {
 515                 prefix = SUTXOPrefix
 516         }
 517
 518         idBytes, err := hex.DecodeString(id)
 519         if err != nil {
 520                 return nil, err
 521         }
 522
 523         accountUtxoIter := store.db.IteratorPrefix(append(prefix, idBytes...))
 524         defer accountUtxoIter.Release()
 525
 526         confirmedUTXOs := []*acc.UTXO{}
 527         for accountUtxoIter.Next() {
 528                 utxo := new(acc.UTXO)
 529                 if err := json.Unmarshal(accountUtxoIter.Value(), utxo); err != nil {
 530                         return nil, err
 531                 }
 532
 533                 confirmedUTXOs = append(confirmedUTXOs, utxo)
 534         }
 535         return confirmedUTXOs, nil
 536 }
 537
 538 func (store *MockWalletStore) ListTransactions(accountID string, StartTxID string, count uint, unconfirmed bool) ([]*query.AnnotatedTx, error) {
 539         annotatedTxs := []*query.AnnotatedTx{}
 540         var startKey []byte
 541         preFix := TxPrefix
 542
 543         if StartTxID != "" {
 544                 if unconfirmed {
 545                         startKey = calcUnconfirmedTxKey(StartTxID)
 546                 } else {
 547                         formatKey := store.db.Get(calcTxIndexKey(StartTxID))
 548                         if formatKey == nil {
 549                                 return nil, errAccntTxIDNotFound
 550                         }
 551                         startKey = calcAnnotatedKey(string(formatKey))
 552                 }
 553         }
 554
 555         if unconfirmed {
 556                 preFix = UnconfirmedTxPrefix
 557         }
 558
 559         itr := store.db.IteratorPrefixWithStart([]byte(preFix), startKey, true)
 560         defer itr.Release()
 561
 562         for txNum := count; itr.Next() && txNum > 0; txNum-- {
 563                 annotatedTx := new(query.AnnotatedTx)
 564                 if err := json.Unmarshal(itr.Value(), &annotatedTx); err != nil {
 565                         return nil, err
 566                 }
 567                 annotatedTxs = append(annotatedTxs, annotatedTx)
 568         }
 569
 570         return annotatedTxs, nil
 571 }
 572
 573 // ListUnconfirmedTransactions get all unconfirmed txs
 574 func (store *MockWalletStore) ListUnconfirmedTransactions() ([]*query.AnnotatedTx, error) {
 575         annotatedTxs := []*query.AnnotatedTx{}
 576         txIter := store.db.IteratorPrefix([]byte(UnconfirmedTxPrefix))
 577         defer txIter.Release()
 578
 579         for txIter.Next() {
 580                 annotatedTx := &query.AnnotatedTx{}
 581                 if err := json.Unmarshal(txIter.Value(), &annotatedTx); err != nil {
 582                         return nil, err
 583                 }
 584                 annotatedTxs = append(annotatedTxs, annotatedTx)
 585         }
 586         return annotatedTxs, nil
 587 }
 588
 589 // SetAssetDefinition set assetID and definition
 590 func (store *MockWalletStore) SetAssetDefinition(assetID *bc.AssetID, definition []byte) {
 591         if store.batch == nil {
 592                 store.db.Set(asset.ExtAssetKey(assetID), definition)
 593         } else {
 594                 store.batch.Set(asset.ExtAssetKey(assetID), definition)
 595         }
 596 }
 597
 598 // SetContractUTXO set standard utxo
 599 func (store *MockWalletStore) SetContractUTXO(outputID bc.Hash, utxo *acc.UTXO) error {
 600         data, err := json.Marshal(utxo)
 601         if err != nil {
 602                 return err
 603         }
 604         if store.batch == nil {
 605                 store.db.Set(ContractUTXOKey(outputID), data)
 606         } else {
 607                 store.batch.Set(ContractUTXOKey(outputID), data)
 608         }
 609         return nil
 610 }
 611
 612 // SetGlobalTransactionIndex set global tx index by blockhash and position
 613 func (store *MockWalletStore) SetGlobalTransactionIndex(globalTxID string, blockHash *bc.Hash, position uint64) {
 614         if store.batch == nil {
 615                 store.db.Set(CalcGlobalTxIndexKey(globalTxID), CalcGlobalTxIndex(blockHash, position))
 616         } else {
 617                 store.batch.Set(CalcGlobalTxIndexKey(globalTxID), CalcGlobalTxIndex(blockHash, position))
 618         }
 619 }
 620
 621 // SetRecoveryStatus set recovery status
 622 func (store *MockWalletStore) SetRecoveryStatus(recoveryState *RecoveryState) error {
 623         rawStatus, err := json.Marshal(recoveryState)
 624         if err != nil {
 625                 return err
 626         }
 627         if store.batch == nil {
 628                 store.db.Set(RecoveryKey, rawStatus)
 629         } else {
 630                 store.batch.Set(RecoveryKey, rawStatus)
 631         }
 632         return nil
 633 }
 634
 635 // SetTransaction set raw transaction by block height and tx position
 636 func (store *MockWalletStore) SetTransaction(height uint64, tx *query.AnnotatedTx) error {
 637         batch := store.db.NewBatch()
 638         if store.batch != nil {
 639                 batch = store.batch
 640         }
 641
 642         rawTx, err := json.Marshal(tx)
 643         if err != nil {
 644                 return err
 645         }
 646         batch.Set(calcAnnotatedKey(formatKey(height, tx.Position)), rawTx)
 647         batch.Set(calcTxIndexKey(tx.ID.String()), []byte(formatKey(height, tx.Position)))
 648
 649         if store.batch == nil {
 650                 batch.Write()
 651         }
 652         return nil
 653 }
 654
 655 // SetUnconfirmedTransaction set unconfirmed tx by txID
 656 func (store *MockWalletStore) SetUnconfirmedTransaction(txID string, tx *query.AnnotatedTx) error {
 657         rawTx, err := json.Marshal(tx)
 658         if err != nil {
 659                 return err
 660         }
 661         if store.batch == nil {
 662                 store.db.Set(calcUnconfirmedTxKey(txID), rawTx)
 663         } else {
 664                 store.batch.Set(calcUnconfirmedTxKey(txID), rawTx)
 665         }
 666         return nil
 667 }
 668
 669 // SetWalletInfo get wallet information
 670 func (store *MockWalletStore) SetWalletInfo(status *StatusInfo) error {
 671         rawWallet, err := json.Marshal(status)
 672         if err != nil {
 673                 return err
 674         }
 675
 676         if store.batch == nil {
 677                 store.db.Set([]byte(WalletKey), rawWallet)
 678         } else {
 679                 store.batch.Set([]byte(WalletKey), rawWallet)
 680         }
 681         return nil
 682 }
 683
 684 type MockAccountStore struct {
 685         db    dbm.DB
 686         batch dbm.Batch
 687 }
 688
 689 // NewAccountStore create new MockAccountStore.
 690 func NewMockAccountStore(db dbm.DB) *MockAccountStore {
 691         return &MockAccountStore{
 692                 db:    db,
 693                 batch: nil,
 694         }
 695 }
 696
 697 // InitBatch initial new account store
 698 func (store *MockAccountStore) InitBatch() acc.AccountStore {
 699         newStore := NewMockAccountStore(store.db)
 700         newStore.batch = newStore.db.NewBatch()
 701         return newStore
 702 }
 703
 704 // CommitBatch commit batch
 705 func (store *MockAccountStore) CommitBatch() error {
 706         if store.batch == nil {
 707                 return errors.New("MockAccountStore commit fail, store batch is nil.")
 708         }
 709         store.batch.Write()
 710         store.batch = nil
 711         return nil
 712 }
 713
 714 // DeleteAccount set account account ID, account alias and raw account.
 715 func (store *MockAccountStore) DeleteAccount(account *acc.Account) error {
 716         batch := store.db.NewBatch()
 717         if store.batch != nil {
 718                 batch = store.batch
 719         }
 720
 721         // delete account utxos
 722         store.deleteAccountUTXOs(account.ID, batch)
 723
 724         // delete account control program
 725         if err := store.deleteAccountControlPrograms(account.ID, batch); err != nil {
 726                 return err
 727         }
 728
 729         // delete bip44 contract index
 730         batch.Delete(Bip44ContractIndexKey(account.ID, false))
 731         batch.Delete(Bip44ContractIndexKey(account.ID, true))
 732
 733         // delete contract index
 734         batch.Delete(contractIndexKey(account.ID))
 735
 736         // delete account id
 737         batch.Delete(AccountIDKey(account.ID))
 738         batch.Delete(accountAliasKey(account.Alias))
 739         if store.batch == nil {
 740                 batch.Write()
 741         }
 742         return nil
 743 }
 744
 745 // deleteAccountUTXOs delete account utxos by accountID
 746 func (store *MockAccountStore) deleteAccountUTXOs(accountID string, batch dbm.Batch) error {
 747         accountUtxoIter := store.db.IteratorPrefix([]byte(UTXOPrefix))
 748         defer accountUtxoIter.Release()
 749
 750         for accountUtxoIter.Next() {
 751                 accountUtxo := new(acc.UTXO)
 752                 if err := json.Unmarshal(accountUtxoIter.Value(), accountUtxo); err != nil {
 753                         return err
 754                 }
 755
 756                 if accountID == accountUtxo.AccountID {
 757                         batch.Delete(StandardUTXOKey(accountUtxo.OutputID))
 758                 }
 759         }
 760
 761         return nil
 762 }
 763
 764 // deleteAccountControlPrograms deletes account control program
 765 func (store *MockAccountStore) deleteAccountControlPrograms(accountID string, batch dbm.Batch) error {
 766         cps, err := store.ListControlPrograms()
 767         if err != nil {
 768                 return err
 769         }
 770
 771         var hash [32]byte
 772         for _, cp := range cps {
 773                 if cp.AccountID == accountID {
 774                         sha3pool.Sum256(hash[:], cp.ControlProgram)
 775                         batch.Delete(ContractKey(bc.NewHash(hash)))
 776                 }
 777         }
 778         return nil
 779 }
 780
 781 // DeleteStandardUTXO delete utxo by outpu id
 782 func (store *MockAccountStore) DeleteStandardUTXO(outputID bc.Hash) {
 783         if store.batch == nil {
 784                 store.db.Delete(StandardUTXOKey(outputID))
 785         } else {
 786                 store.batch.Delete(StandardUTXOKey(outputID))
 787         }
 788 }
 789
 790 // GetAccountByAlias get account by account alias
 791 func (store *MockAccountStore) GetAccountByAlias(accountAlias string) (*acc.Account, error) {
 792         accountID := store.db.Get(accountAliasKey(accountAlias))
 793         if accountID == nil {
 794                 return nil, acc.ErrFindAccount
 795         }
 796         return store.GetAccountByID(string(accountID))
 797 }
 798
 799 // GetAccountByID get account by accountID
 800 func (store *MockAccountStore) GetAccountByID(accountID string) (*acc.Account, error) {
 801         rawAccount := store.db.Get(AccountIDKey(accountID))
 802         if rawAccount == nil {
 803                 return nil, acc.ErrFindAccount
 804         }
 805         account := new(acc.Account)
 806         if err := json.Unmarshal(rawAccount, account); err != nil {
 807                 return nil, err
 808         }
 809         return account, nil
 810 }
 811
 812 // GetAccountIndex get account index by account xpubs
 813 func (store *MockAccountStore) GetAccountIndex(xpubs []chainkd.XPub) uint64 {
 814         currentIndex := uint64(0)
 815         if rawIndexBytes := store.db.Get(accountIndexKey(xpubs)); rawIndexBytes != nil {
 816                 currentIndex = common.BytesToUnit64(rawIndexBytes)
 817         }
 818         return currentIndex
 819 }
 820
 821 // GetBip44ContractIndex get bip44 contract index
 822 func (store *MockAccountStore) GetBip44ContractIndex(accountID string, change bool) uint64 {
 823         index := uint64(0)
 824         if rawIndexBytes := store.db.Get(Bip44ContractIndexKey(accountID, change)); rawIndexBytes != nil {
 825                 index = common.BytesToUnit64(rawIndexBytes)
 826         }
 827         return index
 828 }
 829
 830 // GetCoinbaseArbitrary get coinbase arbitrary
 831 func (store *MockAccountStore) GetCoinbaseArbitrary() []byte {
 832         return store.db.Get([]byte(CoinbaseAbKey))
 833 }
 834
 835 // GetContractIndex get contract index
 836 func (store *MockAccountStore) GetContractIndex(accountID string) uint64 {
 837         index := uint64(0)
 838         if rawIndexBytes := store.db.Get(contractIndexKey(accountID)); rawIndexBytes != nil {
 839                 index = common.BytesToUnit64(rawIndexBytes)
 840         }
 841         return index
 842 }
 843
 844 // GetControlProgram get control program
 845 func (store *MockAccountStore) GetControlProgram(hash bc.Hash) (*acc.CtrlProgram, error) {
 846         rawProgram := store.db.Get(ContractKey(hash))
 847         if rawProgram == nil {
 848                 return nil, acc.ErrFindCtrlProgram
 849         }
 850         cp := new(acc.CtrlProgram)
 851         if err := json.Unmarshal(rawProgram, cp); err != nil {
 852                 return nil, err
 853         }
 854         return cp, nil
 855 }
 856
 857 // GetMiningAddress get mining address
 858 func (store *MockAccountStore) GetMiningAddress() (*acc.CtrlProgram, error) {
 859         rawCP := store.db.Get([]byte(MiningAddressKey))
 860         if rawCP == nil {
 861                 return nil, acc.ErrFindMiningAddress
 862         }
 863         cp := new(acc.CtrlProgram)
 864         if err := json.Unmarshal(rawCP, cp); err != nil {
 865                 return nil, err
 866         }
 867         return cp, nil
 868 }
 869
 870 // GetUTXO get standard utxo by id
 871 func (store *MockAccountStore) GetUTXO(outid bc.Hash) (*acc.UTXO, error) {
 872         u := new(acc.UTXO)
 873         if data := store.db.Get(StandardUTXOKey(outid)); data != nil {
 874                 return u, json.Unmarshal(data, u)
 875         }
 876         if data := store.db.Get(ContractUTXOKey(outid)); data != nil {
 877                 return u, json.Unmarshal(data, u)
 878         }
 879         return nil, acc.ErrMatchUTXO
 880 }
 881
 882 // ListAccounts get all accounts which name prfix is id.
 883 func (store *MockAccountStore) ListAccounts(id string) ([]*acc.Account, error) {
 884         accounts := []*acc.Account{}
 885         accountIter := store.db.IteratorPrefix(AccountIDKey(strings.TrimSpace(id)))
 886         defer accountIter.Release()
 887
 888         for accountIter.Next() {
 889                 account := new(acc.Account)
 890                 if err := json.Unmarshal(accountIter.Value(), &account); err != nil {
 891                         return nil, err
 892                 }
 893                 accounts = append(accounts, account)
 894         }
 895         return accounts, nil
 896 }
 897
 898 // ListControlPrograms get all local control programs
 899 func (store *MockAccountStore) ListControlPrograms() ([]*acc.CtrlProgram, error) {
 900         cps := []*acc.CtrlProgram{}
 901         cpIter := store.db.IteratorPrefix([]byte(ContractPrefix))
 902         defer cpIter.Release()
 903
 904         for cpIter.Next() {
 905                 cp := new(acc.CtrlProgram)
 906                 if err := json.Unmarshal(cpIter.Value(), cp); err != nil {
 907                         return nil, err
 908                 }
 909                 cps = append(cps, cp)
 910         }
 911         return cps, nil
 912 }
 913
 914 // ListUTXOs get utxos by accountID
 915 func (store *MockAccountStore) ListUTXOs() ([]*acc.UTXO, error) {
 916         utxoIter := store.db.IteratorPrefix([]byte(UTXOPrefix))
 917         defer utxoIter.Release()
 918
 919         utxos := []*acc.UTXO{}
 920         for utxoIter.Next() {
 921                 utxo := new(acc.UTXO)
 922                 if err := json.Unmarshal(utxoIter.Value(), utxo); err != nil {
 923                         return nil, err
 924                 }
 925                 utxos = append(utxos, utxo)
 926         }
 927         return utxos, nil
 928 }
 929
 930 // SetAccount set account account ID, account alias and raw account.
 931 func (store *MockAccountStore) SetAccount(account *acc.Account) error {
 932         rawAccount, err := json.Marshal(account)
 933         if err != nil {
 934                 return acc.ErrMarshalAccount
 935         }
 936
 937         batch := store.db.NewBatch()
 938         if store.batch != nil {
 939                 batch = store.batch
 940         }
 941
 942         batch.Set(AccountIDKey(account.ID), rawAccount)
 943         batch.Set(accountAliasKey(account.Alias), []byte(account.ID))
 944
 945         if store.batch == nil {
 946                 batch.Write()
 947         }
 948         return nil
 949 }
 950
 951 // SetAccountIndex update account index
 952 func (store *MockAccountStore) SetAccountIndex(account *acc.Account) {
 953         currentIndex := store.GetAccountIndex(account.XPubs)
 954         if account.KeyIndex > currentIndex {
 955                 if store.batch == nil {
 956                         store.db.Set(accountIndexKey(account.XPubs), common.Unit64ToBytes(account.KeyIndex))
 957                 } else {
 958                         store.batch.Set(accountIndexKey(account.XPubs), common.Unit64ToBytes(account.KeyIndex))
 959                 }
 960         }
 961 }
 962
 963 // SetBip44ContractIndex set contract index
 964 func (store *MockAccountStore) SetBip44ContractIndex(accountID string, change bool, index uint64) {
 965         if store.batch == nil {
 966                 store.db.Set(Bip44ContractIndexKey(accountID, change), common.Unit64ToBytes(index))
 967         } else {
 968                 store.batch.Set(Bip44ContractIndexKey(accountID, change), common.Unit64ToBytes(index))
 969         }
 970 }
 971
 972 // SetCoinbaseArbitrary set coinbase arbitrary
 973 func (store *MockAccountStore) SetCoinbaseArbitrary(arbitrary []byte) {
 974         if store.batch == nil {
 975                 store.db.Set([]byte(CoinbaseAbKey), arbitrary)
 976         } else {
 977                 store.batch.Set([]byte(CoinbaseAbKey), arbitrary)
 978         }
 979 }
 980
 981 // SetContractIndex set contract index
 982 func (store *MockAccountStore) SetContractIndex(accountID string, index uint64) {
 983         if store.batch == nil {
 984                 store.db.Set(contractIndexKey(accountID), common.Unit64ToBytes(index))
 985         } else {
 986                 store.batch.Set(contractIndexKey(accountID), common.Unit64ToBytes(index))
 987         }
 988 }
 989
 990 // SetControlProgram set raw program
 991 func (store *MockAccountStore) SetControlProgram(hash bc.Hash, program *acc.CtrlProgram) error {
 992         accountCP, err := json.Marshal(program)
 993         if err != nil {
 994                 return err
 995         }
 996         if store.batch == nil {
 997                 store.db.Set(ContractKey(hash), accountCP)
 998         } else {
 999                 store.batch.Set(ContractKey(hash), accountCP)
1000         }
1001         return nil
1002 }
1003
1004 // SetMiningAddress set mining address
1005 func (store *MockAccountStore) SetMiningAddress(program *acc.CtrlProgram) error {
1006         rawProgram, err := json.Marshal(program)
1007         if err != nil {
1008                 return err
1009         }
1010
1011         if store.batch == nil {
1012                 store.db.Set([]byte(MiningAddressKey), rawProgram)
1013         } else {
1014                 store.batch.Set([]byte(MiningAddressKey), rawProgram)
1015         }
1016         return nil
1017 }
1018
1019 // SetStandardUTXO set standard utxo
1020 func (store *MockAccountStore) SetStandardUTXO(outputID bc.Hash, utxo *acc.UTXO) error {
1021         data, err := json.Marshal(utxo)
1022         if err != nil {
1023                 return err
1024         }
1025         if store.batch == nil {
1026                 store.db.Set(StandardUTXOKey(outputID), data)
1027         } else {
1028                 store.batch.Set(StandardUTXOKey(outputID), data)
1029         }
1030         return nil
1031 }