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/bytom/account"
  11         "github.com/bytom/asset"
  12         "github.com/bytom/blockchain/pseudohsm"
  13         "github.com/bytom/protocol"
  14         "github.com/bytom/protocol/bc"
  15         "github.com/bytom/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         rescanCh   chan struct{}
  43 }
  44
  45 //NewWallet return a new wallet instance
  46 func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, hsm *pseudohsm.HSM, chain *protocol.Chain) (*Wallet, error) {
  47         w := &Wallet{
  48                 DB:         walletDB,
  49                 AccountMgr: account,
  50                 AssetReg:   asset,
  51                 chain:      chain,
  52                 Hsm:        hsm,
  53                 rescanCh:   make(chan struct{}, 1),
  54         }
  55
  56         if err := w.loadWalletInfo(); err != nil {
  57                 return nil, err
  58         }
  59
  60         go w.walletUpdater()
  61         return w, nil
  62 }
  63
  64 //GetWalletInfo return stored wallet info and nil,if error,
  65 //return initial wallet info and err
  66 func (w *Wallet) loadWalletInfo() error {
  67         if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
  68                 return json.Unmarshal(rawWallet, &w.status)
  69         }
  70
  71         block, err := w.chain.GetBlockByHeight(0)
  72         if err != nil {
  73                 return err
  74         }
  75         return w.AttachBlock(block)
  76 }
  77
  78 func (w *Wallet) commitWalletInfo(batch db.Batch) error {
  79         rawWallet, err := json.Marshal(w.status)
  80         if err != nil {
  81                 log.WithField("err", err).Error("save wallet info")
  82                 return err
  83         }
  84
  85         batch.Set(walletKey, rawWallet)
  86         batch.Write()
  87         return nil
  88 }
  89
  90 // AttachBlock attach a new block
  91 func (w *Wallet) AttachBlock(block *types.Block) error {
  92         w.rw.Lock()
  93         defer w.rw.Unlock()
  94
  95         if block.PreviousBlockHash != w.status.WorkHash {
  96                 log.Warn("wallet skip attachBlock due to status hash not equal to previous hash")
  97                 return nil
  98         }
  99
 100         blockHash := block.Hash()
 101         txStatus, err := w.chain.GetTransactionStatus(&blockHash)
 102         if err != nil {
 103                 return err
 104         }
 105
 106         storeBatch := w.DB.NewBatch()
 107         w.indexTransactions(storeBatch, block, txStatus)
 108         w.attachUtxos(storeBatch, block, txStatus)
 109
 110         w.status.WorkHeight = block.Height
 111         w.status.WorkHash = block.Hash()
 112         if w.status.WorkHeight >= w.status.BestHeight {
 113                 w.status.BestHeight = w.status.WorkHeight
 114                 w.status.BestHash = w.status.WorkHash
 115         }
 116         return w.commitWalletInfo(storeBatch)
 117 }
 118
 119 // DetachBlock detach a block and rollback state
 120 func (w *Wallet) DetachBlock(block *types.Block) error {
 121         w.rw.Lock()
 122         defer w.rw.Unlock()
 123
 124         blockHash := block.Hash()
 125         txStatus, err := w.chain.GetTransactionStatus(&blockHash)
 126         if err != nil {
 127                 return err
 128         }
 129
 130         storeBatch := w.DB.NewBatch()
 131         w.detachUtxos(storeBatch, block, txStatus)
 132         w.deleteTransactions(storeBatch, w.status.BestHeight)
 133
 134         w.status.BestHeight = block.Height - 1
 135         w.status.BestHash = block.PreviousBlockHash
 136
 137         if w.status.WorkHeight > w.status.BestHeight {
 138                 w.status.WorkHeight = w.status.BestHeight
 139                 w.status.WorkHash = w.status.BestHash
 140         }
 141
 142         return w.commitWalletInfo(storeBatch)
 143 }
 144
 145 //WalletUpdate process every valid block and reverse every invalid block which need to rollback
 146 func (w *Wallet) walletUpdater() {
 147         for {
 148                 w.getRescanNotification()
 149                 for !w.chain.InMainChain(w.status.BestHash) {
 150                         block, err := w.chain.GetBlockByHash(&w.status.BestHash)
 151                         if err != nil {
 152                                 log.WithField("err", err).Error("walletUpdater GetBlockByHash")
 153                                 return
 154                         }
 155
 156                         if err := w.DetachBlock(block); err != nil {
 157                                 log.WithField("err", err).Error("walletUpdater detachBlock stop")
 158                                 return
 159                         }
 160                 }
 161
 162                 block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight + 1)
 163                 if block == nil {
 164                         w.walletBlockWaiter()
 165                         continue
 166                 }
 167
 168                 if err := w.AttachBlock(block); err != nil {
 169                         log.WithField("err", err).Error("walletUpdater AttachBlock stop")
 170                         return
 171                 }
 172         }
 173 }
 174
 175 //RescanBlocks provide a trigger to rescan blocks
 176 func (w *Wallet) RescanBlocks() {
 177         select {
 178         case w.rescanCh <- struct{}{}:
 179         default:
 180                 return
 181         }
 182 }
 183
 184 func (w *Wallet) getRescanNotification() {
 185         select {
 186         case <-w.rescanCh:
 187                 w.setRescanStatus()
 188         default:
 189                 return
 190         }
 191 }
 192
 193 func (w *Wallet) setRescanStatus() {
 194         block, _ := w.chain.GetBlockByHeight(0)
 195         w.status.WorkHash = bc.Hash{}
 196         w.AttachBlock(block)
 197 }
 198
 199 func (w *Wallet) walletBlockWaiter() {
 200         select {
 201         case <-w.chain.BlockWaiter(w.status.WorkHeight + 1):
 202         case <-w.rescanCh:
 203                 w.setRescanStatus()
 204         }
 205 }
 206
 207 // GetWalletStatusInfo return current wallet StatusInfo
 208 func (w *Wallet) GetWalletStatusInfo() StatusInfo {
 209         w.rw.RLock()
 210         defer w.rw.RUnlock()
 211
 212         return w.status
 213 }