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                 case "issue":
  89                         _, err := ctx.createAsset(input.AccountAlias, input.AssetAlias)
  90                         if err != nil {
  91                                 return nil, err
  92                         }
  93                         if err := generator.AddIssuanceInput(input.AssetAlias, input.Amount); err != nil {
  94                                 return nil, err
  95                         }
  96                 }
  97         }
  98
  99         for _, output := range t.Outputs {
 100                 switch output.Type {
 101                 case "output":
 102                         if err := generator.AddTxOutput(output.AccountAlias, output.AssetAlias, output.Amount); err != nil {
 103                                 return nil, err
 104                         }
 105                 case "retire":
 106                         if err := generator.AddRetirement(output.AssetAlias, output.Amount); err != nil {
 107                                 return nil, err
 108                         }
 109                 }
 110         }
 111         return generator.Sign(t.Passwords)
 112 }
 113
 114 type action struct {
 115         Type         string `json:"type"`
 116         AccountAlias string `json:"name"`
 117         AssetAlias   string `json:"asset"`
 118         Amount       uint64 `json:"amount"`
 119 }
 120
 121 type accountBalance struct {
 122         AssetAlias   string `json:"asset"`
 123         AccountAlias string `json:"name"`
 124         Amount       uint64 `json:"amount"`
 125 }
 126
 127 type walletTestContext struct {
 128         Wallet *w.Wallet
 129         Chain  *protocol.Chain
 130 }
 131
 132 func (ctx *walletTestContext) createControlProgram(accountName string, change bool) (*account.CtrlProgram, error) {
 133         acc, err := ctx.Wallet.AccountMgr.FindByAlias(accountName)
 134         if err != nil {
 135                 return nil, err
 136         }
 137
 138         return ctx.Wallet.AccountMgr.CreateAddress(acc.ID, change)
 139 }
 140
 141 func (ctx *walletTestContext) getPubkey(keyAlias string) *chainkd.XPub {
 142         pubKeys := ctx.Wallet.Hsm.ListKeys()
 143         for i, key := range pubKeys {
 144                 if key.Alias == keyAlias {
 145                         return &pubKeys[i].XPub
 146                 }
 147         }
 148         return nil
 149 }
 150
 151 func (ctx *walletTestContext) createAsset(accountAlias string, assetAlias string) (*asset.Asset, error) {
 152         acc, err := ctx.Wallet.AccountMgr.FindByAlias(accountAlias)
 153         if err != nil {
 154                 return nil, err
 155         }
 156         return ctx.Wallet.AssetReg.Define(acc.XPubs, len(acc.XPubs), nil, 0, assetAlias, nil)
 157 }
 158
 159 func (ctx *walletTestContext) newBlock(txs []*types.Tx, coinbaseAccount string) (*types.Block, error) {
 160         controlProgram, err := ctx.createControlProgram(coinbaseAccount, true)
 161         if err != nil {
 162                 return nil, err
 163         }
 164         return NewBlock(ctx.Chain, txs, controlProgram.ControlProgram)
 165 }
 166
 167 func (ctx *walletTestContext) createKey(name string, password string) error {
 168         _, _, err := ctx.Wallet.Hsm.XCreate(name, password, "en")
 169         return err
 170 }
 171
 172 func (ctx *walletTestContext) createAccount(name string, keys []string, quorum int) error {
 173         xpubs := []chainkd.XPub{}
 174         for _, alias := range keys {
 175                 xpub := ctx.getPubkey(alias)
 176                 if xpub == nil {
 177                         return fmt.Errorf("can't find pubkey for %s", alias)
 178                 }
 179                 xpubs = append(xpubs, *xpub)
 180         }
 181         _, err := ctx.Wallet.AccountMgr.Create(xpubs, quorum, name, signers.BIP0044)
 182         return err
 183 }
 184
 185 func (ctx *walletTestContext) update(block *types.Block) error {
 186         if _, err := ctx.Chain.ProcessBlock(block); err != nil {
 187                 return err
 188         }
 189         if err := ctx.Wallet.AttachBlock(block); err != nil {
 190                 return err
 191         }
 192         return nil
 193 }
 194
 195 func (ctx *walletTestContext) getBalance(accountAlias string, assetAlias string) (uint64, error) {
 196         balances, _ := ctx.Wallet.GetAccountBalances("", "")
 197         for _, balance := range balances {
 198                 if balance.Alias == accountAlias && balance.AssetAlias == assetAlias {
 199                         return balance.Amount, nil
 200                 }
 201         }
 202         return 0, nil
 203 }
 204
 205 func (ctx *walletTestContext) getAccBalances() map[string]map[string]uint64 {
 206         accBalances := make(map[string]map[string]uint64)
 207         balances, _ := ctx.Wallet.GetAccountBalances("", "")
 208         for _, balance := range balances {
 209                 if accBalance, ok := accBalances[balance.Alias]; ok {
 210                         if _, ok := accBalance[balance.AssetAlias]; ok {
 211                                 accBalance[balance.AssetAlias] += balance.Amount
 212                                 continue
 213                         }
 214                         accBalance[balance.AssetAlias] = balance.Amount
 215                         continue
 216                 }
 217                 accBalances[balance.Alias] = map[string]uint64{balance.AssetAlias: balance.Amount}
 218         }
 219         return accBalances
 220 }
 221
 222 func (ctx *walletTestContext) getDetachedBlocks(height uint64) ([]*types.Block, error) {
 223         currentHeight := ctx.Chain.BestBlockHeight()
 224         detachedBlocks := make([]*types.Block, 0, currentHeight-height)
 225         for i := currentHeight; i > height; i-- {
 226                 block, err := ctx.Chain.GetBlockByHeight(i)
 227                 if err != nil {
 228                         return detachedBlocks, err
 229                 }
 230                 detachedBlocks = append(detachedBlocks, block)
 231         }
 232         return detachedBlocks, nil
 233 }
 234
 235 func (ctx *walletTestContext) validateRollback(oldAccBalances map[string]map[string]uint64) error {
 236         accBalances := ctx.getAccBalances()
 237         if reflect.DeepEqual(oldAccBalances, accBalances) {
 238                 return nil
 239         }
 240         return fmt.Errorf("different account balances after rollback")
 241 }
 242
 243 func (cfg *walletTestConfig) Run() error {
 244         dirPath, err := ioutil.TempDir(".", "pseudo_hsm")
 245         if err != nil {
 246                 return err
 247         }
 248         defer os.RemoveAll(dirPath)
 249         hsm, err := pseudohsm.New(dirPath)
 250         if err != nil {
 251                 return err
 252         }
 253
 254         db := dbm.NewDB("wallet_test_db", "leveldb", path.Join(dirPath, "wallet_test_db"))
 255         chain, _, _, err := MockChain(db)
 256         if err != nil {
 257                 return err
 258         }
 259         walletDB := dbm.NewDB("wallet", "leveldb", path.Join(dirPath, "wallet_db"))
 260         accountManager := account.NewManager(walletDB, chain)
 261         assets := asset.NewRegistry(walletDB, chain)
 262         dispatcher := event.NewDispatcher()
 263         wallet, err := w.NewWallet(walletDB, accountManager, assets, hsm, chain, dispatcher, false)
 264         if err != nil {
 265                 return err
 266         }
 267         ctx := &walletTestContext{
 268                 Wallet: wallet,
 269                 Chain:  chain,
 270         }
 271
 272         for _, key := range cfg.Keys {
 273                 if err := ctx.createKey(key.Name, key.Password); err != nil {
 274                         return err
 275                 }
 276         }
 277
 278         for _, acc := range cfg.Accounts {
 279                 if err := ctx.createAccount(acc.Name, acc.Keys, acc.Quorum); err != nil {
 280                         return err
 281                 }
 282         }
 283
 284         var accBalances map[string]map[string]uint64
 285         var rollbackBlock *types.Block
 286         for _, blk := range cfg.Blocks {
 287                 block, err := blk.create(ctx)
 288                 if err != nil {
 289                         return err
 290                 }
 291                 if err := ctx.update(block); err != nil {
 292                         return err
 293                 }
 294                 if err := blk.verifyPostStates(ctx); err != nil {
 295                         return err
 296                 }
 297                 if block.Height <= cfg.RollbackTo && cfg.RollbackTo <= block.Height+blk.Append {
 298                         accBalances = ctx.getAccBalances()
 299                         rollbackBlock = block
 300                 }
 301                 if err := AppendBlocks(ctx.Chain, blk.Append); err != nil {
 302                         return err
 303                 }
 304         }
 305
 306         if rollbackBlock == nil {
 307                 return nil
 308         }
 309
 310         // rollback and validate
 311         detachedBlocks, err := ctx.getDetachedBlocks(rollbackBlock.Height)
 312         if err != nil {
 313                 return err
 314         }
 315
 316         forkPath, err := ioutil.TempDir(".", "forked_chain")
 317         if err != nil {
 318                 return err
 319         }
 320
 321         forkedChain, err := declChain(forkPath, ctx.Chain, rollbackBlock.Height, ctx.Chain.BestBlockHeight()+1)
 322         defer os.RemoveAll(forkPath)
 323         if err != nil {
 324                 return err
 325         }
 326
 327         if err := merge(forkedChain, ctx.Chain); err != nil {
 328                 return err
 329         }
 330
 331         for _, block := range detachedBlocks {
 332                 if err := ctx.Wallet.DetachBlock(block); err != nil {
 333                         return err
 334                 }
 335         }
 336         return ctx.validateRollback(accBalances)
 337 }