test/wallet_test_util.go

   1 package test
   2
   3 import (
   4         "fmt"
   5         "io/ioutil"
   6         "os"
   7         "path"
   8         "reflect"
   9
  10         "github.com/vapor/account"
  11         "github.com/vapor/asset"
  12         "github.com/vapor/blockchain/pseudohsm"
  13         "github.com/vapor/blockchain/signers"
  14         "github.com/vapor/crypto/ed25519/chainkd"
  15         "github.com/vapor/database"
  16         dbm "github.com/vapor/database/leveldb"
  17         "github.com/vapor/event"
  18         "github.com/vapor/protocol"
  19         "github.com/vapor/protocol/bc/types"
  20         w "github.com/vapor/wallet"
  21 )
  22
  23 type walletTestConfig struct {
  24         Keys       []*keyInfo     `json:"keys"`
  25         Accounts   []*accountInfo `json:"accounts"`
  26         Blocks     []*wtBlock     `json:"blocks"`
  27         RollbackTo uint64         `json:"rollback_to"`
  28 }
  29
  30 type keyInfo struct {
  31         Name     string `json:"name"`
  32         Password string `json:"password"`
  33 }
  34
  35 type accountInfo struct {
  36         Name   string   `json:"name"`
  37         Keys   []string `json:"keys"`
  38         Quorum int      `json:"quorum"`
  39 }
  40
  41 type wtBlock struct {
  42         CoinbaseAccount string            `json:"coinbase_account"`
  43         Transactions    []*wtTransaction  `json:"transactions"`
  44         PostStates      []*accountBalance `json:"post_states"`
  45         Append          uint64            `json:"append"`
  46 }
  47
  48 func (b *wtBlock) create(ctx *walletTestContext) (*types.Block, error) {
  49         transactions := make([]*types.Tx, 0, len(b.Transactions))
  50         for _, t := range b.Transactions {
  51                 tx, err := t.create(ctx)
  52                 if err != nil {
  53                         continue
  54                 }
  55                 transactions = append(transactions, tx)
  56         }
  57         return ctx.newBlock(transactions, b.CoinbaseAccount)
  58 }
  59
  60 func (b *wtBlock) verifyPostStates(ctx *walletTestContext) error {
  61         for _, state := range b.PostStates {
  62                 balance, err := ctx.getBalance(state.AccountAlias, state.AssetAlias)
  63                 if err != nil {
  64                         return err
  65                 }
  66
  67                 if balance != state.Amount {
  68                         return fmt.Errorf("AccountAlias: %s, AssetAlias: %s, expected: %d, have: %d", state.AccountAlias, state.AssetAlias, state.Amount, balance)
  69                 }
  70         }
  71         return nil
  72 }
  73
  74 type wtTransaction struct {
  75         Passwords []string  `json:"passwords"`
  76         Inputs    []*action `json:"inputs"`
  77         Outputs   []*action `json:"outputs"`
  78 }
  79
  80 // create signed transaction
  81 func (t *wtTransaction) create(ctx *walletTestContext) (*types.Tx, error) {
  82         generator := NewTxGenerator(ctx.Wallet.AccountMgr, ctx.Wallet.AssetReg, ctx.Wallet.Hsm)
  83         for _, input := range t.Inputs {
  84                 switch input.Type {
  85                 case "spend_account":
  86                         if err := generator.AddSpendInput(input.AccountAlias, input.AssetAlias, input.Amount); err != nil {
  87                                 return nil, err
  88                         }
  89                 }
  90         }
  91
  92         for _, output := range t.Outputs {
  93                 switch output.Type {
  94                 case "output":
  95                         if err := generator.AddTxOutput(output.AccountAlias, output.AssetAlias, output.Amount); err != nil {
  96                                 return nil, err
  97                         }
  98                 case "retire":
  99                         if err := generator.AddRetirement(output.AssetAlias, output.Amount); err != nil {
 100                                 return nil, err
 101                         }
 102                 }
 103         }
 104         return generator.Sign(t.Passwords)
 105 }
 106
 107 type action struct {
 108         Type         string `json:"type"`
 109         AccountAlias string `json:"name"`
 110         AssetAlias   string `json:"asset"`
 111         Amount       uint64 `json:"amount"`
 112 }
 113
 114 type accountBalance struct {
 115         AssetAlias   string `json:"asset"`
 116         AccountAlias string `json:"name"`
 117         Amount       uint64 `json:"amount"`
 118 }
 119
 120 type walletTestContext struct {
 121         Wallet *w.Wallet
 122         Chain  *protocol.Chain
 123 }
 124
 125 func (ctx *walletTestContext) createControlProgram(accountName string, change bool) (*account.CtrlProgram, error) {
 126         acc, err := ctx.Wallet.AccountMgr.FindByAlias(accountName)
 127         if err != nil {
 128                 return nil, err
 129         }
 130
 131         return ctx.Wallet.AccountMgr.CreateAddress(acc.ID, change)
 132 }
 133
 134 func (ctx *walletTestContext) getPubkey(keyAlias string) *chainkd.XPub {
 135         pubKeys := ctx.Wallet.Hsm.ListKeys()
 136         for i, key := range pubKeys {
 137                 if key.Alias == keyAlias {
 138                         return &pubKeys[i].XPub
 139                 }
 140         }
 141         return nil
 142 }
 143
 144 func (ctx *walletTestContext) newBlock(txs []*types.Tx, coinbaseAccount string) (*types.Block, error) {
 145         controlProgram, err := ctx.createControlProgram(coinbaseAccount, true)
 146         if err != nil {
 147                 return nil, err
 148         }
 149         return NewBlock(ctx.Chain, txs, controlProgram.ControlProgram)
 150 }
 151
 152 func (ctx *walletTestContext) createKey(name string, password string) error {
 153         _, _, err := ctx.Wallet.Hsm.XCreate(name, password, "en")
 154         return err
 155 }
 156
 157 func (ctx *walletTestContext) createAccount(name string, keys []string, quorum int) error {
 158         xpubs := []chainkd.XPub{}
 159         for _, alias := range keys {
 160                 xpub := ctx.getPubkey(alias)
 161                 if xpub == nil {
 162                         return fmt.Errorf("can't find pubkey for %s", alias)
 163                 }
 164                 xpubs = append(xpubs, *xpub)
 165         }
 166         _, err := ctx.Wallet.AccountMgr.Create(xpubs, quorum, name, signers.BIP0044)
 167         return err
 168 }
 169
 170 func (ctx *walletTestContext) update(block *types.Block) error {
 171         if _, err := ctx.Chain.ProcessBlock(block); err != nil {
 172                 return err
 173         }
 174         if err := ctx.Wallet.AttachBlock(block); err != nil {
 175                 return err
 176         }
 177         return nil
 178 }
 179
 180 func (ctx *walletTestContext) getBalance(accountAlias string, assetAlias string) (uint64, error) {
 181         balances, _ := ctx.Wallet.GetAccountBalances("", "")
 182         for _, balance := range balances {
 183                 if balance.Alias == accountAlias && balance.AssetAlias == assetAlias {
 184                         return balance.Amount, nil
 185                 }
 186         }
 187         return 0, nil
 188 }
 189
 190 func (ctx *walletTestContext) getAccBalances() map[string]map[string]uint64 {
 191         accBalances := make(map[string]map[string]uint64)
 192         balances, _ := ctx.Wallet.GetAccountBalances("", "")
 193         for _, balance := range balances {
 194                 if accBalance, ok := accBalances[balance.Alias]; ok {
 195                         if _, ok := accBalance[balance.AssetAlias]; ok {
 196                                 accBalance[balance.AssetAlias] += balance.Amount
 197                                 continue
 198                         }
 199                         accBalance[balance.AssetAlias] = balance.Amount
 200                         continue
 201                 }
 202                 accBalances[balance.Alias] = map[string]uint64{balance.AssetAlias: balance.Amount}
 203         }
 204         return accBalances
 205 }
 206
 207 func (ctx *walletTestContext) getDetachedBlocks(height uint64) ([]*types.Block, error) {
 208         currentHeight := ctx.Chain.BestBlockHeight()
 209         detachedBlocks := make([]*types.Block, 0, currentHeight-height)
 210         for i := currentHeight; i > height; i-- {
 211                 block, err := ctx.Chain.GetBlockByHeight(i)
 212                 if err != nil {
 213                         return detachedBlocks, err
 214                 }
 215                 detachedBlocks = append(detachedBlocks, block)
 216         }
 217         return detachedBlocks, nil
 218 }
 219
 220 func (ctx *walletTestContext) validateRollback(oldAccBalances map[string]map[string]uint64) error {
 221         accBalances := ctx.getAccBalances()
 222         if reflect.DeepEqual(oldAccBalances, accBalances) {
 223                 return nil
 224         }
 225         return fmt.Errorf("different account balances after rollback")
 226 }
 227
 228 func (cfg *walletTestConfig) Run() error {
 229         dirPath, err := ioutil.TempDir(".", "pseudo_hsm")
 230         if err != nil {
 231                 return err
 232         }
 233         defer os.RemoveAll(dirPath)
 234         hsm, err := pseudohsm.New(dirPath)
 235         if err != nil {
 236                 return err
 237         }
 238
 239         db := dbm.NewDB("wallet_test_db", "leveldb", path.Join(dirPath, "wallet_test_db"))
 240         chain, _, _, err := MockChain(db)
 241         if err != nil {
 242                 return err
 243         }
 244         walletDB := dbm.NewDB("wallet", "leveldb", path.Join(dirPath, "wallet_db"))
 245         walletStore := database.NewWalletStore(walletDB)
 246         accountStore := database.NewAccountStore(walletDB)
 247         accountManager := account.NewManager(accountStore, chain)
 248         assets := asset.NewRegistry(walletDB, chain)
 249         dispatcher := event.NewDispatcher()
 250         wallet, err := w.NewWallet(walletStore, accountManager, assets, hsm, chain, dispatcher, false)
 251         if err != nil {
 252                 return err
 253         }
 254         ctx := &walletTestContext{
 255                 Wallet: wallet,
 256                 Chain:  chain,
 257         }
 258
 259         for _, key := range cfg.Keys {
 260                 if err := ctx.createKey(key.Name, key.Password); err != nil {
 261                         return err
 262                 }
 263         }
 264
 265         for _, acc := range cfg.Accounts {
 266                 if err := ctx.createAccount(acc.Name, acc.Keys, acc.Quorum); err != nil {
 267                         return err
 268                 }
 269         }
 270
 271         var accBalances map[string]map[string]uint64
 272         var rollbackBlock *types.Block
 273         for _, blk := range cfg.Blocks {
 274                 block, err := blk.create(ctx)
 275                 if err != nil {
 276                         return err
 277                 }
 278                 if err := ctx.update(block); err != nil {
 279                         return err
 280                 }
 281                 if err := blk.verifyPostStates(ctx); err != nil {
 282                         return err
 283                 }
 284                 if block.Height <= cfg.RollbackTo && cfg.RollbackTo <= block.Height+blk.Append {
 285                         accBalances = ctx.getAccBalances()
 286                         rollbackBlock = block
 287                 }
 288                 if err := AppendBlocks(ctx.Chain, blk.Append); err != nil {
 289                         return err
 290                 }
 291         }
 292
 293         if rollbackBlock == nil {
 294                 return nil
 295         }
 296
 297         // rollback and validate
 298         detachedBlocks, err := ctx.getDetachedBlocks(rollbackBlock.Height)
 299         if err != nil {
 300                 return err
 301         }
 302
 303         forkPath, err := ioutil.TempDir(".", "forked_chain")
 304         if err != nil {
 305                 return err
 306         }
 307
 308         forkedChain, err := declChain(forkPath, ctx.Chain, rollbackBlock.Height, ctx.Chain.BestBlockHeight()+1)
 309         defer os.RemoveAll(forkPath)
 310         if err != nil {
 311                 return err
 312         }
 313
 314         if err := merge(forkedChain, ctx.Chain); err != nil {
 315                 return err
 316         }
 317
 318         for _, block := range detachedBlocks {
 319                 if err := ctx.Wallet.DetachBlock(block); err != nil {
 320                         return err
 321                 }
 322         }
 323         return ctx.validateRollback(accBalances)
 324 }