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