OSDN Git Service

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