database/wallet_store.go

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