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