wallet/wallet.go

   1 package wallet
   2
   3 import (
   4         "encoding/json"
   5         "sync"
   6
   7         log "github.com/sirupsen/logrus"
   8         "github.com/tendermint/tmlibs/db"
   9
  10         "github.com/vapor/account"
  11         "github.com/vapor/asset"
  12         "github.com/vapor/blockchain/pseudohsm"
  13         "github.com/vapor/common"
  14         "github.com/vapor/protocol"
  15         "github.com/vapor/protocol/bc"
  16         "github.com/vapor/protocol/bc/types"
  17 )
  18
  19 const (
  20         //SINGLE single sign
  21         SINGLE = 1
  22 )
  23
  24 var walletKey = []byte("walletInfo")
  25
  26 //StatusInfo is base valid block info to handle orphan block rollback
  27 type StatusInfo struct {
  28         WorkHeight uint64
  29         WorkHash   bc.Hash
  30         BestHeight uint64
  31         BestHash   bc.Hash
  32 }
  33
  34 //Wallet is related to storing account unspent outputs
  35 type Wallet struct {
  36         DB          db.DB
  37         rw          sync.RWMutex
  38         status      StatusInfo
  39         AccountMgr  *account.Manager
  40         AssetReg    *asset.Registry
  41         Hsm         *pseudohsm.HSM
  42         chain       *protocol.Chain
  43         RecoveryMgr *recoveryManager
  44         rescanCh    chan struct{}
  45         dposAddress common.Address
  46 }
  47
  48 //NewWallet return a new wallet instance
  49 func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, hsm *pseudohsm.HSM, chain *protocol.Chain, dposAddress common.Address) (*Wallet, error) {
  50         w := &Wallet{
  51                 DB:          walletDB,
  52                 AccountMgr:  account,
  53                 AssetReg:    asset,
  54                 chain:       chain,
  55                 Hsm:         hsm,
  56                 RecoveryMgr: newRecoveryManager(walletDB, account),
  57                 rescanCh:    make(chan struct{}, 1),
  58                 dposAddress: dposAddress,
  59         }
  60
  61         if err := w.loadWalletInfo(); err != nil {
  62                 return nil, err
  63         }
  64
  65         if err := w.RecoveryMgr.LoadStatusInfo(); err != nil {
  66                 return nil, err
  67         }
  68
  69         go w.walletUpdater()
  70         go w.delUnconfirmedTx()
  71         return w, nil
  72 }
  73
  74 //GetWalletInfo return stored wallet info and nil,if error,
  75 //return initial wallet info and err
  76 func (w *Wallet) loadWalletInfo() error {
  77         if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
  78                 return json.Unmarshal(rawWallet, &w.status)
  79         }
  80
  81         block, err := w.chain.GetBlockByHeight(0)
  82         if err != nil {
  83                 return err
  84         }
  85         return w.AttachBlock(block)
  86 }
  87
  88 func (w *Wallet) commitWalletInfo(batch db.Batch) error {
  89         rawWallet, err := json.Marshal(w.status)
  90         if err != nil {
  91                 log.WithField("err", err).Error("save wallet info")
  92                 return err
  93         }
  94
  95         batch.Set(walletKey, rawWallet)
  96         batch.Write()
  97         return nil
  98 }
  99
 100 // AttachBlock attach a new block
 101 func (w *Wallet) AttachBlock(block *types.Block) error {
 102         w.rw.Lock()
 103         defer w.rw.Unlock()
 104
 105         if block.PreviousBlockHash != w.status.WorkHash {
 106                 log.Warn("wallet skip attachBlock due to status hash not equal to previous hash")
 107                 return nil
 108         }
 109
 110         blockHash := block.Hash()
 111         txStatus, err := w.chain.GetTransactionStatus(&blockHash)
 112         if err != nil {
 113                 return err
 114         }
 115
 116         if err := w.RecoveryMgr.FilterRecoveryTxs(block); err != nil {
 117                 return err
 118         }
 119
 120         storeBatch := w.DB.NewBatch()
 121         if err := w.indexTransactions(storeBatch, block, txStatus); err != nil {
 122                 return err
 123         }
 124
 125         w.attachUtxos(storeBatch, block, txStatus)
 126         w.status.WorkHeight = block.Height
 127         w.status.WorkHash = block.Hash()
 128         if w.status.WorkHeight >= w.status.BestHeight {
 129                 w.status.BestHeight = w.status.WorkHeight
 130                 w.status.BestHash = w.status.WorkHash
 131         }
 132         return w.commitWalletInfo(storeBatch)
 133 }
 134
 135 // DetachBlock detach a block and rollback state
 136 func (w *Wallet) DetachBlock(block *types.Block) error {
 137         w.rw.Lock()
 138         defer w.rw.Unlock()
 139
 140         blockHash := block.Hash()
 141         txStatus, err := w.chain.GetTransactionStatus(&blockHash)
 142         if err != nil {
 143                 return err
 144         }
 145
 146         storeBatch := w.DB.NewBatch()
 147         w.detachUtxos(storeBatch, block, txStatus)
 148         w.deleteTransactions(storeBatch, w.status.BestHeight)
 149
 150         w.status.BestHeight = block.Height - 1
 151         w.status.BestHash = block.PreviousBlockHash
 152
 153         if w.status.WorkHeight > w.status.BestHeight {
 154                 w.status.WorkHeight = w.status.BestHeight
 155                 w.status.WorkHash = w.status.BestHash
 156         }
 157
 158         return w.commitWalletInfo(storeBatch)
 159 }
 160
 161 //WalletUpdate process every valid block and reverse every invalid block which need to rollback
 162 func (w *Wallet) walletUpdater() {
 163         for {
 164                 w.getRescanNotification()
 165                 for !w.chain.InMainChain(w.status.BestHash) {
 166                         block, err := w.chain.GetBlockByHash(&w.status.BestHash)
 167                         if err != nil {
 168                                 log.WithField("err", err).Error("walletUpdater GetBlockByHash")
 169                                 return
 170                         }
 171
 172                         if err := w.DetachBlock(block); err != nil {
 173                                 log.WithField("err", err).Error("walletUpdater detachBlock stop")
 174                                 return
 175                         }
 176                 }
 177
 178                 block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight + 1)
 179                 if block == nil {
 180                         w.walletBlockWaiter()
 181                         continue
 182                 }
 183
 184                 if err := w.AttachBlock(block); err != nil {
 185                         log.WithField("err", err).Error("walletUpdater AttachBlock stop")
 186                         return
 187                 }
 188         }
 189 }
 190
 191 //RescanBlocks provide a trigger to rescan blocks
 192 func (w *Wallet) RescanBlocks() {
 193         select {
 194         case w.rescanCh <- struct{}{}:
 195         default:
 196                 return
 197         }
 198 }
 199
 200 // deleteAccountTxs deletes all txs in wallet
 201 func (w *Wallet) deleteAccountTxs() {
 202         storeBatch := w.DB.NewBatch()
 203
 204         txIter := w.DB.IteratorPrefix([]byte(TxPrefix))
 205         defer txIter.Release()
 206
 207         for txIter.Next() {
 208                 storeBatch.Delete(txIter.Key())
 209         }
 210
 211         txIndexIter := w.DB.IteratorPrefix([]byte(TxIndexPrefix))
 212         defer txIndexIter.Release()
 213
 214         for txIndexIter.Next() {
 215                 storeBatch.Delete(txIndexIter.Key())
 216         }
 217
 218         storeBatch.Write()
 219 }
 220
 221 // DeleteAccount deletes account matching accountID, then rescan wallet
 222 func (w *Wallet) DeleteAccount(accountID string) (err error) {
 223         w.rw.Lock()
 224         defer w.rw.Unlock()
 225
 226         if err := w.AccountMgr.DeleteAccount(accountID); err != nil {
 227                 return err
 228         }
 229
 230         w.deleteAccountTxs()
 231         w.RescanBlocks()
 232         return nil
 233 }
 234
 235 func (w *Wallet) UpdateAccountAlias(accountID string, newAlias string) (err error) {
 236         w.rw.Lock()
 237         defer w.rw.Unlock()
 238
 239         if err := w.AccountMgr.UpdateAccountAlias(accountID, newAlias); err != nil {
 240                 return err
 241         }
 242
 243         w.deleteAccountTxs()
 244         w.RescanBlocks()
 245         return nil
 246 }
 247
 248 func (w *Wallet) getRescanNotification() {
 249         select {
 250         case <-w.rescanCh:
 251                 w.setRescanStatus()
 252         default:
 253                 return
 254         }
 255 }
 256
 257 func (w *Wallet) setRescanStatus() {
 258         block, _ := w.chain.GetBlockByHeight(0)
 259         w.status.WorkHash = bc.Hash{}
 260         w.AttachBlock(block)
 261 }
 262
 263 func (w *Wallet) walletBlockWaiter() {
 264         select {
 265         case <-w.chain.BlockWaiter(w.status.WorkHeight + 1):
 266         case <-w.rescanCh:
 267                 w.setRescanStatus()
 268         }
 269 }
 270
 271 // GetWalletStatusInfo return current wallet StatusInfo
 272 func (w *Wallet) GetWalletStatusInfo() StatusInfo {
 273         w.rw.RLock()
 274         defer w.rw.RUnlock()
 275
 276         return w.status
 277 }