gecko/xputtest/chainwallet/wallet.go

// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package chainwallet

import (
	"errors"
	"fmt"

	"github.com/ava-labs/gecko/ids"
	"github.com/ava-labs/gecko/snow"
	"github.com/ava-labs/gecko/utils/crypto"
	"github.com/ava-labs/gecko/utils/logging"
	"github.com/ava-labs/gecko/vms/spchainvm"
)

// Wallet is a holder for keys and UTXOs.
type Wallet struct {
	networkID uint32
	chainID   ids.ID

	log logging.Logger

	keychain   *spchainvm.Keychain            // Mapping from public address to the SigningKeys
	accountSet map[[20]byte]spchainvm.Account // Mapping from addresses to accounts
	balance    uint64

	txs []*spchainvm.Tx
}

// NewWallet ...
func NewWallet(log logging.Logger, networkID uint32, chainID ids.ID) *Wallet {
	return &Wallet{
		networkID:  networkID,
		chainID:    chainID,
		log:        log,
		keychain:   spchainvm.NewKeychain(networkID, chainID),
		accountSet: make(map[[20]byte]spchainvm.Account),
	}
}

// CreateAddress returns a brand new address! Ready to receive funds!
func (w *Wallet) CreateAddress() (ids.ShortID, error) {
	sk, err := w.keychain.New()
	if err != nil {
		return ids.ShortID{}, err
	}
	return sk.PublicKey().Address(), nil
}

// ImportKey imports a private key into this wallet
func (w *Wallet) ImportKey(sk *crypto.PrivateKeySECP256K1R) { w.keychain.Add(sk) }

// AddAccount adds a new account to this wallet, if this wallet can spend it.
func (w *Wallet) AddAccount(account spchainvm.Account) {
	if account.Balance() > 0 {
		w.accountSet[account.ID().Key()] = account
		w.balance += account.Balance()
	}
}

// Balance returns the amount of the assets in this wallet
func (w *Wallet) Balance() uint64 { return w.balance }

// GenerateTxs generates the transactions that will be sent
// during the test
// Generate them all on test initialization so tx generation is not bottleneck
// in testing
func (w *Wallet) GenerateTxs(numTxs int) error {
	w.log.Info("Generating %d transactions", numTxs)

	ctx := snow.DefaultContextTest()
	ctx.NetworkID = w.networkID
	ctx.ChainID = w.chainID

	frequency := numTxs / 50
	if frequency > 1000 {
		frequency = 1000
	}

	w.txs = make([]*spchainvm.Tx, numTxs)
	for i := range w.txs {
		tx, err := w.MakeTx()
		if err != nil {
			return err
		}

		if numGenerated := i + 1; numGenerated%frequency == 0 {
			w.log.Info("Generated %d out of %d transactions", numGenerated, numTxs)
		}

		w.txs[i] = tx
	}

	w.log.Info("Finished generating %d transactions", numTxs)

	return nil
}

// NextTx returns the next tx to be sent as part of xput test
func (w *Wallet) NextTx() *spchainvm.Tx {
	if len(w.txs) == 0 {
		return nil
	}
	tx := w.txs[0]
	w.txs = w.txs[1:]
	return tx
}

// MakeTx creates a new transaction and update the state to after the tx is accepted
func (w *Wallet) MakeTx() (*spchainvm.Tx, error) {
	ctx := snow.DefaultContextTest()
	ctx.NetworkID = w.networkID
	ctx.ChainID = w.chainID

	for _, account := range w.accountSet {
		accountID := account.ID()
		key, exists := w.keychain.Get(accountID)
		if !exists {
			return nil, errors.New("missing account")
		}

		amount := uint64(1)
		tx, sendAccount, err := account.CreateTx(amount, accountID, ctx, key)
		if err != nil {
			continue
		}

		newAccount, err := sendAccount.Receive(tx, ctx)
		if err != nil {
			return nil, err
		}
		w.accountSet[accountID.Key()] = newAccount
		return tx, nil
	}
	return nil, errors.New("empty")
}

func (w Wallet) String() string {
	return fmt.Sprintf(
		"Keychain:\n"+
			"%s",
		w.keychain.PrefixedString("    "))
}