Merge branch 'master' into atomic

api/keystore/service.go

@@ -21,10 +21,30 @@ import (
 	"github.com/ava-labs/gecko/vms/components/codec"
 
 	jsoncodec "github.com/ava-labs/gecko/utils/json"
+	zxcvbn "github.com/nbutton23/zxcvbn-go"
+)
+
+const (
+	// maxUserPassLen is the maximum length of the username or password allowed
+	maxUserPassLen = 1024
+
+	// requiredPassScore defines the score a password must achieve to be accepted
+	// as a password with strong characteristics by the zxcvbn package
+	//
+	// The scoring mechanism defined is as follows;
+	//
+	// 0 # too guessable: risky password. (guesses < 10^3)
+	// 1 # very guessable: protection from throttled online attacks. (guesses < 10^6)
+	// 2 # somewhat guessable: protection from unthrottled online attacks. (guesses < 10^8)
+	// 3 # safely unguessable: moderate protection from offline slow-hash scenario. (guesses < 10^10)
+	// 4 # very unguessable: strong protection from offline slow-hash scenario. (guesses >= 10^10)
+	requiredPassScore = 2
 )
 
 var (
 	errEmptyUsername     = errors.New("username can't be the empty string")
+	errUserPassMaxLength = fmt.Errorf("CreateUser call rejected due to username or password exceeding maximum length of %d chars", maxUserPassLen)
+	errWeakPassword      = errors.New("Failed to create user as the given password is too weak. A stronger password is one of 8 or more characters containing attributes of upper and lowercase letters, numbers, and/or special characters")
 )
 
 // KeyValuePair ...
@@ -114,7 +134,11 @@ func (ks *Keystore) CreateUser(_ *http.Request, args *CreateUserArgs, reply *Cre
 	ks.lock.Lock()
 	defer ks.lock.Unlock()
 
-	ks.log.Verbo("CreateUser called with %s", args.Username)
+	ks.log.Verbo("CreateUser called with %.*s", maxUserPassLen, args.Username)
+
+	if len(args.Username) > maxUserPassLen || len(args.Password) > maxUserPassLen {
+		return errUserPassMaxLength
+	}
 
 	if args.Username == "" {
 		return errEmptyUsername
@@ -123,6 +147,10 @@ func (ks *Keystore) CreateUser(_ *http.Request, args *CreateUserArgs, reply *Cre
 		return fmt.Errorf("user already exists: %s", args.Username)
 	}
 
+	if zxcvbn.PasswordStrength(args.Password, nil).Score < requiredPassScore {
+		return errWeakPassword
+	}
+
 	usr := &User{}
 	if err := usr.Initialize(args.Password); err != nil {
 		return err

api/keystore/service_test.go

@@ -5,6 +5,8 @@ package keystore
 
 import (
 	"bytes"
+	"fmt"
+	"math/rand"
 	"testing"
 
 	"github.com/ava-labs/gecko/database/memdb"
@@ -12,6 +14,12 @@ import (
 	"github.com/ava-labs/gecko/utils/logging"
 )
 
+var (
+	// strongPassword defines a password used for the following tests that
+	// scores high enough to pass the password strength scoring system
+	strongPassword = "N_+=_jJ;^(<;{4,:*m6CET}'&N;83FYK.wtNpwp-Jt"
+)
+
 func TestServiceListNoUsers(t *testing.T) {
 	ks := Keystore{}
 	ks.Initialize(logging.NoLog{}, memdb.New())
@@ -33,7 +41,7 @@ func TestServiceCreateUser(t *testing.T) {
 		reply := CreateUserReply{}
 		if err := ks.CreateUser(nil, &CreateUserArgs{
 			Username: "bob",
-			Password: "launch",
+			Password: strongPassword,
 		}, &reply); err != nil {
 			t.Fatal(err)
 		}
@@ -56,6 +64,78 @@ func TestServiceCreateUser(t *testing.T) {
 	}
 }
 
+// genStr returns a string of given length
+func genStr(n int) string {
+	b := make([]byte, n)
+	rand.Read(b)
+	return fmt.Sprintf("%x", b)[:n]
+}
+
+// TestServiceCreateUserArgsChecks generates excessively long usernames or
+// passwords to assure the santity checks on string length are not exceeded
+func TestServiceCreateUserArgsCheck(t *testing.T) {
+	ks := Keystore{}
+	ks.Initialize(logging.NoLog{}, memdb.New())
+
+	{
+		reply := CreateUserReply{}
+		err := ks.CreateUser(nil, &CreateUserArgs{
+			Username: genStr(maxUserPassLen + 1),
+			Password: strongPassword,
+		}, &reply)
+
+		if reply.Success || err != errUserPassMaxLength {
+			t.Fatal("User was created when it should have been rejected due to too long a Username, err =", err)
+		}
+	}
+
+	{
+		reply := CreateUserReply{}
+		err := ks.CreateUser(nil, &CreateUserArgs{
+			Username: "shortuser",
+			Password: genStr(maxUserPassLen + 1),
+		}, &reply)
+
+		if reply.Success || err != errUserPassMaxLength {
+			t.Fatal("User was created when it should have been rejected due to too long a Password, err =", err)
+		}
+	}
+
+	{
+		reply := ListUsersReply{}
+		if err := ks.ListUsers(nil, &ListUsersArgs{}, &reply); err != nil {
+			t.Fatal(err)
+		}
+
+		if len(reply.Users) > 0 {
+			t.Fatalf("A user exists when there should be none")
+		}
+	}
+}
+
+// TestServiceCreateUserWeakPassword tests creating a new user with a weak
+// password to ensure the password strength check is working
+func TestServiceCreateUserWeakPassword(t *testing.T) {
+	ks := Keystore{}
+	ks.Initialize(logging.NoLog{}, memdb.New())
+
+	{
+		reply := CreateUserReply{}
+		err := ks.CreateUser(nil, &CreateUserArgs{
+			Username: "bob",
+			Password: "weak",
+		}, &reply)
+
+		if err != errWeakPassword {
+			t.Error("Unexpected error occurred when testing weak password:", err)
+		}
+
+		if reply.Success {
+			t.Fatal("User was created when it should have been rejected due to weak password")
+		}
+	}
+}
+
 func TestServiceCreateDuplicate(t *testing.T) {
 	ks := Keystore{}
 	ks.Initialize(logging.NoLog{}, memdb.New())
@@ -64,7 +144,7 @@ func TestServiceCreateDuplicate(t *testing.T) {
 		reply := CreateUserReply{}
 		if err := ks.CreateUser(nil, &CreateUserArgs{
 			Username: "bob",
-			Password: "launch",
+			Password: strongPassword,
 		}, &reply); err != nil {
 			t.Fatal(err)
 		}
@@ -77,7 +157,7 @@ func TestServiceCreateDuplicate(t *testing.T) {
 		reply := CreateUserReply{}
 		if err := ks.CreateUser(nil, &CreateUserArgs{
 			Username: "bob",
-			Password: "launch!",
+			Password: strongPassword,
 		}, &reply); err == nil {
 			t.Fatalf("Should have errored due to the username already existing")
 		}
@@ -90,7 +170,7 @@ func TestServiceCreateUserNoName(t *testing.T) {
 
 	reply := CreateUserReply{}
 	if err := ks.CreateUser(nil, &CreateUserArgs{
-		Password: "launch",
+		Password: strongPassword,
 	}, &reply); err == nil {
 		t.Fatalf("Shouldn't have allowed empty username")
 	}
@@ -104,7 +184,7 @@ func TestServiceUseBlockchainDB(t *testing.T) {
 		reply := CreateUserReply{}
 		if err := ks.CreateUser(nil, &CreateUserArgs{
 			Username: "bob",
-			Password: "launch",
+			Password: strongPassword,
 		}, &reply); err != nil {
 			t.Fatal(err)
 		}
@@ -114,7 +194,7 @@ func TestServiceUseBlockchainDB(t *testing.T) {
 	}
 
 	{
-		db, err := ks.GetDatabase(ids.Empty, "bob", "launch")
+		db, err := ks.GetDatabase(ids.Empty, "bob", strongPassword)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -124,7 +204,7 @@ func TestServiceUseBlockchainDB(t *testing.T) {
 	}
 
 	{
-		db, err := ks.GetDatabase(ids.Empty, "bob", "launch")
+		db, err := ks.GetDatabase(ids.Empty, "bob", strongPassword)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -144,7 +224,7 @@ func TestServiceExportImport(t *testing.T) {
 		reply := CreateUserReply{}
 		if err := ks.CreateUser(nil, &CreateUserArgs{
 			Username: "bob",
-			Password: "launch",
+			Password: strongPassword,
 		}, &reply); err != nil {
 			t.Fatal(err)
 		}
@@ -154,7 +234,7 @@ func TestServiceExportImport(t *testing.T) {
 	}
 
 	{
-		db, err := ks.GetDatabase(ids.Empty, "bob", "launch")
+		db, err := ks.GetDatabase(ids.Empty, "bob", strongPassword)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -166,7 +246,7 @@ func TestServiceExportImport(t *testing.T) {
 	exportReply := ExportUserReply{}
 	if err := ks.ExportUser(nil, &ExportUserArgs{
 		Username: "bob",
-		Password: "launch",
+		Password: strongPassword,
 	}, &exportReply); err != nil {
 		t.Fatal(err)
 	}
@@ -178,7 +258,7 @@ func TestServiceExportImport(t *testing.T) {
 		reply := ImportUserReply{}
 		if err := newKS.ImportUser(nil, &ImportUserArgs{
 			Username: "bob",
-			Password: "launch",
+			Password: strongPassword,
 			User:     exportReply.User,
 		}, &reply); err != nil {
 			t.Fatal(err)
@@ -189,7 +269,7 @@ func TestServiceExportImport(t *testing.T) {
 	}
 
 	{
-		db, err := newKS.GetDatabase(ids.Empty, "bob", "launch")
+		db, err := newKS.GetDatabase(ids.Empty, "bob", strongPassword)
 		if err != nil {
 			t.Fatal(err)
 		}

genesis/genesis.go

@@ -504,6 +504,9 @@ func VMGenesis(networkID uint32, vmID ids.ID) *platformvm.CreateChainTx {
 	genesisBytes := Genesis(networkID)
 	genesis := platformvm.Genesis{}
 	platformvm.Codec.Unmarshal(genesisBytes, &genesis)
+	if err := genesis.Initialize(); err != nil {
+		panic(err)
+	}
 	for _, chain := range genesis.Chains {
 		if chain.VMID.Equals(vmID) {
 			return chain

ids/bag.go

@@ -95,7 +95,7 @@ func (b *Bag) Mode() (ID, int) { return b.mode, b.modeFreq }
 func (b *Bag) Threshold() Set { return b.metThreshold }
 
 // Filter returns the bag of ids with the same counts as this bag, except all
-// the ids in the returned bag must have the same bits in the range [start, end]
+// the ids in the returned bag must have the same bits in the range [start, end)
 // as id.
 func (b *Bag) Filter(start, end int, id ID) Bag {
 	newBag := Bag{}

snow/consensus/snowball/tree.go

@@ -409,13 +409,18 @@ func (u *unaryNode) RecordPoll(votes ids.Bag, reset bool) node {
 		u.snowball.RecordSuccessfulPoll()
 
 		if u.child != nil {
-			decidedPrefix := u.child.DecidedPrefix()
+			// We are guaranteed that u.commonPrefix will equal
-			filteredVotes := votes.Filter(u.commonPrefix, decidedPrefix, u.preference)
+			// u.child.DecidedPrefix(). Otherwise, there must have been a
+			// decision under this node, which isn't possible because
+			// beta1 <= beta2. That means that filtering the votes between
+			// u.commonPrefix and u.child.DecidedPrefix() would always result in
+			// the same set being returned.
+
 			// If I'm now decided, return my child
 			if u.Finalized() {
-				return u.child.RecordPoll(filteredVotes, u.shouldReset)
+				return u.child.RecordPoll(votes, u.shouldReset)
 			}
-			u.child = u.child.RecordPoll(filteredVotes, u.shouldReset)
+			u.child = u.child.RecordPoll(votes, u.shouldReset)
 			// The child's preference may have changed
 			u.preference = u.child.Preference()
 		}

snow/consensus/snowball/tree_test.go

@@ -256,7 +256,7 @@ func TestSnowballTrinary(t *testing.T) {
 	tree.RecordPoll(redBag)
 
 	if pref := tree.Preference(); !Blue.Equals(pref) {
-		t.Fatalf("Wrong preference. Expected %s got %s", Green, pref)
+		t.Fatalf("Wrong preference. Expected %s got %s", Blue, pref)
 	} else if tree.Finalized() {
 		t.Fatalf("Finalized too early")
 	}
@@ -533,3 +533,108 @@ func TestSnowballConsistent(t *testing.T) {
 		t.Fatalf("Network agreed on inconsistent values")
 	}
 }
+
+func TestSnowballFilterBinaryChildren(t *testing.T) {
+	c0000 := ids.NewID([32]byte{0b00000000})
+	c1000 := ids.NewID([32]byte{0b00000001})
+	c0100 := ids.NewID([32]byte{0b00000010})
+	c0010 := ids.NewID([32]byte{0b00000100})
+
+	params := Parameters{
+		Metrics: prometheus.NewRegistry(),
+		K:       1, Alpha: 1, BetaVirtuous: 1, BetaRogue: 2,
+	}
+	tree := Tree{}
+	tree.Initialize(params, c0000)
+	{
+		expected := "SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [0, 256)"
+		if pref := tree.Preference(); !c0000.Equals(pref) {
+			t.Fatalf("Wrong preference. Expected %s got %s", c0000, pref)
+		} else if tree.Finalized() {
+			t.Fatalf("Finalized too early")
+		} else if str := tree.String(); expected != str {
+			t.Fatalf("Wrong string. Expected:\n%s\ngot:\n%s", expected, str)
+		}
+	}
+
+	tree.Add(c1000)
+	{
+		expected := "SB(Preference = 0, NumSuccessfulPolls[0] = 0, NumSuccessfulPolls[1] = 0, SF = SF(Preference = 0, Confidence = 0, Finalized = false)) Bit = 0\n"+
+		"    SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [1, 256)\n"+
+		"    SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [1, 256)"
+		if pref := tree.Preference(); !c0000.Equals(pref) {
+			t.Fatalf("Wrong preference. Expected %s got %s", c0000, pref)
+		} else if tree.Finalized() {
+			t.Fatalf("Finalized too early")
+		} else if str := tree.String(); expected != str {
+			t.Fatalf("Wrong string. Expected:\n%s\ngot:\n%s", expected, str)
+		}
+	}
+
+	tree.Add(c0010)
+	{
+		expected := "SB(Preference = 0, NumSuccessfulPolls[0] = 0, NumSuccessfulPolls[1] = 0, SF = SF(Preference = 0, Confidence = 0, Finalized = false)) Bit = 0\n"+
+		"    SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [1, 2)\n"+
+		"        SB(Preference = 0, NumSuccessfulPolls[0] = 0, NumSuccessfulPolls[1] = 0, SF = SF(Preference = 0, Confidence = 0, Finalized = false)) Bit = 2\n"+
+		"            SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [3, 256)\n"+
+		"            SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [3, 256)\n"+
+		"    SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [1, 256)"
+		if pref := tree.Preference(); !c0000.Equals(pref) {
+			t.Fatalf("Wrong preference. Expected %s got %s", c0000, pref)
+		} else if tree.Finalized() {
+			t.Fatalf("Finalized too early")
+		} else if str := tree.String(); expected != str {
+			t.Fatalf("Wrong string. Expected:\n%s\ngot:\n%s", expected, str)
+		}
+	}
+
+	c0000Bag := ids.Bag{}
+	c0000Bag.Add(c0000)
+	tree.RecordPoll(c0000Bag)
+	{
+		expected := "SB(Preference = 0, NumSuccessfulPolls[0] = 1, NumSuccessfulPolls[1] = 0, SF = SF(Preference = 0, Confidence = 1, Finalized = false)) Bit = 0\n"+
+		"    SB(Preference = 0, NumSuccessfulPolls[0] = 1, NumSuccessfulPolls[1] = 0, SF = SF(Preference = 0, Confidence = 1, Finalized = false)) Bit = 2\n"+
+		"        SB(NumSuccessfulPolls = 1, Confidence = 1, Finalized = true) Bits = [3, 256)\n"+
+		"        SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [3, 256)\n"+
+		"    SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [1, 256)"
+		if pref := tree.Preference(); !c0000.Equals(pref) {
+			t.Fatalf("Wrong preference. Expected %s got %s", c0000, pref)
+		} else if tree.Finalized() {
+			t.Fatalf("Finalized too early")
+		} else if str := tree.String(); expected != str {
+			t.Fatalf("Wrong string. Expected:\n%s\ngot:\n%s", expected, str)
+		}
+	}
+
+	tree.Add(c0100)
+	{
+		expected := "SB(Preference = 0, NumSuccessfulPolls[0] = 1, NumSuccessfulPolls[1] = 0, SF = SF(Preference = 0, Confidence = 1, Finalized = false)) Bit = 0\n"+
+		"    SB(Preference = 0, NumSuccessfulPolls[0] = 1, NumSuccessfulPolls[1] = 0, SF = SF(Preference = 0, Confidence = 1, Finalized = false)) Bit = 2\n"+
+		"        SB(NumSuccessfulPolls = 1, Confidence = 1, Finalized = true) Bits = [3, 256)\n"+
+		"        SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [3, 256)\n"+
+		"    SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [1, 256)"
+		if pref := tree.Preference(); !c0000.Equals(pref) {
+			t.Fatalf("Wrong preference. Expected %s got %s", c0000, pref)
+		} else if tree.Finalized() {
+			t.Fatalf("Finalized too early")
+		} else if str := tree.String(); expected != str {
+			t.Fatalf("Wrong string. Expected:\n%s\ngot:\n%s", expected, str)
+		}
+	}
+
+	c0100Bag := ids.Bag{}
+	c0100Bag.Add(c0100)
+	tree.RecordPoll(c0100Bag)
+	{
+		expected := "SB(Preference = 0, NumSuccessfulPolls[0] = 1, NumSuccessfulPolls[1] = 0, SF = SF(Preference = 0, Confidence = 0, Finalized = false)) Bit = 2\n"+
+		"    SB(NumSuccessfulPolls = 1, Confidence = 1, Finalized = true) Bits = [3, 256)\n"+
+		"    SB(NumSuccessfulPolls = 0, Confidence = 0, Finalized = false) Bits = [3, 256)"
+		if pref := tree.Preference(); !c0000.Equals(pref) {
+			t.Fatalf("Wrong preference. Expected %s got %s", c0000, pref)
+		} else if tree.Finalized() {
+			t.Fatalf("Finalized too early")
+		} else if str := tree.String(); expected != str {
+			t.Fatalf("Wrong string. Expected:\n%s\ngot:\n%s", expected, str)
+		}
+	}
+}

utils/logging/log.go

@@ -171,7 +171,7 @@ func (l *Log) format(level Level, format string, args ...interface{}) string {
 
 	return fmt.Sprintf("%s[%s]%s %s\n",
 		level,
-		time.Now().Format("01-02|15:04:05.000"),
+		time.Now().Format("01-02|15:04:05"),
 		prefix,
 		text)
 }

vms/avm/vm.go

@@ -251,7 +251,10 @@ func (vm *VM) GetTx(txID ids.ID) (snowstorm.Tx, error) {
  ******************************************************************************
  */
 
-// IssueTx attempts to send a transaction to consensus
+// IssueTx attempts to send a transaction to consensus.
+// If onDecide is specified, the function will be called when the transaction is
+// either accepted or rejected with the appropriate status. This function will
+// go out of scope when the transaction is removed from memory.
 func (vm *VM) IssueTx(b []byte, onDecide func(choices.Status)) (ids.ID, error) {
 	tx, err := vm.parseTx(b)
 	if err != nil {

vms/avm/vm_test.go

@@ -548,6 +548,8 @@ func TestGenesisGetUTXOs(t *testing.T) {
 	}
 }
 
+// Test issuing a transaction that consumes a currently pending UTXO. The
+// transaction should be issued successfully.
 func TestIssueDependentTx(t *testing.T) {
 	genesisBytes := BuildGenesisTest(t)
 

vms/platformvm/service.go

@@ -770,6 +770,8 @@ func (service *Service) signAddNonDefaultSubnetValidatorTx(tx *addNonDefaultSubn
 		return nil, errors.New("no place for key to sign")
 	}
 
+	crypto.SortSECP2561RSigs(tx.ControlSigs)
+
 	return tx, nil
 }
 

vms/spchainvm/keychain.go

@@ -20,31 +20,42 @@ var (
 
 // Keychain is a collection of keys that can be used to spend utxos
 type Keychain struct {
+	factory   crypto.FactorySECP256K1R
 	networkID uint32
 	chainID   ids.ID
-	// This can be used to iterate over. However, it should not be modified externally.
+
+	// Key: The id of a private key (namely, [privKey].PublicKey().Address().Key())
+	// Value: The index in Keys of that private key
 	keyMap map[[20]byte]int
+
+	// Each element is an address controlled by a key in [Keys]
+	// This can be used to iterate over. It should not be modified externally.
 	Addrs ids.ShortSet
+
+	// List of keys this keychain manages
+	// This can be used to iterate over. It should not be modified externally.
 	Keys []*crypto.PrivateKeySECP256K1R
 }
 
 // NewKeychain creates a new keychain for a chain
 func NewKeychain(networkID uint32, chainID ids.ID) *Keychain {
 	return &Keychain{
+		networkID: networkID,
 		chainID:   chainID,
 		keyMap:    make(map[[20]byte]int),
 	}
 }
 
 // New returns a newly generated private key
-func (kc *Keychain) New() *crypto.PrivateKeySECP256K1R {
+func (kc *Keychain) New() (*crypto.PrivateKeySECP256K1R, error) {
-	factory := &crypto.FactorySECP256K1R{}
+	skGen, err := kc.factory.NewPrivateKey()
-
+	if err != nil {
-	skGen, _ := factory.NewPrivateKey()
+		return nil, err
+	}
 
 	sk := skGen.(*crypto.PrivateKeySECP256K1R)
 	kc.Add(sk)
-	return sk
+	return sk, nil
 }
 
 // Add a new key to the key chain

vms/spdagvm/keychain.go

@@ -20,29 +20,35 @@ var (
 
 // Keychain is a collection of keys that can be used to spend utxos
 type Keychain struct {
-	// This can be used to iterate over. However, it should not be modified externally.
+	factory   crypto.FactorySECP256K1R
+	networkID uint32
+	chainID   ids.ID
+
 	// Key: The id of a private key (namely, [privKey].PublicKey().Address().Key())
 	// Value: The index in Keys of that private key
 	keyMap map[[20]byte]int
 
 	// Each element is an address controlled by a key in [Keys]
+	// This can be used to iterate over. It should not be modified externally.
 	Addrs ids.ShortSet
 
 	// List of keys this keychain manages
+	// This can be used to iterate over. It should not be modified externally.
 	Keys []*crypto.PrivateKeySECP256K1R
 }
 
-func (kc *Keychain) init() {
+// NewKeychain creates a new keychain for a chain
-	if kc.keyMap == nil {
+func NewKeychain(networkID uint32, chainID ids.ID) *Keychain {
-		kc.keyMap = make(map[[20]byte]int)
+	return &Keychain{
+		networkID: networkID,
+		chainID:   chainID,
+		keyMap:    make(map[[20]byte]int),
 	}
 }
 
 // Add a new key to the key chain.
 // If [key] is already in the keychain, does nothing.
 func (kc *Keychain) Add(key *crypto.PrivateKeySECP256K1R) {
-	kc.init()
-
 	addr := key.PublicKey().Address() // The address controlled by [key]
 	addrHash := addr.Key()
 	if _, ok := kc.keyMap[addrHash]; !ok {
@@ -53,9 +59,7 @@ func (kc *Keychain) Add(key *crypto.PrivateKeySECP256K1R) {
 }
 
 // Get a key from the keychain. If the key is unknown, the second return value is false.
-func (kc Keychain) Get(id ids.ShortID) (*crypto.PrivateKeySECP256K1R, bool) {
+func (kc *Keychain) Get(id ids.ShortID) (*crypto.PrivateKeySECP256K1R, bool) {
-	kc.init()
-
 	if i, ok := kc.keyMap[id.Key()]; ok {
 		return kc.Keys[i], true
 	}
@@ -63,15 +67,13 @@ func (kc Keychain) Get(id ids.ShortID) (*crypto.PrivateKeySECP256K1R, bool) {
 }
 
 // Addresses returns a list of addresses this keychain manages
-func (kc Keychain) Addresses() ids.ShortSet { return kc.Addrs }
+func (kc *Keychain) Addresses() ids.ShortSet { return kc.Addrs }
 
 // New returns a newly generated private key.
 // The key and the address it controls are added to
 // [kc.Keys] and [kc.Addrs], respectively
 func (kc *Keychain) New() (*crypto.PrivateKeySECP256K1R, error) {
-	factory := crypto.FactorySECP256K1R{}
+	skGen, err := kc.factory.NewPrivateKey()
-
-	skGen, err := factory.NewPrivateKey()
 	if err != nil {
 		return nil, err
 	}
@@ -84,8 +86,8 @@ func (kc *Keychain) New() (*crypto.PrivateKeySECP256K1R, error) {
 // Spend attempts to create an input
 func (kc *Keychain) Spend(utxo *UTXO, time uint64) (Input, *InputSigner, error) {
 	builder := Builder{
-		NetworkID: 0,
+		NetworkID: kc.networkID,
-		ChainID:   ids.Empty,
+		ChainID:   kc.chainID,
 	}
 
 	switch out := utxo.Out().(type) {
@@ -148,8 +150,8 @@ func (kc *Keychain) GetSigsAndKeys(addresses []ids.ShortID, threshold int) ([]*S
 	sigs := []*Sig{}
 	keys := []*crypto.PrivateKeySECP256K1R{}
 	builder := Builder{
-		NetworkID: 0,
+		NetworkID: kc.networkID,
-		ChainID:   ids.Empty,
+		ChainID:   kc.chainID,
 	}
 	for i := uint32(0); i < uint32(len(addresses)) && len(keys) < threshold; i++ {
 		if key, exists := kc.Get(addresses[i]); exists {

vms/spdagvm/vm.go

@@ -315,7 +315,7 @@ func (vm *VM) Send(amount uint64, assetID, toAddrStr string, fromPKs []string) (
 	}
 
 	// Add all of the keys in [fromPKs] to a keychain
-	keychain := Keychain{}
+	keychain := NewKeychain(vm.ctx.NetworkID, vm.ctx.ChainID)
 	factory := crypto.FactorySECP256K1R{}
 	cb58 := formatting.CB58{}
 	for _, fpk := range fromPKs {
@@ -359,7 +359,7 @@ func (vm *VM) Send(amount uint64, assetID, toAddrStr string, fromPKs []string) (
 		ChainID:   vm.ctx.ChainID,
 	}
 	currentTime := vm.clock.Unix()
-	tx, err := builder.NewTxFromUTXOs(&keychain, utxos, amount, vm.TxFee, 0, 1, toAddrs, outAddr, currentTime)
+	tx, err := builder.NewTxFromUTXOs(keychain, utxos, amount, vm.TxFee, 0, 1, toAddrs, outAddr, currentTime)
 	if err != nil {
 		return "", err
 	}

xputtest/README.md

@@ -0,0 +1,17 @@
+# Throughput testing
+
+A throughput test is run in two parts. First a network must be running with at least one of the nodes running a throughput server. To start a throughput server when running a node the `--xput-server-enabled=true` flag should be passed.
+
+An example single node network can be started with:
+
+```sh
+./build/ava --public-ip=127.0.0.1 --xput-server-port=9652 --xput-server-enabled=true --db-enabled=false --staking-tls-enabled=false --snow-sample-size=1 --snow-quorum-size=1
+```
+
+The thoughput node can be started with:
+
+```sh
+./build/xputtest --ip=127.0.0.1 --port=9652 --sp-chain
+```
+
+The above example with run a throughput test on the simple payment chain. Tests can be run with `--sp-dag` to run throughput tests on the simple payment dag. Tests can be run with `--avm` to run throughput tests on the AVA virtual machine.

xputtest/avm.go

@@ -19,12 +19,10 @@ import (
 	"github.com/ava-labs/gecko/xputtest/avmwallet"
 )
 
-func (n *network) benchmarkAVM(genesisState *platformvm.Genesis) {
+// benchmark an instance of the avm
-	avmChain := genesisState.Chains[0]
+func (n *network) benchmarkAVM(chain *platformvm.CreateChainTx) {
-	n.log.AssertTrue(avmChain.ChainName == "AVM", "wrong chain name")
+	genesisBytes := chain.GenesisData
-	genesisBytes := avmChain.GenesisData
+	wallet, err := avmwallet.NewWallet(n.log, n.networkID, chain.ID(), config.AvaTxFee)
-
-	wallet, err := avmwallet.NewWallet(n.networkID, avmChain.ID(), config.AvaTxFee)
 	n.log.AssertNoError(err)
 
 	cb58 := formatting.CB58{}
@@ -56,9 +54,10 @@ func (n *network) benchmarkAVM(genesisState *platformvm.Genesis) {
 
 	n.log.AssertNoError(wallet.GenerateTxs(config.NumTxs, assetID))
 
-	go n.log.RecoverAndPanic(func() { n.IssueAVM(avmChain.ID(), assetID, wallet) })
+	go n.log.RecoverAndPanic(func() { n.IssueAVM(chain.ID(), assetID, wallet) })
 }
 
+// issue transactions to the instance of the avm funded by the provided wallet
 func (n *network) IssueAVM(chainID ids.ID, assetID ids.ID, wallet *avmwallet.Wallet) {
 	n.log.Debug("Issuing with %d", wallet.Balance(assetID))
 	numAccepted := 0
@@ -66,11 +65,15 @@ func (n *network) IssueAVM(chainID ids.ID, assetID ids.ID, wallet *avmwallet.Wal
 
 	n.decided <- ids.ID{}
 
+	// track the last second of transactions
 	meter := timer.TimedMeter{Duration: time.Second}
 	for d := range n.decided {
+		// display the TPS every 1000 txs
 		if numAccepted%1000 == 0 {
 			n.log.Info("TPS: %d", meter.Ticks())
 		}
+
+		// d is the ID of the tx that was accepted
 		if !d.IsZero() {
 			meter.Tick()
 			n.log.Debug("Finalized %s", d)
@@ -78,10 +81,12 @@ func (n *network) IssueAVM(chainID ids.ID, assetID ids.ID, wallet *avmwallet.Wal
 			numPending--
 		}
 
+		// Issue all the txs that we can right now
 		for numPending < config.MaxOutstandingTxs && wallet.Balance(assetID) > 0 && numAccepted+numPending < config.NumTxs {
 			tx := wallet.NextTx()
 			n.log.AssertTrue(tx != nil, "Tx creation failed")
 
+			// send the IssueTx message
 			it, err := n.build.IssueTx(chainID, tx.Bytes())
 			n.log.AssertNoError(err)
 			ds := it.DataStream()
@@ -97,8 +102,11 @@ func (n *network) IssueAVM(chainID ids.ID, assetID ids.ID, wallet *avmwallet.Wal
 			numPending++
 			n.log.Debug("Sent tx, pending = %d, accepted = %d", numPending, numAccepted)
 		}
+
+		// If we are done issuing txs, return from the function
 		if numAccepted+numPending >= config.NumTxs {
 			n.log.Info("done with test")
+			net.ec.Stop()
 			return
 		}
 	}

xputtest/avmwallet/utxo_set.go

@@ -13,8 +13,12 @@ import (
 
 // UTXOSet ...
 type UTXOSet struct {
-	// This can be used to iterate over. However, it should not be modified externally.
+	// Key: The id of a UTXO
+	// Value: The index in UTXOs of that UTXO
 	utxoMap map[[32]byte]int
+
+	// List of UTXOs in this set
+	// This can be used to iterate over. It should not be modified externally.
 	UTXOs []*avm.UTXO
 }
 

xputtest/avmwallet/wallet.go

@@ -13,6 +13,7 @@ import (
 	"github.com/ava-labs/gecko/snow"
 	"github.com/ava-labs/gecko/utils/crypto"
 	"github.com/ava-labs/gecko/utils/hashing"
+	"github.com/ava-labs/gecko/utils/logging"
 	"github.com/ava-labs/gecko/utils/math"
 	"github.com/ava-labs/gecko/utils/timer"
 	"github.com/ava-labs/gecko/utils/wrappers"
@@ -25,10 +26,14 @@ import (
 type Wallet struct {
 	networkID uint32
 	chainID   ids.ID
+
 	clock timer.Clock
 	codec codec.Codec
+	log   logging.Logger
+
 	keychain *secp256k1fx.Keychain // Mapping from public address to the SigningKeys
 	utxoSet  *UTXOSet              // Mapping from utxoIDs to UTXOs
+
 	balance map[[32]byte]uint64
 	txFee   uint64
 
@@ -37,7 +42,7 @@ type Wallet struct {
 }
 
 // NewWallet returns a new Wallet
-func NewWallet(networkID uint32, chainID ids.ID, txFee uint64) (*Wallet, error) {
+func NewWallet(log logging.Logger, networkID uint32, chainID ids.ID, txFee uint64) (*Wallet, error) {
 	c := codec.NewDefault()
 	errs := wrappers.Errs{}
 	errs.Add(
@@ -54,6 +59,7 @@ func NewWallet(networkID uint32, chainID ids.ID, txFee uint64) (*Wallet, error)
 		networkID: networkID,
 		chainID:   chainID,
 		codec:     c,
+		log:       log,
 		keychain:  secp256k1fx.NewKeychain(),
 		utxoSet:   &UTXOSet{},
 		balance:   make(map[[32]byte]uint64),
@@ -249,10 +255,17 @@ func (w *Wallet) CreateTx(assetID ids.ID, amount uint64, destAddr ids.ShortID) (
 // Generate them all on test initialization so tx generation is not bottleneck
 // in testing
 func (w *Wallet) GenerateTxs(numTxs int, assetID ids.ID) 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([]*avm.Tx, numTxs)
 	for i := 0; i < numTxs; i++ {
 		addr, err := w.CreateAddress()
@@ -271,8 +284,14 @@ func (w *Wallet) GenerateTxs(numTxs int, assetID ids.ID) error {
 			w.AddUTXO(utxo)
 		}
 
+		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
 }
 

xputtest/avmwallet/wallet_test.go

@@ -11,15 +11,15 @@ import (
 	"github.com/ava-labs/gecko/snow"
 	"github.com/ava-labs/gecko/utils/crypto"
 	"github.com/ava-labs/gecko/utils/formatting"
+	"github.com/ava-labs/gecko/utils/logging"
 	"github.com/ava-labs/gecko/utils/units"
 	"github.com/ava-labs/gecko/vms/avm"
-	"github.com/ava-labs/gecko/vms/platformvm"
 	"github.com/ava-labs/gecko/vms/secp256k1fx"
 )
 
 func TestNewWallet(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -30,7 +30,7 @@ func TestNewWallet(t *testing.T) {
 
 func TestWalletGetAddress(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -46,7 +46,7 @@ func TestWalletGetAddress(t *testing.T) {
 
 func TestWalletGetMultipleAddresses(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -66,7 +66,7 @@ func TestWalletGetMultipleAddresses(t *testing.T) {
 
 func TestWalletEmptyBalance(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -78,7 +78,7 @@ func TestWalletEmptyBalance(t *testing.T) {
 
 func TestWalletAddUTXO(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -100,7 +100,7 @@ func TestWalletAddUTXO(t *testing.T) {
 
 func TestWalletAddInvalidUTXO(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -119,7 +119,7 @@ func TestWalletAddInvalidUTXO(t *testing.T) {
 
 func TestWalletCreateTx(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -169,7 +169,7 @@ func TestWalletCreateTx(t *testing.T) {
 
 func TestWalletImportKey(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -194,7 +194,7 @@ func TestWalletImportKey(t *testing.T) {
 
 func TestWalletString(t *testing.T) {
 	chainID := ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
-	w, err := NewWallet(12345, chainID, 0)
+	w, err := NewWallet(logging.NoLog{}, 12345, chainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -226,7 +226,7 @@ func TestWalletWithGenesis(t *testing.T) {
 	ctx.NetworkID = 12345
 	ctx.ChainID = ids.NewID([32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
 
-	w, err := NewWallet(ctx.NetworkID, ctx.ChainID, 0)
+	w, err := NewWallet(logging.NoLog{}, ctx.NetworkID, ctx.ChainID, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -245,20 +245,7 @@ func TestWalletWithGenesis(t *testing.T) {
 		w.ImportKey(sk.(*crypto.PrivateKeySECP256K1R))
 	}
 
-	platformGenesisBytes := genesis.Genesis(genesis.LocalID)
+	avmChain := genesis.VMGenesis(ctx.NetworkID, avm.ID)
-	genesisState := &platformvm.Genesis{}
-	err = platformvm.Codec.Unmarshal(platformGenesisBytes, genesisState)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if err := genesisState.Initialize(); err != nil {
-		t.Fatal(err)
-	}
-
-	avmChain := genesisState.Chains[0]
-	if name := avmChain.ChainName; name != "AVM" {
-		t.Fatalf("wrong chain name")
-	}
 	genesisBytes := avmChain.GenesisData
 
 	genesis := avm.Genesis{}
@@ -282,8 +269,8 @@ func TestWalletWithGenesis(t *testing.T) {
 
 	assetID := genesisTx.ID()
 
-	if balance := w.Balance(assetID); balance != 45*units.MegaAva {
+	if balance := w.Balance(assetID); balance == 0 {
-		t.Fatalf("balance of %d was expected but got %d", 45*units.MegaAva, balance)
+		t.Fatalf("expected a positive balance")
 	}
 
 	for i := 1; i <= 1000; i++ {

xputtest/chainwallet/wallet.go

@@ -10,6 +10,7 @@ import (
 	"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"
 )
 
@@ -17,6 +18,9 @@ import (
 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
@@ -25,17 +29,24 @@ type Wallet struct {
 }
 
 // NewWallet ...
-func NewWallet(networkID uint32, chainID ids.ID) Wallet {
+func NewWallet(log logging.Logger, networkID uint32, chainID ids.ID) *Wallet {
-	return 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 { return w.keychain.New().PublicKey().Address() }
+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) }
@@ -56,65 +67,36 @@ func (w *Wallet) Balance() uint64 { return w.balance }
 // 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 := 0; i < numTxs; {
+	for i := range w.txs {
-		for _, account := range w.accountSet {
+		tx, err := w.MakeTx()
-			if i >= numTxs {
-				break
-			}
-
-			accountID := account.ID()
-			key, exists := w.keychain.Get(accountID)
-			if !exists {
-				return errors.New("missing account")
-			}
-
-			amount := uint64(1)
-			tx, sendAccount, err := account.CreateTx(amount, accountID, ctx, key)
 		if err != nil {
 			return err
 		}
 
-			newAccount, err := sendAccount.Receive(tx, ctx)
+		if numGenerated := i + 1; numGenerated%frequency == 0 {
-			if err != nil {
+			w.log.Info("Generated %d out of %d transactions", numGenerated, numTxs)
-				return err
 		}
-			w.accountSet[accountID.Key()] = newAccount
+
 		w.txs[i] = tx
-			i++
-		}
 	}
+
+	w.log.Info("Finished generating %d transactions", numTxs)
+
 	return nil
 }
 
-/*
-// Send a new transaction
-func (w *Wallet) Send() *spchainvm.Tx {
-	ctx := snow.DefaultContextTest()
-	ctx.NetworkID = w.networkID
-	ctx.ChainID = w.chainID
-
-	for _, account := range w.accountSet {
-		accountID := account.ID()
-		if key, exists := w.keychain.Get(accountID); exists {
-			amount := uint64(1)
-			if tx, sendAccount, err := account.CreateTx(amount, accountID, ctx, key); err == nil {
-				newAccount, err := sendAccount.Receive(tx, ctx)
-				if err == nil {
-					w.accountSet[accountID.Key()] = newAccount
-					return tx
-				}
-			}
-		}
-	}
-	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 {
@@ -125,6 +107,35 @@ func (w *Wallet) NextTx() *spchainvm.Tx {
 	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"+

xputtest/dagwallet/utxo_set.go

@@ -13,8 +13,12 @@ import (
 
 // UTXOSet ...
 type UTXOSet struct {
-	// This can be used to iterate over. However, it should not be modified externally.
+	// Key: The id of a UTXO
+	// Value: The index in UTXOs of that UTXO
 	utxoMap map[[32]byte]int
+
+	// List of UTXOs in this set
+	// This can be used to iterate over. It should not be modified externally.
 	UTXOs []*spdagvm.UTXO
 }
 

xputtest/dagwallet/wallet.go

@@ -25,11 +25,11 @@ type Wallet struct {
 }
 
 // NewWallet returns a new Wallet
-func NewWallet(networkID uint32, chainID ids.ID, txFee uint64) Wallet {
+func NewWallet(networkID uint32, chainID ids.ID, txFee uint64) *Wallet {
-	return Wallet{
+	return &Wallet{
 		networkID: networkID,
 		chainID:   chainID,
-		keychain:  &spdagvm.Keychain{},
+		keychain:  spdagvm.NewKeychain(networkID, chainID),
 		utxoSet:   &UTXOSet{},
 		txFee:     txFee,
 	}

xputtest/main.go

@@ -16,7 +16,9 @@ import (
 	"github.com/ava-labs/gecko/ids"
 	"github.com/ava-labs/gecko/utils/crypto"
 	"github.com/ava-labs/gecko/utils/logging"
-	"github.com/ava-labs/gecko/vms/platformvm"
+	"github.com/ava-labs/gecko/vms/avm"
+	"github.com/ava-labs/gecko/vms/spchainvm"
+	"github.com/ava-labs/gecko/vms/spdagvm"
 )
 
 func main() {
@@ -24,6 +26,7 @@ func main() {
 		fmt.Printf("Failed to parse arguments: %s\n", err)
 	}
 
+	// set up logging
 	config.LoggingConfig.Directory = path.Join(config.LoggingConfig.Directory, "client")
 	log, err := logging.New(config.LoggingConfig)
 	if err != nil {
@@ -33,6 +36,7 @@ func main() {
 
 	defer log.Stop()
 
+	// initialize state based on CLI args
 	net.log = log
 	crypto.EnableCrypto = config.EnableCrypto
 	net.decided = make(chan ids.ID, config.MaxOutstandingTxs)
@@ -42,11 +46,13 @@ func main() {
 		return
 	}
 
+	// Init the network
 	log.AssertNoError(net.Initialize())
 
 	net.net.Start()
 	defer net.net.Stop()
 
+	// connect to the node
 	serr := salticidae.NewError()
 	remoteIP := salticidae.NewNetAddrFromIPPortString(config.RemoteIP.String(), true, &serr)
 	if code := serr.GetCode(); code != 0 {
@@ -60,6 +66,7 @@ func main() {
 		return
 	}
 
+	// start a cpu profile
 	file, gErr := os.Create("cpu_client.profile")
 	log.AssertNoError(gErr)
 	gErr = pprof.StartCPUProfile(file)
@@ -71,22 +78,19 @@ func main() {
 
 	net.networkID = config.NetworkID
 
-	platformGenesisBytes := genesis.Genesis(net.networkID)
+	// start the benchmark we want to run
-	genesisState := &platformvm.Genesis{}
-	log.AssertNoError(platformvm.Codec.Unmarshal(platformGenesisBytes, genesisState))
-	log.AssertNoError(genesisState.Initialize())
-
 	switch config.Chain {
 	case spChain:
-		net.benchmarkSPChain(genesisState)
+		net.benchmarkSPChain(genesis.VMGenesis(config.NetworkID, spchainvm.ID))
 	case spDAG:
-		net.benchmarkSPDAG(genesisState)
+		net.benchmarkSPDAG(genesis.VMGenesis(config.NetworkID, spdagvm.ID))
 	case avmDAG:
-		net.benchmarkAVM(genesisState)
+		net.benchmarkAVM(genesis.VMGenesis(config.NetworkID, avm.ID))
 	default:
 		log.Fatal("did not specify whether to test dag or chain. Exiting")
 		return
 	}
 
+	// start processing network messages
 	net.ec.Dispatch()
 }

xputtest/spchain.go

@@ -19,12 +19,10 @@ import (
 	"github.com/ava-labs/gecko/xputtest/chainwallet"
 )
 
-func (n *network) benchmarkSPChain(genesisState *platformvm.Genesis) {
+// benchmark an instance of the sp chain
-	spchainChain := genesisState.Chains[3]
+func (n *network) benchmarkSPChain(chain *platformvm.CreateChainTx) {
-	n.log.AssertTrue(spchainChain.ChainName == "Simple Chain Payments", "wrong chain name")
+	genesisBytes := chain.GenesisData
-	genesisBytes := spchainChain.GenesisData
+	wallet := chainwallet.NewWallet(n.log, n.networkID, chain.ID())
-
-	wallet := chainwallet.NewWallet(n.networkID, spchainChain.ID())
 
 	codec := spchainvm.Codec{}
 	accounts, err := codec.UnmarshalGenesis(genesisBytes)
@@ -47,10 +45,10 @@ func (n *network) benchmarkSPChain(genesisState *platformvm.Genesis) {
 
 	n.log.AssertNoError(wallet.GenerateTxs(config.NumTxs))
 
-	go n.log.RecoverAndPanic(func() { n.IssueSPChain(spchainChain.ID(), wallet) })
+	go n.log.RecoverAndPanic(func() { n.IssueSPChain(chain.ID(), wallet) })
 }
 
-func (n *network) IssueSPChain(chainID ids.ID, wallet chainwallet.Wallet) {
+func (n *network) IssueSPChain(chainID ids.ID, wallet *chainwallet.Wallet) {
 	n.log.Debug("Issuing with %d", wallet.Balance())
 	numAccepted := 0
 	numPending := 0
@@ -90,6 +88,7 @@ func (n *network) IssueSPChain(chainID ids.ID, wallet chainwallet.Wallet) {
 		}
 		if numAccepted+numPending >= config.NumTxs {
 			n.log.Info("done with test")
+			net.ec.Stop()
 			return
 		}
 	}

xputtest/spdag.go

@@ -19,12 +19,10 @@ import (
 	"github.com/ava-labs/gecko/xputtest/dagwallet"
 )
 
-func (n *network) benchmarkSPDAG(genesisState *platformvm.Genesis) {
+// benchmark an instance of the sp dag
-	spDAGChain := genesisState.Chains[2]
+func (n *network) benchmarkSPDAG(chain *platformvm.CreateChainTx) {
-	n.log.AssertTrue(spDAGChain.ChainName == "Simple DAG Payments", "wrong chain name")
+	genesisBytes := chain.GenesisData
-	genesisBytes := spDAGChain.GenesisData
+	wallet := dagwallet.NewWallet(n.networkID, chain.ID(), config.AvaTxFee)
-
-	wallet := dagwallet.NewWallet(n.networkID, spDAGChain.ID(), config.AvaTxFee)
 
 	codec := spdagvm.Codec{}
 	tx, err := codec.UnmarshalTx(genesisBytes)
@@ -43,10 +41,11 @@ func (n *network) benchmarkSPDAG(genesisState *platformvm.Genesis) {
 		wallet.AddUTXO(utxo)
 	}
 
-	go n.log.RecoverAndPanic(func() { n.IssueSPDAG(spDAGChain.ID(), wallet) })
+	go n.log.RecoverAndPanic(func() { n.IssueSPDAG(chain.ID(), wallet) })
 }
 
-func (n *network) IssueSPDAG(chainID ids.ID, wallet dagwallet.Wallet) {
+// issue transactions to the instance of the spdag funded by the provided wallet
+func (n *network) IssueSPDAG(chainID ids.ID, wallet *dagwallet.Wallet) {
 	n.log.Info("starting avalanche benchmark")
 	pending := make(map[[32]byte]*spdagvm.Tx)
 	canAdd := []*spdagvm.Tx{}