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