wallet/wallet.go

   1 package wallet
   2
   3 import (
   4         "encoding/json"
   5
   6         log "github.com/sirupsen/logrus"
   7         "github.com/tendermint/tmlibs/db"
   8
   9         "github.com/bytom/account"
  10         "github.com/bytom/asset"
  11         "github.com/bytom/blockchain/pseudohsm"
  12         "github.com/bytom/protocol"
  13         "github.com/bytom/protocol/bc"
  14         "github.com/bytom/protocol/bc/types"
  15 )
  16
  17 //SINGLE single sign
  18 const SINGLE = 1
  19
  20 var walletKey = []byte("walletInfo")
  21
  22 //StatusInfo is base valid block info to handle orphan block rollback
  23 type StatusInfo struct {
  24         WorkHeight uint64
  25         WorkHash   bc.Hash
  26         BestHeight uint64
  27         BestHash   bc.Hash
  28 }
  29
  30 //Wallet is related to storing account unspent outputs
  31 type Wallet struct {
  32         DB         db.DB
  33         status     StatusInfo
  34         AccountMgr *account.Manager
  35         AssetReg   *asset.Registry
  36         Hsm        *pseudohsm.HSM
  37         chain      *protocol.Chain
  38         rescanCh   chan struct{}
  39 }
  40
  41 //NewWallet return a new wallet instance
  42 func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, hsm *pseudohsm.HSM, chain *protocol.Chain) (*Wallet, error) {
  43         w := &Wallet{
  44                 DB:         walletDB,
  45                 AccountMgr: account,
  46                 AssetReg:   asset,
  47                 chain:      chain,
  48                 Hsm:        hsm,
  49                 rescanCh:   make(chan struct{}, 1),
  50         }
  51
  52         if err := w.loadWalletInfo(); err != nil {
  53                 return nil, err
  54         }
  55
  56         go w.walletUpdater()
  57
  58         return w, nil
  59 }
  60
  61 //GetWalletInfo return stored wallet info and nil,if error,
  62 //return initial wallet info and err
  63 func (w *Wallet) loadWalletInfo() error {
  64         if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
  65                 return json.Unmarshal(rawWallet, &w.status)
  66         }
  67
  68         block, err := w.chain.GetBlockByHeight(0)
  69         if err != nil {
  70                 return err
  71         }
  72         return w.AttachBlock(block)
  73 }
  74
  75 func (w *Wallet) commitWalletInfo(batch db.Batch) error {
  76         rawWallet, err := json.Marshal(w.status)
  77         if err != nil {
  78                 log.WithField("err", err).Error("save wallet info")
  79                 return err
  80         }
  81
  82         batch.Set(walletKey, rawWallet)
  83         batch.Write()
  84         return nil
  85 }
  86
  87 // AttachBlock attach a new block
  88 func (w *Wallet) AttachBlock(block *types.Block) error {
  89         if block.PreviousBlockHash != w.status.WorkHash {
  90                 log.Warn("wallet skip attachBlock due to status hash not equal to previous hash")
  91                 return nil
  92         }
  93
  94         blockHash := block.Hash()
  95         txStatus, err := w.chain.GetTransactionStatus(&blockHash)
  96         if err != nil {
  97                 return err
  98         }
  99
 100         storeBatch := w.DB.NewBatch()
 101         w.indexTransactions(storeBatch, block, txStatus)
 102         w.buildAccountUTXOs(storeBatch, block, txStatus)
 103
 104         w.status.WorkHeight = block.Height
 105         w.status.WorkHash = block.Hash()
 106         if w.status.WorkHeight >= w.status.BestHeight {
 107                 w.status.BestHeight = w.status.WorkHeight
 108                 w.status.BestHash = w.status.WorkHash
 109         }
 110         return w.commitWalletInfo(storeBatch)
 111 }
 112
 113 // DetachBlock detach a block and rollback state
 114 func (w *Wallet) DetachBlock(block *types.Block) error {
 115         blockHash := block.Hash()
 116         txStatus, err := w.chain.GetTransactionStatus(&blockHash)
 117         if err != nil {
 118                 return err
 119         }
 120
 121         storeBatch := w.DB.NewBatch()
 122         w.reverseAccountUTXOs(storeBatch, block, txStatus)
 123         w.deleteTransactions(storeBatch, w.status.BestHeight)
 124
 125         w.status.BestHeight = block.Height - 1
 126         w.status.BestHash = block.PreviousBlockHash
 127
 128         if w.status.WorkHeight > w.status.BestHeight {
 129                 w.status.WorkHeight = w.status.BestHeight
 130                 w.status.WorkHash = w.status.BestHash
 131         }
 132
 133         return w.commitWalletInfo(storeBatch)
 134 }
 135
 136 //WalletUpdate process every valid block and reverse every invalid block which need to rollback
 137 func (w *Wallet) walletUpdater() {
 138         for {
 139                 w.getRescanNotification()
 140                 for !w.chain.InMainChain(w.status.BestHash) {
 141                         block, err := w.chain.GetBlockByHash(&w.status.BestHash)
 142                         if err != nil {
 143                                 log.WithField("err", err).Error("walletUpdater GetBlockByHash")
 144                                 return
 145                         }
 146
 147                         if err := w.DetachBlock(block); err != nil {
 148                                 log.WithField("err", err).Error("walletUpdater detachBlock")
 149                                 return
 150                         }
 151                 }
 152
 153                 block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight + 1)
 154                 if block == nil {
 155                         <-w.chain.BlockWaiter(w.status.WorkHeight + 1)
 156                         continue
 157                 }
 158
 159                 if err := w.AttachBlock(block); err != nil {
 160                         log.WithField("err", err).Error("walletUpdater stop")
 161                         return
 162                 }
 163         }
 164 }
 165
 166 func (w *Wallet) RescanBlocks() {
 167         select {
 168         case w.rescanCh <- struct{}{}:
 169         default:
 170                 return
 171         }
 172 }
 173
 174 func (w *Wallet) getRescanNotification() {
 175         select {
 176         case <-w.rescanCh:
 177                 block, _ := w.chain.GetBlockByHeight(0)
 178                 w.status.WorkHash = bc.Hash{}
 179                 w.AttachBlock(block)
 180         default:
 181                 return
 182         }
 183 }
 184
 185 func (w *Wallet) createProgram(account *account.Account, XPub *pseudohsm.XPub, index uint64) error {
 186         for i := uint64(0); i < index; i++ {
 187                 if _, err := w.AccountMgr.CreateAddress(nil, account.ID, false); err != nil {
 188                         return err
 189                 }
 190         }
 191         return nil
 192 }