add wallet private Key

ava-labs · Aug 9, 2024 · 9fd32a3 · 9fd32a3
1 parent c35a474
commit 9fd32a3
Show file tree

Hide file tree

Showing 9 changed files with 64 additions and 22 deletions.
diff --git a/client/client.go b/client/client.go
@@ -474,6 +474,7 @@ type Op struct {
 	reassignPortsIfUsed bool
 	dynamicPorts        bool
 	networkID           uint32
+	walletPrivateKey    string
 }
 
 type OpOption func(*Op)
@@ -591,6 +592,12 @@ func WithDynamicPorts(dynamicPorts bool) OpOption {
 	}
 }
 
+func WithWalletPrivateKey(walletPrivateKey string) OpOption {
+	return func(op *Op) {
+		op.walletPrivateKey = walletPrivateKey
+	}
+}
+
 func isClientCanceled(ctxErr error, err error) bool {
 	if ctxErr != nil {
 		return true

diff --git a/cmd/control/control.go b/cmd/control/control.go
@@ -302,6 +302,15 @@ func startFunc(*cobra.Command, []string) error {
 		client.WithNetworkID(networkID),
 	}
 
+	if walletPrivateKey != "" {
+		ux.Print(log, logging.Yellow.Wrap("funding wallet private key provided: %s"), walletPrivateKey)
+		// validate if it's a valid private key
+		if _, err := os.Stat(walletPrivateKey); err != nil {
+			return fmt.Errorf("private key doesn't exist: %w", err)
+		}
+		opts = append(opts, client.WithWalletPrivateKey(walletPrivateKey))
+	}
+
 	if globalNodeConfig != "" {
 		ux.Print(log, logging.Yellow.Wrap("global node config provided, will be applied to all nodes: %s"), globalNodeConfig)
 		// validate it's valid JSON

diff --git a/local/blockchain.go b/local/blockchain.go
@@ -32,6 +32,7 @@ import (
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/avalanchego/utils/constants"
 	"github.com/ava-labs/avalanchego/utils/crypto/bls"
+	"github.com/ava-labs/avalanchego/utils/crypto/secp256k1"
 	"github.com/ava-labs/avalanchego/utils/logging"
 	"github.com/ava-labs/avalanchego/utils/set"
 	"github.com/ava-labs/avalanchego/vms/platformvm"
@@ -282,7 +283,7 @@ func (ln *localNetwork) installCustomChains(
 		}
 	}
 
-	w, err := newWallet(ctx, clientURI, preloadTXs)
+	w, err := newWallet(ctx, clientURI, preloadTXs, ln.walletPrivateKey)
 	if err != nil {
 		return nil, err
 	}
@@ -478,7 +479,7 @@ func (ln *localNetwork) addSubnetValidators(
 		}
 		subnetIDs[i] = subnetID
 	}
-	w, err := newWallet(ctx, clientURI, subnetIDs)
+	w, err := newWallet(ctx, clientURI, subnetIDs, ln.walletPrivateKey)
 	if err != nil {
 		return err
 	}
@@ -551,7 +552,7 @@ func (ln *localNetwork) installSubnets(
 	}
 	platformCli := platformvm.NewClient(clientURI)
 
-	w, err := newWallet(ctx, clientURI, []ids.ID{})
+	w, err := newWallet(ctx, clientURI, []ids.ID{}, ln.walletPrivateKey)
 	if err != nil {
 		return nil, err
 	}
@@ -827,8 +828,13 @@ func newWallet(
 	ctx context.Context,
 	uri string,
 	preloadTXs []ids.ID,
+	walletPrivateKey string,
 ) (*wallet, error) {
-	kc := secp256k1fx.NewKeychain(genesis.EWOQKey)
+	var privateKey *secp256k1.PrivateKey
+	if walletPrivateKey == "" {
+		privateKey = genesis.EWOQKey
+	}
+	kc := secp256k1fx.NewKeychain(privateKey)
 	primaryAVAXState, err := primary.FetchState(ctx, uri, kc.Addresses())
 	if err != nil {
 		return nil, err
@@ -850,7 +856,7 @@ func newWallet(
 	pUTXOs := common.NewChainUTXOs(constants.PlatformChainID, utxos)
 	xUTXOs := common.NewChainUTXOs(xCTX.BlockchainID, utxos)
 	var w wallet
-	w.addr = genesis.EWOQKey.PublicKey().Address()
+	w.addr = privateKey.PublicKey().Address()
 	w.pCTX = pCTX
 	w.pBackend = p.NewBackend(pCTX, pUTXOs, pTXs)
 	w.pBuilder = pbuilder.New(kc.Addresses(), pCTX, w.pBackend)
@@ -1028,7 +1034,7 @@ func (ln *localNetwork) removeSubnetValidators(
 		}
 		preloadTXs[i] = subnetID
 	}
-	w, err := newWallet(ctx, clientURI, preloadTXs)
+	w, err := newWallet(ctx, clientURI, preloadTXs, ln.walletPrivateKey)
 	if err != nil {
 		return err
 	}
@@ -1113,7 +1119,7 @@ func (ln *localNetwork) addPermissionlessDelegators(
 		}
 	}
 	// wallet needs txs for all existent subnets
-	w, err := newWallet(ctx, clientURI, subnetIDs)
+	w, err := newWallet(ctx, clientURI, subnetIDs, ln.walletPrivateKey)
 	if err != nil {
 		return err
 	}
@@ -1214,7 +1220,7 @@ func (ln *localNetwork) addPermissionlessValidators(
 		}
 	}
 	// wallet needs txs for all existent subnets
-	w, err := newWallet(ctx, clientURI, subnetIDs)
+	w, err := newWallet(ctx, clientURI, subnetIDs, ln.walletPrivateKey)
 	if err != nil {
 		return err
 	}
@@ -1351,7 +1357,7 @@ func (ln *localNetwork) transformToElasticSubnets(
 		}
 	}
 	// wallet needs txs for all existent subnets
-	w, err := newWallet(ctx, clientURI, subnetIDs)
+	w, err := newWallet(ctx, clientURI, subnetIDs, ln.walletPrivateKey)
 	if err != nil {
 		return nil, nil, err
 	}

diff --git a/local/network.go b/local/network.go
@@ -121,6 +121,8 @@ type localNetwork struct {
 	subnetID2ElasticSubnetID map[ids.ID]ids.ID
 	// map from blockchain id to blockchain aliases
 	blockchainAliases map[string][]string
+	// wallet private key used. IF nil, genesis ewoq key will be used
+	walletPrivateKey string
 }
 
 type deprecatedFlagEsp struct {
@@ -168,6 +170,7 @@ func NewNetwork(
 	reassignPortsIfUsed bool,
 	redirectStdout bool,
 	redirectStderr bool,
+	walletPrivateKey string,
 ) (network.Network, error) {
 	net, err := newNetwork(
 		log,
@@ -184,6 +187,7 @@ func NewNetwork(
 		reassignPortsIfUsed,
 		redirectStdout,
 		redirectStderr,
+		walletPrivateKey,
 	)
 	if err != nil {
 		return net, err
@@ -204,6 +208,7 @@ func newNetwork(
 	reassignPortsIfUsed bool,
 	redirectStdout bool,
 	redirectStderr bool,
+	walletPrivateKey string,
 ) (*localNetwork, error) {
 	var err error
 	if rootDir == "" {
@@ -246,6 +251,7 @@ func newNetwork(
 		redirectStderr:           redirectStderr,
 		subnetID2ElasticSubnetID: map[ids.ID]ids.ID{},
 		blockchainAliases:        map[string][]string{},
+		walletPrivateKey:         walletPrivateKey,
 	}
 	return net, nil
 }
@@ -276,7 +282,7 @@ func NewDefaultNetwork(
 	if err != nil {
 		return nil, err
 	}
-	return NewNetwork(log, config, "", "", "", reassignPortsIfUsed, redirectStdout, redirectStderr)
+	return NewNetwork(log, config, "", "", "", reassignPortsIfUsed, redirectStdout, redirectStderr, "")
 }
 
 func loadDefaultNetworkFiles() (map[string]interface{}, []byte, []*utils.NodeKeys, error) {

diff --git a/local/network_test.go b/local/network_test.go
@@ -153,6 +153,7 @@ func TestNewNetworkEmpty(t *testing.T) {
 		false,
 		false,
 		false,
+		"",
 	)
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
@@ -219,6 +220,7 @@ func TestNewNetworkOneNode(t *testing.T) {
 		false,
 		false,
 		false,
+		"",
 	)
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
@@ -250,6 +252,7 @@ func TestNewNetworkFailToStartNode(t *testing.T) {
 		false,
 		false,
 		false,
+		"",
 	)
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
@@ -486,7 +489,7 @@ func TestWrongNetworkConfigs(t *testing.T) {
 	require := require.New(t)
 	for name, tt := range tests {
 		t.Run(name, func(*testing.T) {
-			net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+			net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 			require.NoError(err)
 			err = net.loadConfig(context.Background(), tt.config)
 			require.Error(err)
@@ -500,7 +503,7 @@ func TestUnhealthyNetwork(t *testing.T) {
 	t.Parallel()
 	require := require.New(t)
 	networkConfig := testNetworkConfig(t)
-	net, err := newNetwork(logging.NoLog{}, newMockAPIUnhealthy, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPIUnhealthy, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
 	require.NoError(err)
@@ -515,7 +518,7 @@ func TestGeneratedNodesNames(t *testing.T) {
 	for i := range networkConfig.NodeConfigs {
 		networkConfig.NodeConfigs[i].Name = ""
 	}
-	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
 	require.NoError(err)
@@ -536,7 +539,7 @@ func TestGenerateDefaultNetwork(t *testing.T) {
 	binaryPath := "pepito"
 	networkConfig, err := NewDefaultConfig(binaryPath, 0)
 	require.NoError(err)
-	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
 	require.NoError(err)
@@ -586,7 +589,7 @@ func TestNetworkFromConfig(t *testing.T) {
 	t.Parallel()
 	require := require.New(t)
 	networkConfig := testNetworkConfig(t)
-	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
 	require.NoError(err)
@@ -610,7 +613,7 @@ func TestNetworkNodeOps(t *testing.T) {
 	// Start a new, empty network
 	emptyNetworkConfig, err := emptyNetworkConfig()
 	require.NoError(err)
-	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), emptyNetworkConfig)
 	require.NoError(err)
@@ -648,7 +651,7 @@ func TestNodeNotFound(t *testing.T) {
 	emptyNetworkConfig, err := emptyNetworkConfig()
 	require.NoError(err)
 	networkConfig := testNetworkConfig(t)
-	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), emptyNetworkConfig)
 	require.NoError(err)
@@ -681,7 +684,7 @@ func TestStoppedNetwork(t *testing.T) {
 	emptyNetworkConfig, err := emptyNetworkConfig()
 	require.NoError(err)
 	networkConfig := testNetworkConfig(t)
-	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), emptyNetworkConfig)
 	require.NoError(err)
@@ -714,7 +717,7 @@ func TestStoppedNetwork(t *testing.T) {
 func TestGetAllNodes(t *testing.T) {
 	require := require.New(t)
 	networkConfig := testNetworkConfig(t)
-	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
 	require.NoError(err)
@@ -769,6 +772,7 @@ func TestFlags(t *testing.T) {
 		false,
 		false,
 		false,
+		"",
 	)
 	require.NoError(err)
 	err = nw.loadConfig(context.Background(), networkConfig)
@@ -801,6 +805,7 @@ func TestFlags(t *testing.T) {
 		false,
 		false,
 		false,
+		"",
 	)
 	require.NoError(err)
 	err = nw.loadConfig(context.Background(), networkConfig)
@@ -832,6 +837,7 @@ func TestFlags(t *testing.T) {
 		false,
 		false,
 		false,
+		"",
 	)
 	require.NoError(err)
 	err = nw.loadConfig(context.Background(), networkConfig)
@@ -1272,7 +1278,7 @@ func TestRemoveBeacon(t *testing.T) {
 	// create a network with no nodes in it
 	emptyNetworkConfig, err := emptyNetworkConfig()
 	require.NoError(err)
-	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPISuccessful, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), emptyNetworkConfig)
 	require.NoError(err)
@@ -1323,7 +1329,7 @@ func TestHealthyDuringNetworkStop(t *testing.T) {
 	require := require.New(t)
 	networkConfig := testNetworkConfig(t)
 	// Calls to a node's Healthy() function blocks until context cancelled
-	net, err := newNetwork(logging.NoLog{}, newMockAPIHealthyBlocks, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false)
+	net, err := newNetwork(logging.NoLog{}, newMockAPIHealthyBlocks, &localTestSuccessfulNodeProcessCreator{}, "", "", "", false, false, false, "")
 	require.NoError(err)
 	err = net.loadConfig(context.Background(), networkConfig)
 	require.NoError(err)

diff --git a/local/snapshot.go b/local/snapshot.go
@@ -75,6 +75,7 @@ func NewNetworkFromSnapshot(
 	redirectStdout bool,
 	redirectStderr bool,
 	inPlace bool,
+	walletPrivateKey string,
 ) (network.Network, error) {
 	if inPlace {
 		if rootDir != "" {
@@ -97,6 +98,7 @@ func NewNetworkFromSnapshot(
 		reassignPortsIfUsed,
 		redirectStdout,
 		redirectStderr,
+		walletPrivateKey,
 	)
 	if err != nil {
 		return net, err

diff --git a/rpcpb/rpc.pb.go b/rpcpb/rpc.pb.go
diff --git a/server/network.go b/server/network.go
@@ -113,6 +113,8 @@ type localNetworkOptions struct {
 	dynamicPorts bool
 
 	networkID uint32
+	// wallet private key used. IF nil, genesis ewoq key will be used
+	walletPrivateKey string
 }
 
 func newLocalNetwork(opts localNetworkOptions) (*localNetwork, error) {
@@ -249,7 +251,7 @@ func (lc *localNetwork) Start(ctx context.Context) error {
 	}
 
 	ux.Print(lc.log, logging.Blue.Wrap(logging.Bold.Wrap("create and run local network")))
-	nw, err := local.NewNetwork(lc.log, lc.cfg, lc.options.rootDataDir, lc.options.logRootDir, lc.options.snapshotsDir, lc.options.reassignPortsIfUsed, lc.options.redirectNodesOutput, lc.options.redirectNodesOutput)
+	nw, err := local.NewNetwork(lc.log, lc.cfg, lc.options.rootDataDir, lc.options.logRootDir, lc.options.snapshotsDir, lc.options.reassignPortsIfUsed, lc.options.redirectNodesOutput, lc.options.redirectNodesOutput, lc.options.walletPrivateKey)
 	if err != nil {
 		return err
 	}
@@ -576,6 +578,7 @@ func (lc *localNetwork) LoadSnapshot(snapshotName string, inPlace bool) error {
 		lc.options.redirectNodesOutput,
 		lc.options.redirectNodesOutput,
 		inPlace,
+		lc.options.walletPrivateKey,
 	)
 	if err != nil {
 		return err

diff --git a/server/server.go b/server/server.go
@@ -334,6 +334,7 @@ func (s *server) Start(_ context.Context, req *rpcpb.StartRequest) (*rpcpb.Start
 
 	s.network, err = newLocalNetwork(localNetworkOptions{
 		networkID:           req.NetworkId,
+		walletPrivateKey:    req.WalletPrivateKey,
 		execPath:            execPath,
 		rootDataDir:         rootDataDir,
 		logRootDir:          req.GetLogRootDir(),