test/mock/wallet_store.go

   1 package mock
   2
   3 import (
   4         "encoding/binary"
   5         "encoding/json"
   6         "fmt"
   7         "sort"
   8
   9         acc "github.com/vapor/account"
  10         "github.com/vapor/asset"
  11         "github.com/vapor/blockchain/query"
  12         "github.com/vapor/blockchain/signers"
  13         "github.com/vapor/crypto/ed25519/chainkd"
  14         "github.com/vapor/crypto/sha3pool"
  15         dbm "github.com/vapor/database/leveldb"
  16         "github.com/vapor/errors"
  17         "github.com/vapor/protocol/bc"
  18 )
  19
  20 const (
  21         colon            = 0x3a
  22         utxoPrefix  byte = iota //UTXOPrefix is StandardUTXOKey prefix
  23         sutxoPrefix             //SUTXOPrefix is ContractUTXOKey prefix
  24         contractPrefix
  25         contractIndexPrefix
  26         accountPrefix // AccountPrefix is account ID prefix
  27         accountAliasPrefix
  28         accountIndexPrefix
  29         txPrefix            //TxPrefix is wallet database transactions prefix
  30         txIndexPrefix       //TxIndexPrefix is wallet database tx index prefix
  31         unconfirmedTxPrefix //UnconfirmedTxPrefix is txpool unconfirmed transactions prefix
  32         globalTxIndexPrefix //GlobalTxIndexPrefix is wallet database global tx index prefix
  33         walletKey
  34         miningAddressKey
  35         coinbaseAbKey
  36         recoveryKey
  37 )
  38
  39 // leveldb key prefix
  40 var (
  41         UTXOPrefix  = []byte{utxoPrefix, colon}
  42         SUTXOPrefix = []byte{sutxoPrefix, colon}
  43         // ContractPrefix = []byte{contractPrefix, contractPrefix, colon}
  44         ContractPrefix      = "Contract:"
  45         ContractIndexPrefix = []byte{contractIndexPrefix, colon}
  46         AccountPrefix       = []byte{accountPrefix, colon} // AccountPrefix is account ID prefix
  47         AccountAliasPrefix  = []byte{accountAliasPrefix, colon}
  48         AccountIndexPrefix  = []byte{accountIndexPrefix, colon}
  49         TxPrefix            = []byte{txPrefix, colon}            //TxPrefix is wallet database transactions prefix
  50         TxIndexPrefix       = []byte{txIndexPrefix, colon}       //TxIndexPrefix is wallet database tx index prefix
  51         UnconfirmedTxPrefix = []byte{unconfirmedTxPrefix, colon} //UnconfirmedTxPrefix is txpool unconfirmed transactions prefix
  52         GlobalTxIndexPrefix = []byte{globalTxIndexPrefix, colon} //GlobalTxIndexPrefix is wallet database global tx index prefix
  53         WalletKey           = []byte{walletKey}
  54         MiningAddressKey    = []byte{miningAddressKey}
  55         CoinbaseAbKey       = []byte{coinbaseAbKey}
  56         RecoveryKey         = []byte{recoveryKey}
  57 )
  58
  59 // errors
  60 var (
  61         // ErrFindAccount        = errors.New("Failed to find account")
  62         errAccntTxIDNotFound = errors.New("account TXID not found")
  63         errGetAsset          = errors.New("Failed to find asset definition")
  64 )
  65
  66 func accountIndexKey(xpubs []chainkd.XPub) []byte {
  67         var hash [32]byte
  68         var xPubs []byte
  69         cpy := append([]chainkd.XPub{}, xpubs[:]...)
  70         sort.Sort(signers.SortKeys(cpy))
  71         for _, xpub := range cpy {
  72                 xPubs = append(xPubs, xpub[:]...)
  73         }
  74         sha3pool.Sum256(hash[:], xPubs)
  75         return append([]byte(AccountIndexPrefix), hash[:]...)
  76 }
  77
  78 func Bip44ContractIndexKey(accountID string, change bool) []byte {
  79         key := append([]byte(ContractIndexPrefix), accountID...)
  80         if change {
  81                 return append(key, []byte{1}...)
  82         }
  83         return append(key, []byte{0}...)
  84 }
  85
  86 // ContractKey account control promgram store prefix
  87 func ContractKey(hash bc.Hash) []byte {
  88         // h := hash.String()
  89         // return append([]byte(ContractPrefix), []byte(h)...)
  90         return append([]byte(ContractPrefix), hash.Bytes()...)
  91 }
  92
  93 // AccountIDKey account id store prefix
  94 func AccountIDKey(accountID string) []byte {
  95         return append([]byte(AccountPrefix), []byte(accountID)...)
  96 }
  97
  98 // StandardUTXOKey makes an account unspent outputs key to store
  99 func StandardUTXOKey(id bc.Hash) []byte {
 100         name := id.String()
 101         return append(UTXOPrefix, []byte(name)...)
 102 }
 103
 104 // ContractUTXOKey makes a smart contract unspent outputs key to store
 105 func ContractUTXOKey(id bc.Hash) []byte {
 106         name := id.String()
 107         return append(SUTXOPrefix, []byte(name)...)
 108 }
 109
 110 func calcDeleteKey(blockHeight uint64) []byte {
 111         return []byte(fmt.Sprintf("%s%016x", TxPrefix, blockHeight))
 112 }
 113
 114 func calcTxIndexKey(txID string) []byte {
 115         return append(TxIndexPrefix, []byte(txID)...)
 116 }
 117
 118 func calcAnnotatedKey(formatKey string) []byte {
 119         return append(TxPrefix, []byte(formatKey)...)
 120 }
 121
 122 func calcUnconfirmedTxKey(formatKey string) []byte {
 123         return append(UnconfirmedTxPrefix, []byte(formatKey)...)
 124 }
 125
 126 func calcGlobalTxIndexKey(txID string) []byte {
 127         return append(GlobalTxIndexPrefix, []byte(txID)...)
 128 }
 129
 130 func CalcGlobalTxIndex(blockHash *bc.Hash, position uint64) []byte {
 131         txIdx := make([]byte, 40)
 132         copy(txIdx[:32], blockHash.Bytes())
 133         binary.BigEndian.PutUint64(txIdx[32:], position)
 134         return txIdx
 135 }
 136
 137 func formatKey(blockHeight uint64, position uint32) string {
 138         return fmt.Sprintf("%016x%08x", blockHeight, position)
 139 }
 140
 141 func contractIndexKey(accountID string) []byte {
 142         return append([]byte(ContractIndexPrefix), []byte(accountID)...)
 143 }
 144
 145 func accountAliasKey(name string) []byte {
 146         return append([]byte(AccountAliasPrefix), []byte(name)...)
 147 }
 148
 149 // MockWalletStore store wallet using leveldb
 150 type MockWalletStore struct {
 151         walletDB dbm.DB
 152         batch    dbm.Batch
 153 }
 154
 155 // NewMockWalletStore create new MockWalletStore struct
 156 func NewMockWalletStore(db dbm.DB) *MockWalletStore {
 157         return &MockWalletStore{
 158                 walletDB: db,
 159                 batch:    nil,
 160         }
 161 }
 162
 163 // InitBatch initial batch
 164 func (store *MockWalletStore) InitBatch() error {
 165         if store.batch != nil {
 166                 return errors.New("MockWalletStore initail fail, store batch is not nil.")
 167         }
 168         store.batch = store.walletDB.NewBatch()
 169         return nil
 170 }
 171
 172 // CommitBatch commit batch
 173 func (store *MockWalletStore) CommitBatch() error {
 174         if store.batch == nil {
 175                 return errors.New("MockWalletStore commit fail, store batch is nil.")
 176         }
 177         store.batch.Write()
 178         store.batch = nil
 179         return nil
 180 }
 181
 182 // DeleteContractUTXO delete contract utxo by outputID
 183 func (store *MockWalletStore) DeleteContractUTXO(outputID bc.Hash) {
 184         if store.batch == nil {
 185                 store.walletDB.Delete(ContractUTXOKey(outputID))
 186         } else {
 187                 store.batch.Delete(ContractUTXOKey(outputID))
 188         }
 189 }
 190
 191 // DeleteRecoveryStatus delete recovery status
 192 func (store *MockWalletStore) DeleteRecoveryStatus() {
 193         if store.batch == nil {
 194                 store.walletDB.Delete(RecoveryKey)
 195         } else {
 196                 store.batch.Delete(RecoveryKey)
 197         }
 198 }
 199
 200 // DeleteTransactions delete transactions when orphan block rollback
 201 func (store *MockWalletStore) DeleteTransactions(height uint64) {
 202         batch := store.walletDB.NewBatch()
 203         if store.batch != nil {
 204                 batch = store.batch
 205         }
 206         txIter := store.walletDB.IteratorPrefix(calcDeleteKey(height))
 207         defer txIter.Release()
 208
 209         tmpTx := query.AnnotatedTx{}
 210         for txIter.Next() {
 211                 if err := json.Unmarshal(txIter.Value(), &tmpTx); err == nil {
 212                         batch.Delete(calcTxIndexKey(tmpTx.ID.String()))
 213                 }
 214                 batch.Delete(txIter.Key())
 215         }
 216         if store.batch == nil {
 217                 batch.Write()
 218         }
 219 }
 220
 221 // DeleteUnconfirmedTransaction delete unconfirmed tx by txID
 222 func (store *MockWalletStore) DeleteUnconfirmedTransaction(txID string) {
 223         if store.batch == nil {
 224                 store.walletDB.Delete(calcUnconfirmedTxKey(txID))
 225         } else {
 226                 store.batch.Delete(calcUnconfirmedTxKey(txID))
 227         }
 228 }
 229
 230 // DeleteWalletTransactions delete all txs in wallet
 231 func (store *MockWalletStore) DeleteWalletTransactions() {
 232         batch := store.walletDB.NewBatch()
 233         if store.batch != nil {
 234                 batch = store.batch
 235         }
 236         txIter := store.walletDB.IteratorPrefix([]byte(TxPrefix))
 237         defer txIter.Release()
 238
 239         for txIter.Next() {
 240                 batch.Delete(txIter.Key())
 241         }
 242
 243         txIndexIter := store.walletDB.IteratorPrefix([]byte(TxIndexPrefix))
 244         defer txIndexIter.Release()
 245
 246         for txIndexIter.Next() {
 247                 batch.Delete(txIndexIter.Key())
 248         }
 249         if store.batch == nil {
 250                 batch.Write()
 251         }
 252 }
 253
 254 // DeleteWalletUTXOs delete all txs in wallet
 255 func (store *MockWalletStore) DeleteWalletUTXOs() {
 256         batch := store.walletDB.NewBatch()
 257         if store.batch != nil {
 258                 batch = store.batch
 259         }
 260         ruIter := store.walletDB.IteratorPrefix([]byte(UTXOPrefix))
 261         defer ruIter.Release()
 262         for ruIter.Next() {
 263                 batch.Delete(ruIter.Key())
 264         }
 265
 266         suIter := store.walletDB.IteratorPrefix([]byte(SUTXOPrefix))
 267         defer suIter.Release()
 268         for suIter.Next() {
 269                 batch.Delete(suIter.Key())
 270         }
 271         if store.batch == nil {
 272                 batch.Write()
 273         }
 274 }
 275
 276 // GetAsset get asset by assetID
 277 func (store *MockWalletStore) GetAsset(assetID *bc.AssetID) (*asset.Asset, error) {
 278         definitionByte := store.walletDB.Get(asset.ExtAssetKey(assetID))
 279         if definitionByte == nil {
 280                 return nil, errGetAsset
 281         }
 282         definitionMap := make(map[string]interface{})
 283         if err := json.Unmarshal(definitionByte, &definitionMap); err != nil {
 284                 return nil, err
 285         }
 286         alias := assetID.String()
 287         externalAsset := &asset.Asset{
 288                 AssetID:           *assetID,
 289                 Alias:             &alias,
 290                 DefinitionMap:     definitionMap,
 291                 RawDefinitionByte: definitionByte,
 292         }
 293         return externalAsset, nil
 294 }
 295
 296 // GetControlProgram get raw program by hash
 297 func (store *MockWalletStore) GetControlProgram(hash bc.Hash) (*acc.CtrlProgram, error) {
 298         rawProgram := store.walletDB.Get(ContractKey(hash))
 299         if rawProgram == nil {
 300                 return nil, acc.ErrFindCtrlProgram
 301         }
 302         accountCP := new(acc.CtrlProgram)
 303         if err := json.Unmarshal(rawProgram, &accountCP); err != nil {
 304                 return nil, err
 305         }
 306         return accountCP, nil
 307 }
 308
 309 // GetGlobalTransactionIndex get global tx by txID
 310 func (store *MockWalletStore) GetGlobalTransactionIndex(txID string) []byte {
 311         return store.walletDB.Get(calcGlobalTxIndexKey(txID))
 312 }
 313
 314 // GetStandardUTXO get standard utxo by id
 315 func (store *MockWalletStore) GetStandardUTXO(outid bc.Hash) (*acc.UTXO, error) {
 316         rawUTXO := store.walletDB.Get(StandardUTXOKey(outid))
 317         if rawUTXO == nil {
 318                 return nil, fmt.Errorf("failed get standard UTXO, outputID: %s ", outid.String())
 319         }
 320         UTXO := new(acc.UTXO)
 321         if err := json.Unmarshal(rawUTXO, UTXO); err != nil {
 322                 return nil, err
 323         }
 324         return UTXO, nil
 325 }
 326
 327 // GetTransaction get tx by txid
 328 func (store *MockWalletStore) GetTransaction(txID string) (*query.AnnotatedTx, error) {
 329         formatKey := store.walletDB.Get(calcTxIndexKey(txID))
 330         if formatKey == nil {
 331                 return nil, errAccntTxIDNotFound
 332         }
 333         rawTx := store.walletDB.Get(calcAnnotatedKey(string(formatKey)))
 334         tx := new(query.AnnotatedTx)
 335         if err := json.Unmarshal(rawTx, tx); err != nil {
 336                 return nil, err
 337         }
 338         return tx, nil
 339 }
 340
 341 // GetUnconfirmedTransaction get unconfirmed tx by txID
 342 func (store *MockWalletStore) GetUnconfirmedTransaction(txID string) (*query.AnnotatedTx, error) {
 343         rawUnconfirmedTx := store.walletDB.Get(calcUnconfirmedTxKey(txID))
 344         if rawUnconfirmedTx == nil {
 345                 return nil, fmt.Errorf("failed get unconfirmed tx, txID: %s ", txID)
 346         }
 347         tx := new(query.AnnotatedTx)
 348         if err := json.Unmarshal(rawUnconfirmedTx, tx); err != nil {
 349                 return nil, err
 350         }
 351         return tx, nil
 352 }
 353
 354 // GetRecoveryStatus delete recovery status
 355 func (store *MockWalletStore) GetRecoveryStatus(recoveryKey []byte) []byte {
 356         return store.walletDB.Get(recoveryKey)
 357 }
 358
 359 // GetWalletInfo get wallet information
 360 func (store *MockWalletStore) GetWalletInfo() []byte {
 361         return store.walletDB.Get([]byte(WalletKey))
 362 }
 363
 364 // ListAccountUTXOs get all account unspent outputs
 365 func (store *MockWalletStore) ListAccountUTXOs(key string) ([]*acc.UTXO, error) {
 366         accountUtxoIter := store.walletDB.IteratorPrefix([]byte(key))
 367         defer accountUtxoIter.Release()
 368
 369         confirmedUTXOs := []*acc.UTXO{}
 370         for accountUtxoIter.Next() {
 371                 utxo := new(acc.UTXO)
 372                 if err := json.Unmarshal(accountUtxoIter.Value(), utxo); err != nil {
 373                         return nil, err
 374                 }
 375                 confirmedUTXOs = append(confirmedUTXOs, utxo)
 376         }
 377         return confirmedUTXOs, nil
 378 }
 379
 380 func (store *MockWalletStore) ListTransactions(accountID string, StartTxID string, count uint, unconfirmed bool) ([]*query.AnnotatedTx, error) {
 381         annotatedTxs := []*query.AnnotatedTx{}
 382         var startKey []byte
 383         preFix := TxPrefix
 384
 385         if StartTxID != "" {
 386                 if unconfirmed {
 387                         startKey = calcUnconfirmedTxKey(StartTxID)
 388                 } else {
 389                         formatKey := store.walletDB.Get(calcTxIndexKey(StartTxID))
 390                         if formatKey == nil {
 391                                 return nil, errAccntTxIDNotFound
 392                         }
 393                         startKey = calcAnnotatedKey(string(formatKey))
 394                 }
 395         }
 396
 397         if unconfirmed {
 398                 preFix = UnconfirmedTxPrefix
 399         }
 400
 401         itr := store.walletDB.IteratorPrefixWithStart([]byte(preFix), startKey, true)
 402         defer itr.Release()
 403
 404         for txNum := count; itr.Next() && txNum > 0; txNum-- {
 405                 annotatedTx := new(query.AnnotatedTx)
 406                 if err := json.Unmarshal(itr.Value(), &annotatedTx); err != nil {
 407                         return nil, err
 408                 }
 409                 annotatedTxs = append(annotatedTxs, annotatedTx)
 410         }
 411
 412         return annotatedTxs, nil
 413 }
 414
 415 // ListUnconfirmedTransactions get all unconfirmed txs
 416 func (store *MockWalletStore) ListUnconfirmedTransactions() ([]*query.AnnotatedTx, error) {
 417         annotatedTxs := []*query.AnnotatedTx{}
 418         txIter := store.walletDB.IteratorPrefix([]byte(UnconfirmedTxPrefix))
 419         defer txIter.Release()
 420
 421         for txIter.Next() {
 422                 annotatedTx := &query.AnnotatedTx{}
 423                 if err := json.Unmarshal(txIter.Value(), &annotatedTx); err != nil {
 424                         return nil, err
 425                 }
 426                 annotatedTxs = append(annotatedTxs, annotatedTx)
 427         }
 428         return annotatedTxs, nil
 429 }
 430
 431 // SetAssetDefinition set assetID and definition
 432 func (store *MockWalletStore) SetAssetDefinition(assetID *bc.AssetID, definition []byte) {
 433         if store.batch == nil {
 434                 store.walletDB.Set(asset.ExtAssetKey(assetID), definition)
 435         } else {
 436                 store.batch.Set(asset.ExtAssetKey(assetID), definition)
 437         }
 438 }
 439
 440 // SetContractUTXO set standard utxo
 441 func (store *MockWalletStore) SetContractUTXO(outputID bc.Hash, utxo *acc.UTXO) error {
 442         data, err := json.Marshal(utxo)
 443         if err != nil {
 444                 return err
 445         }
 446         if store.batch == nil {
 447                 store.walletDB.Set(ContractUTXOKey(outputID), data)
 448         } else {
 449                 store.batch.Set(ContractUTXOKey(outputID), data)
 450         }
 451         return nil
 452 }
 453
 454 // SetGlobalTransactionIndex set global tx index by blockhash and position
 455 func (store *MockWalletStore) SetGlobalTransactionIndex(globalTxID string, blockHash *bc.Hash, position uint64) {
 456         if store.batch == nil {
 457                 store.walletDB.Set(calcGlobalTxIndexKey(globalTxID), CalcGlobalTxIndex(blockHash, position))
 458         } else {
 459                 store.batch.Set(calcGlobalTxIndexKey(globalTxID), CalcGlobalTxIndex(blockHash, position))
 460         }
 461 }
 462
 463 // SetRecoveryStatus set recovery status
 464 func (store *MockWalletStore) SetRecoveryStatus(recoveryKey, rawStatus []byte) {
 465         if store.batch == nil {
 466                 store.walletDB.Set(recoveryKey, rawStatus)
 467         } else {
 468                 store.batch.Set(recoveryKey, rawStatus)
 469         }
 470 }
 471
 472 // SetTransaction set raw transaction by block height and tx position
 473 func (store *MockWalletStore) SetTransaction(height uint64, tx *query.AnnotatedTx) error {
 474         batch := store.walletDB.NewBatch()
 475         if store.batch != nil {
 476                 batch = store.batch
 477         }
 478
 479         rawTx, err := json.Marshal(tx)
 480         if err != nil {
 481                 return err
 482         }
 483         batch.Set(calcAnnotatedKey(formatKey(height, tx.Position)), rawTx)
 484         batch.Set(calcTxIndexKey(tx.ID.String()), []byte(formatKey(height, tx.Position)))
 485
 486         if store.batch == nil {
 487                 batch.Write()
 488         }
 489         return nil
 490 }
 491
 492 // SetUnconfirmedTransaction set unconfirmed tx by txID
 493 func (store *MockWalletStore) SetUnconfirmedTransaction(txID string, tx *query.AnnotatedTx) error {
 494         rawTx, err := json.Marshal(tx)
 495         if err != nil {
 496                 return err
 497         }
 498         if store.batch == nil {
 499                 store.walletDB.Set(calcUnconfirmedTxKey(txID), rawTx)
 500         } else {
 501                 store.batch.Set(calcUnconfirmedTxKey(txID), rawTx)
 502         }
 503         return nil
 504 }
 505
 506 // SetWalletInfo get wallet information
 507 func (store *MockWalletStore) SetWalletInfo(rawWallet []byte) {
 508         if store.batch == nil {
 509                 store.walletDB.Set([]byte(WalletKey), rawWallet)
 510         } else {
 511                 store.batch.Set([]byte(WalletKey), rawWallet)
 512         }
 513 }