wallet/wallet.go

   1 package wallet
   2
   3 import (
   4         "encoding/json"
   5         "fmt"
   6
   7         log "github.com/sirupsen/logrus"
   8         "github.com/tendermint/go-wire/data/base58"
   9         "github.com/tendermint/tmlibs/db"
  10
  11         "github.com/bytom/account"
  12         "github.com/bytom/asset"
  13         "github.com/bytom/blockchain/pseudohsm"
  14         "github.com/bytom/crypto/ed25519/chainkd"
  15         "github.com/bytom/crypto/sha3pool"
  16         "github.com/bytom/protocol"
  17         "github.com/bytom/protocol/bc"
  18         "github.com/bytom/protocol/bc/types"
  19 )
  20
  21 //SINGLE single sign
  22 const SINGLE = 1
  23
  24 //RecoveryIndex walletdb recovery cp number
  25 const RecoveryIndex = 5000
  26
  27 var walletKey = []byte("walletInfo")
  28 var privKeyKey = []byte("keysInfo")
  29
  30 //StatusInfo is base valid block info to handle orphan block rollback
  31 type StatusInfo struct {
  32         WorkHeight uint64
  33         WorkHash   bc.Hash
  34         BestHeight uint64
  35         BestHash   bc.Hash
  36 }
  37
  38 //KeyInfo is key import status
  39 type KeyInfo struct {
  40         Alias    string       `json:"alias"`
  41         XPub     chainkd.XPub `json:"xpub"`
  42         Percent  uint8        `json:"percent"`
  43         Complete bool         `json:"complete"`
  44 }
  45
  46 //Wallet is related to storing account unspent outputs
  47 type Wallet struct {
  48         DB             db.DB
  49         status         StatusInfo
  50         AccountMgr     *account.Manager
  51         AssetReg       *asset.Registry
  52         Hsm            *pseudohsm.HSM
  53         chain          *protocol.Chain
  54         rescanProgress chan struct{}
  55         ImportPrivKey  bool
  56         keysInfo       []KeyInfo
  57 }
  58
  59 //NewWallet return a new wallet instance
  60 func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, hsm *pseudohsm.HSM, chain *protocol.Chain) (*Wallet, error) {
  61         w := &Wallet{
  62                 DB:             walletDB,
  63                 AccountMgr:     account,
  64                 AssetReg:       asset,
  65                 chain:          chain,
  66                 Hsm:            hsm,
  67                 rescanProgress: make(chan struct{}, 1),
  68                 keysInfo:       make([]KeyInfo, 0),
  69         }
  70
  71         if err := w.loadWalletInfo(); err != nil {
  72                 return nil, err
  73         }
  74
  75         if err := w.loadKeysInfo(); err != nil {
  76                 return nil, err
  77         }
  78
  79         w.ImportPrivKey = w.getImportKeyFlag()
  80
  81         go w.walletUpdater()
  82
  83         return w, nil
  84 }
  85
  86 //GetWalletInfo return stored wallet info and nil,if error,
  87 //return initial wallet info and err
  88 func (w *Wallet) loadWalletInfo() error {
  89         if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
  90                 return json.Unmarshal(rawWallet, &w.status)
  91         }
  92
  93         block, err := w.chain.GetBlockByHeight(0)
  94         if err != nil {
  95                 return err
  96         }
  97         return w.AttachBlock(block)
  98 }
  99
 100 func (w *Wallet) commitWalletInfo(batch db.Batch) error {
 101         rawWallet, err := json.Marshal(w.status)
 102         if err != nil {
 103                 log.WithField("err", err).Error("save wallet info")
 104                 return err
 105         }
 106
 107         batch.Set(walletKey, rawWallet)
 108         batch.Write()
 109         return nil
 110 }
 111
 112 //GetWalletInfo return stored wallet info and nil,if error,
 113 //return initial wallet info and err
 114 func (w *Wallet) loadKeysInfo() error {
 115         if rawKeyInfo := w.DB.Get(privKeyKey); rawKeyInfo != nil {
 116                 json.Unmarshal(rawKeyInfo, &w.keysInfo)
 117                 return nil
 118         }
 119         return nil
 120 }
 121
 122 func (w *Wallet) commitkeysInfo() error {
 123         rawKeysInfo, err := json.Marshal(w.keysInfo)
 124         if err != nil {
 125                 log.WithField("err", err).Error("save wallet info")
 126                 return err
 127         }
 128         w.DB.Set(privKeyKey, rawKeysInfo)
 129         return nil
 130 }
 131
 132 func (w *Wallet) getImportKeyFlag() bool {
 133         for _, v := range w.keysInfo {
 134                 if v.Complete == false {
 135                         return true
 136                 }
 137         }
 138         return false
 139 }
 140
 141 // AttachBlock attach a new block
 142 func (w *Wallet) AttachBlock(block *types.Block) error {
 143         if block.PreviousBlockHash != w.status.WorkHash {
 144                 log.Warn("wallet skip attachBlock due to status hash not equal to previous hash")
 145                 return nil
 146         }
 147
 148         blockHash := block.Hash()
 149         txStatus, err := w.chain.GetTransactionStatus(&blockHash)
 150         if err != nil {
 151                 return err
 152         }
 153
 154         storeBatch := w.DB.NewBatch()
 155         w.indexTransactions(storeBatch, block, txStatus)
 156         w.buildAccountUTXOs(storeBatch, block, txStatus)
 157
 158         w.status.WorkHeight = block.Height
 159         w.status.WorkHash = block.Hash()
 160         if w.status.WorkHeight >= w.status.BestHeight {
 161                 w.status.BestHeight = w.status.WorkHeight
 162                 w.status.BestHash = w.status.WorkHash
 163         }
 164         return w.commitWalletInfo(storeBatch)
 165 }
 166
 167 // DetachBlock detach a block and rollback state
 168 func (w *Wallet) DetachBlock(block *types.Block) error {
 169         blockHash := block.Hash()
 170         txStatus, err := w.chain.GetTransactionStatus(&blockHash)
 171         if err != nil {
 172                 return err
 173         }
 174
 175         storeBatch := w.DB.NewBatch()
 176         w.reverseAccountUTXOs(storeBatch, block, txStatus)
 177         w.deleteTransactions(storeBatch, w.status.BestHeight)
 178
 179         w.status.BestHeight = block.Height - 1
 180         w.status.BestHash = block.PreviousBlockHash
 181
 182         if w.status.WorkHeight > w.status.BestHeight {
 183                 w.status.WorkHeight = w.status.BestHeight
 184                 w.status.WorkHash = w.status.BestHash
 185         }
 186
 187         return w.commitWalletInfo(storeBatch)
 188 }
 189
 190 //WalletUpdate process every valid block and reverse every invalid block which need to rollback
 191 func (w *Wallet) walletUpdater() {
 192         for {
 193                 getRescanNotification(w)
 194                 checkRescanStatus(w)
 195                 for !w.chain.InMainChain(w.status.BestHash) {
 196                         block, err := w.chain.GetBlockByHash(&w.status.BestHash)
 197                         if err != nil {
 198                                 log.WithField("err", err).Error("walletUpdater GetBlockByHash")
 199                                 return
 200                         }
 201
 202                         if err := w.DetachBlock(block); err != nil {
 203                                 log.WithField("err", err).Error("walletUpdater detachBlock")
 204                                 return
 205                         }
 206                 }
 207
 208                 block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight + 1)
 209                 if block == nil {
 210                         <-w.chain.BlockWaiter(w.status.WorkHeight + 1)
 211                         continue
 212                 }
 213
 214                 if err := w.AttachBlock(block); err != nil {
 215                         log.WithField("err", err).Error("walletUpdater stop")
 216                         return
 217                 }
 218         }
 219 }
 220
 221 func getRescanNotification(w *Wallet) {
 222         select {
 223         case <-w.rescanProgress:
 224                 w.status.WorkHeight = 0
 225                 block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight)
 226                 w.status.WorkHash = block.Hash()
 227         default:
 228                 return
 229         }
 230 }
 231
 232 // ExportAccountPrivKey exports the account private key as a WIF for encoding as a string
 233 // in the Wallet Import Formt.
 234 func (w *Wallet) ExportAccountPrivKey(xpub chainkd.XPub, auth string) (*string, error) {
 235         xprv, err := w.Hsm.LoadChainKDKey(xpub, auth)
 236         if err != nil {
 237                 return nil, err
 238         }
 239         var hashed [32]byte
 240         sha3pool.Sum256(hashed[:], xprv[:])
 241
 242         tmp := append(xprv[:], hashed[:4]...)
 243         res := base58.Encode(tmp)
 244         return &res, nil
 245 }
 246
 247 // ImportAccountPrivKey imports the account key in the Wallet Import Formt.
 248 func (w *Wallet) ImportAccountPrivKey(xprv chainkd.XPrv, keyAlias, auth string, index uint64, accountAlias string) (*pseudohsm.XPub, error) {
 249         if w.Hsm.HasAlias(keyAlias) {
 250                 return nil, pseudohsm.ErrDuplicateKeyAlias
 251         }
 252         if w.Hsm.HasKey(xprv) {
 253                 return nil, pseudohsm.ErrDuplicateKey
 254         }
 255
 256         if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
 257                 return nil, account.ErrDuplicateAlias
 258         }
 259
 260         xpub, _, err := w.Hsm.ImportXPrvKey(auth, keyAlias, xprv)
 261         if err != nil {
 262                 return nil, err
 263         }
 264
 265         newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias, nil)
 266         if err != nil {
 267                 return nil, err
 268         }
 269         if err := w.recoveryAccountWalletDB(newAccount, xpub, index, keyAlias); err != nil {
 270                 return nil, err
 271         }
 272         return xpub, nil
 273 }
 274
 275 // ImportAccountXpubKey imports the account key in the Wallet Import Formt.
 276 func (w *Wallet) ImportAccountXpubKey(xpubIndex int, xpub pseudohsm.XPub, cpIndex uint64) error {
 277         accountAlias := fmt.Sprintf("recovery_%d", xpubIndex)
 278
 279         if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
 280                 return account.ErrDuplicateAlias
 281         }
 282
 283         newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias, nil)
 284         if err != nil {
 285                 return err
 286         }
 287
 288         return w.recoveryAccountWalletDB(newAccount, &xpub, cpIndex, xpub.Alias)
 289 }
 290
 291 func (w *Wallet) recoveryAccountWalletDB(account *account.Account, XPub *pseudohsm.XPub, index uint64, keyAlias string) error {
 292         if err := w.createProgram(account, XPub, index); err != nil {
 293                 return err
 294         }
 295         w.ImportPrivKey = true
 296         tmp := KeyInfo{
 297                 Alias:    keyAlias,
 298                 XPub:     XPub.XPub,
 299                 Complete: false,
 300         }
 301         w.keysInfo = append(w.keysInfo, tmp)
 302         w.commitkeysInfo()
 303         w.rescanBlocks()
 304
 305         return nil
 306 }
 307
 308 func (w *Wallet) createProgram(account *account.Account, XPub *pseudohsm.XPub, index uint64) error {
 309         for i := uint64(0); i < index; i++ {
 310                 if _, err := w.AccountMgr.CreateAddress(nil, account.ID); err != nil {
 311                         return err
 312                 }
 313         }
 314         return nil
 315 }
 316
 317 func (w *Wallet) rescanBlocks() {
 318         select {
 319         case w.rescanProgress <- struct{}{}:
 320         default:
 321                 return
 322         }
 323 }
 324
 325 //GetRescanStatus return key import rescan status
 326 func (w *Wallet) GetRescanStatus() ([]KeyInfo, error) {
 327         keysInfo := make([]KeyInfo, len(w.keysInfo))
 328
 329         if rawKeyInfo := w.DB.Get(privKeyKey); rawKeyInfo != nil {
 330                 if err := json.Unmarshal(rawKeyInfo, &keysInfo); err != nil {
 331                         return nil, err
 332                 }
 333         }
 334
 335         var status StatusInfo
 336         if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
 337                 if err := json.Unmarshal(rawWallet, &status); err != nil {
 338                         return nil, err
 339                 }
 340         }
 341
 342         for i := range keysInfo {
 343                 if keysInfo[i].Complete == true || status.BestHeight == 0 {
 344                         keysInfo[i].Percent = 100
 345                         continue
 346                 }
 347
 348                 keysInfo[i].Percent = uint8(status.WorkHeight * 100 / status.BestHeight)
 349                 if keysInfo[i].Percent == 100 {
 350                         keysInfo[i].Complete = true
 351                 }
 352         }
 353         return keysInfo, nil
 354 }
 355
 356 //checkRescanStatus mark private key import process `Complete` if rescan finished
 357 func checkRescanStatus(w *Wallet) {
 358         if !w.ImportPrivKey || w.status.WorkHeight < w.status.BestHeight {
 359                 return
 360         }
 361
 362         w.ImportPrivKey = false
 363         for _, keyInfo := range w.keysInfo {
 364                 keyInfo.Complete = true
 365         }
 366         w.commitkeysInfo()
 367 }