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