test/wallet_test_util.go

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