OSDN Git Service

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