test/wallet_test_util.go

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