Repository: coinbase/rosetta-cli
Pull request #105 (2020-08-12 19:50:41)

Support for importing prefunded accounts
### Motivation
Supports importing accounts into the database by adding them as `PrefundedAccount` in the config file. 

### Solution
<!--
What is the solution here from a high level. What are the key technical decisions and why were they made?
-->

### Open questions
<!--
(optional) Any open questions or feedback on design desired?
-->



Patch:
diff --git a/configuration/configuration.go b/configuration/configuration.go
index 54e6a2d..cb0ce14 100644
--- a/configuration/configuration.go
+++ b/configuration/configuration.go
@@ -15,6 +15,7 @@
 package configuration
 
 import (
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"log"
@@ -220,6 +221,16 @@ type ConstructionConfiguration struct {
 	// MaxAddresses is the maximum number of addresses
 	// to generate while testing.
 	MaxAddresses int `json:"max_addresses"`
+
+	// PrefundedAccounts is an array of prefunded accounts
+	// to use while testing.
+	PrefundedAccounts []*PrefundedAccount `json:"prefunded_accounts"`
+}
+
+type PrefundedAccount struct {
+	PrivateKeyHex string          `json:"privkey"`
+	Address       string          `json:"address"`
+	CurveType     types.CurveType `json:"curve_type"`
 }
 
 // DefaultConstructionConfiguration returns the *ConstructionConfiguration
@@ -581,6 +592,25 @@ func assertConstructionConfiguration(config *ConstructionConfiguration) error {
 		return fmt.Errorf("%w: invalid value for MaximumFee", err)
 	}
 
+	for _, account := range config.PrefundedAccounts {
+		// Checks that privkey is hex encoded
+		_, err := hex.DecodeString(account.PrivateKeyHex)
+		if err != nil {
+			return fmt.Errorf("%s is not hex encoded", account.PrivateKeyHex)
+		}
+
+		// Checks if valid curvetype
+		err = asserter.CurveType(account.CurveType)
+		if err != nil {
+			return fmt.Errorf("invalid CurveType %s", err)
+		}
+
+		// Checks if address is not empty string
+		if account.Address == "" {
+			return fmt.Errorf("Account.Address is missing")
+		}
+	}
+
 	return nil
 }
 
diff --git a/go.sum b/go.sum
index 2d10d3e..2e9c46c 100644
--- a/go.sum
+++ b/go.sum
@@ -194,6 +194,7 @@ github.com/mattn/go-isatty v0.0.10/go.mod h1:qgIWMr58cqv1PHHyhnkY9lrL7etaEgOFcME
 github.com/mattn/go-isatty v0.0.11 h1:FxPOTFNqGkuDUGi3H/qkUbQO4ZiBa2brKq5r0l8TGeM=
 github.com/mattn/go-isatty v0.0.11/go.mod h1:PhnuNfih5lzO57/f3n+odYbM4JtupLOxQOAqxQCu2WE=
 github.com/mattn/go-runewidth v0.0.3/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
+github.com/mattn/go-runewidth v0.0.4 h1:2BvfKmzob6Bmd4YsL0zygOqfdFnK7GR4QL06Do4/p7Y=
 github.com/mattn/go-runewidth v0.0.4/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4=
diff --git a/mocks/constructor/handler.go b/mocks/constructor/handler.go
index 8cef173..71c2e46 100644
--- a/mocks/constructor/handler.go
+++ b/mocks/constructor/handler.go
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Code generated by mockery v1.0.0. DO NOT EDIT.
+// Code generated by mockery v2.0.0-alpha.2. DO NOT EDIT.
 
 package constructor
 
diff --git a/mocks/constructor/helper.go b/mocks/constructor/helper.go
index d29bd13..acc8139 100644
--- a/mocks/constructor/helper.go
+++ b/mocks/constructor/helper.go
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Code generated by mockery v1.0.0. DO NOT EDIT.
+// Code generated by mockery v2.0.0-alpha.2. DO NOT EDIT.
 
 package constructor
 
diff --git a/pkg/storage/key_storage.go b/pkg/storage/key_storage.go
index 1b33b8f..5bcaebb 100644
--- a/pkg/storage/key_storage.go
+++ b/pkg/storage/key_storage.go
@@ -22,6 +22,7 @@ import (
 	"math/rand"
 	"time"
 
+	"github.com/coinbase/rosetta-cli/configuration"
 	"github.com/coinbase/rosetta-sdk-go/keys"
 	"github.com/coinbase/rosetta-sdk-go/types"
 )
@@ -37,6 +38,10 @@ const (
 	keyNamespace = "key"
 )
 
+var (
+	ErrAddrExists = errors.New("Address already exists")
+)
+
 func getAddressKey(address string) []byte {
 	return []byte(
 		fmt.Sprintf("%s/%s", keyNamespace, address),
@@ -58,13 +63,13 @@ func NewKeyStorage(
 	}
 }
 
-type key struct {
+type Key struct {
 	Address string        `json:"address"`
 	KeyPair *keys.KeyPair `json:"keypair"`
 }
 
-func parseKey(buf []byte) (*key, error) {
-	var k key
+func parseKey(buf []byte) (*Key, error) {
+	var k Key
 	err := getDecoder(bytes.NewReader(buf)).Decode(&k)
 	if err != nil {
 		return nil, err
@@ -73,7 +78,7 @@ func parseKey(buf []byte) (*key, error) {
 	return &k, nil
 }
 
-func serializeKey(k *key) ([]byte, error) {
+func serializeKey(k *Key) ([]byte, error) {
 	buf := new(bytes.Buffer)
 	err := getEncoder(buf).Encode(k)
 	if err != nil {
@@ -95,10 +100,10 @@ func (k *KeyStorage) Store(ctx context.Context, address string, keyPair *keys.Ke
 	}
 
 	if exists {
-		return fmt.Errorf("address %s already exists", address)
+		return fmt.Errorf("%w: address %s already exists", ErrAddrExists, address)
 	}
 
-	val, err := serializeKey(&key{
+	val, err := serializeKey(&Key{
 		Address: address,
 		KeyPair: keyPair,
 	})
@@ -205,3 +210,23 @@ func (k *KeyStorage) RandomAddress(ctx context.Context) (string, error) {
 
 	return addresses[rand.Intn(len(addresses))], nil
 }
+
+func (k *KeyStorage) ImportAccounts(ctx context.Context, accounts []*configuration.PrefundedAccount) error {
+	// Import prefunded account and save to database
+	for _, acc := range accounts {
+		keyPair, err := keys.ImportPrivKey(acc.PrivateKeyHex, acc.CurveType)
+		if err != nil {
+			return fmt.Errorf("%w: unable to import prefunded account", err)
+		}
+
+		// Skip if key already exists
+		err = k.Store(ctx, acc.Address, keyPair)
+		if errors.Is(err, ErrAddrExists) {
+			continue
+		}
+		if err != nil {
+			return fmt.Errorf("%w: unable to store prefunded account", err)
+		}
+	}
+	return nil
+}


Test patch:
diff --git a/configuration/configuration_test.go b/configuration/configuration_test.go
index 9fceb2a..138176f 100644
--- a/configuration/configuration_test.go
+++ b/configuration/configuration_test.go
@@ -105,6 +105,16 @@ var (
 			MaximumFee: "hello",
 		},
 	}
+
+	invalidPrefundedAccounts = &Configuration{
+		Construction: &ConstructionConfiguration{
+			PrefundedAccounts: []*PrefundedAccount{
+				&PrefundedAccount{
+					PrivateKeyHex: "hello",
+				},
+			},
+		},
+	}
 	invalidStartIndex = &Configuration{
 		Data: &DataConfiguration{
 			StartIndex: &badStartIndex,
@@ -180,6 +190,10 @@ func TestLoadConfiguration(t *testing.T) {
 			provided: invalidMaximumFee,
 			err:      true,
 		},
+		"invalid prefunded accounts": {
+			provided: invalidPrefundedAccounts,
+			err:      true,
+		},
 		"invalid start index": {
 			provided: invalidStartIndex,
 			err:      true,
diff --git a/pkg/storage/key_storage_test.go b/pkg/storage/key_storage_test.go
index f66c3df..4e09f35 100644
--- a/pkg/storage/key_storage_test.go
+++ b/pkg/storage/key_storage_test.go
@@ -20,6 +20,7 @@ import (
 
 	"github.com/coinbase/rosetta-cli/pkg/utils"
 
+	"github.com/coinbase/rosetta-cli/configuration"
 	"github.com/coinbase/rosetta-sdk-go/keys"
 	"github.com/coinbase/rosetta-sdk-go/types"
 	"github.com/ethereum/go-ethereum/common"
@@ -142,4 +143,58 @@ func TestKeyStorage(t *testing.T) {
 		assert.Error(t, err)
 		assert.Nil(t, sigs)
 	})
+
+	t.Run("imports accounts", func(t *testing.T) {
+		addresses, err := k.GetAllAddresses(ctx)
+		assert.NoError(t, err)
+		startingLen := len(addresses)
+
+		prefundedAccs := []*configuration.PrefundedAccount{
+			&configuration.PrefundedAccount{
+				PrivateKeyHex: "0e842a16b2d39a4dff5c63688513cb2109e30c3c30bc4eb502cc54f4614493f6",
+				Address:       "add1",
+				CurveType:     types.Edwards25519,
+			},
+			&configuration.PrefundedAccount{
+				PrivateKeyHex: "42efc44bdf7b2d4d45ddd6ddb727ed498c91e7070914c9ed0d80af680ff42b3e",
+				Address:       "add2",
+				CurveType:     types.Edwards25519,
+			},
+			&configuration.PrefundedAccount{
+				PrivateKeyHex: "01ea48249742650907004331e85536f868e2d3959434ba751d8aa230138a9707",
+				Address:       "add3",
+				CurveType:     types.Edwards25519,
+			},
+		}
+
+		err = k.ImportAccounts(ctx, prefundedAccs)
+		assert.NoError(t, err)
+		addresses, _ = k.GetAllAddresses(ctx)
+		endLen := len(addresses)
+		assert.Equal(t, endLen, startingLen+len(prefundedAccs))
+	})
+
+	t.Run("does not import same key twice", func(t *testing.T) {
+		prefundedAccs := []*configuration.PrefundedAccount{
+			&configuration.PrefundedAccount{
+				PrivateKeyHex: "17d08f5fe8c77af811caa0c9a187e668ce3b74a99acc3f6d976f075fa8e0be55",
+				Address:       "badadd",
+				CurveType:     types.Edwards25519,
+			},
+		}
+
+		// No error when importing key for the first time
+		err = k.ImportAccounts(ctx, prefundedAccs)
+		assert.NoError(t, err)
+		addresses, _ := k.GetAllAddresses(ctx)
+		startingLen := len(addresses)
+
+		// No error when trying to import the key again
+		err = k.ImportAccounts(ctx, prefundedAccs)
+		assert.NoError(t, err)
+		addresses, _ = k.GetAllAddresses(ctx)
+		endLen := len(addresses)
+
+		assert.Equal(t, endLen, startingLen)
+	})
 }
diff --git a/pkg/tester/construction.go b/pkg/tester/construction.go
index 5c4e07e..4ae3d75 100644
--- a/pkg/tester/construction.go
+++ b/pkg/tester/construction.go
@@ -136,6 +136,13 @@ func InitializeConstruction(
 		fetcher.WithTimeout(time.Duration(config.HTTPTimeout)*time.Second),
 	)
 
+	// Import prefunded account and save to database
+	// TODO: load balances of these imported addresses
+	err = keyStorage.ImportAccounts(ctx, config.Construction.PrefundedAccounts)
+	if err != nil {
+		return nil, err
+	}
+
 	// Load all accounts for network
 	addresses, err := keyStorage.GetAllAddresses(ctx)
 	if err != nil {