feat: Add a key file to the fund command (#672)

* feat: add key-file to fund command

* fund key-file flag changes

* review adjusts

* make lint

Author: tclemos

cmd/fund/cmd.go

@@ -31,6 +31,8 @@ type cmdFundParams struct {
 	FundingAmountInWei *big.Int
 	OutputFile         *string
 
+	KeyFile *string
+
 	FunderAddress *string
 }
 
@@ -71,12 +73,16 @@ func init() {
 	p.WalletAddresses = flagSet.StringSlice("addresses", nil, "Comma-separated list of wallet addresses to fund")
 	p.FundingAmountInWei = defaultFundingInWei
 	flagSet.Var(&flag_loader.BigIntValue{Val: p.FundingAmountInWei}, "eth-amount", "The amount of wei to send to each wallet")
+	p.KeyFile = flagSet.String("key-file", "", "The file containing the accounts private keys, one per line.")
 
 	p.OutputFile = flagSet.StringP("file", "f", "wallets.json", "The output JSON file path for storing the addresses and private keys of funded wallets")
 
 	// Marking flags as mutually exclusive
 	FundCmd.MarkFlagsMutuallyExclusive("addresses", "number")
 	FundCmd.MarkFlagsMutuallyExclusive("addresses", "hd-derivation")
+	FundCmd.MarkFlagsMutuallyExclusive("key-file", "addresses")
+	FundCmd.MarkFlagsMutuallyExclusive("key-file", "number")
+	FundCmd.MarkFlagsMutuallyExclusive("key-file", "hd-derivation")
 
 	// Funder contract parameters.
 	p.FunderAddress = flagSet.String("contract-address", "", "The address of a pre-deployed Funder contract")
@@ -98,10 +104,14 @@ func checkFlags() error {
 		return errors.New("the private key is empty")
 	}
 
-	// Check wallet flags.
-	if params.WalletsNumber != nil && *params.WalletsNumber == 0 {
-		return errors.New("the number of wallets to fund is set to zero")
+	// Check that exactly one method is used to specify target accounts
+	hasAddresses := params.WalletAddresses != nil && len(*params.WalletAddresses) > 0
+	hasKeyFile := params.KeyFile != nil && *params.KeyFile != ""
+	hasNumberFlag := params.WalletsNumber != nil && *params.WalletsNumber > 0
+	if !hasAddresses && !hasKeyFile && !hasNumberFlag {
+		return errors.New("must specify target accounts via --addresses, --key-file, or --number")
 	}
+
 	minValue := big.NewInt(1000000000)
 	if params.FundingAmountInWei != nil && params.FundingAmountInWei.Cmp(minValue) <= 0 {
 		return errors.New("the funding amount must be greater than 1000000000")

cmd/fund/fund.go

@@ -6,6 +6,7 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"math/big"
 	"os"
 	"strings"
@@ -49,21 +50,14 @@ func runFunding(ctx context.Context) error {
 		return err
 	}
 
-	// Derive or generate a set of wallets.
 	var addresses []common.Address
-	if params.WalletAddresses != nil && *params.WalletAddresses != nil {
-		log.Info().Msg("Using addresses provided by the user")
-		addresses = make([]common.Address, len(*params.WalletAddresses))
-		for i, address := range *params.WalletAddresses {
-			addresses[i] = common.HexToAddress(address)
-		}
-	} else if *params.UseHDDerivation {
-		log.Info().Msg("Deriving wallets from the default mnemonic")
-		addresses, err = deriveHDWallets(int(*params.WalletsNumber))
-	} else {
-		log.Info().Msg("Generating random wallets")
-		addresses, err = generateWallets(int(*params.WalletsNumber))
+
+	if len(*params.KeyFile) > 0 { // get addresses from key-file
+		addresses, err = getAddressesFromKeyFile(*params.KeyFile)
+	} else { // get addresses from private key
+		addresses, err = getAddressesFromPrivateKey(ctx, c)
 	}
+	// check errors after getting addresses
 	if err != nil {
 		return err
 	}
@@ -85,6 +79,57 @@ func runFunding(ctx context.Context) error {
 	return nil
 }
 
+func getAddressesFromKeyFile(keyFilePath string) ([]common.Address, error) {
+	if len(keyFilePath) == 0 {
+		return nil, errors.New("the key file path is empty")
+	}
+
+	log.Trace().
+		Str("keyFilePath", keyFilePath).
+		Msg("getting addresses from key file")
+
+	privateKeys, iErr := util.ReadPrivateKeysFromFile(keyFilePath)
+	if iErr != nil {
+		log.Error().
+			Err(iErr).
+			Msg("Unable to read private keys from key file")
+		return nil, fmt.Errorf("unable to read private keys from key file. %w", iErr)
+	}
+	addresses := make([]common.Address, len(privateKeys))
+	for i, privateKey := range privateKeys {
+		addresses[i] = util.GetAddress(context.Background(), privateKey)
+		log.Trace().
+			Interface("address", addresses[i]).
+			Str("privateKey", hex.EncodeToString(privateKey.D.Bytes())).
+			Msg("New wallet derived from key file")
+	}
+	log.Info().Int("count", len(addresses)).Msg("Wallet(s) derived from key file")
+	return addresses, nil
+}
+
+func getAddressesFromPrivateKey(ctx context.Context, c *ethclient.Client) ([]common.Address, error) {
+	// Derive or generate a set of wallets.
+	var addresses []common.Address
+	var err error
+	if params.WalletAddresses != nil && *params.WalletAddresses != nil {
+		log.Info().Msg("Using addresses provided by the user")
+		addresses = make([]common.Address, len(*params.WalletAddresses))
+		for i, address := range *params.WalletAddresses {
+			addresses[i] = common.HexToAddress(address)
+		}
+	} else if *params.UseHDDerivation {
+		log.Info().Msg("Deriving wallets from the default mnemonic")
+		addresses, err = deriveHDWallets(int(*params.WalletsNumber))
+	} else {
+		log.Info().Msg("Generating random wallets")
+		addresses, err = generateWallets(int(*params.WalletsNumber))
+	}
+	if err != nil {
+		return nil, err
+	}
+	return addresses, nil
+}
+
 // dialRpc dials the Ethereum RPC server and return an Ethereum client.
 func dialRpc(ctx context.Context) (*ethclient.Client, error) {
 	rpc, err := rpc.DialContext(ctx, *params.RpcUrl)

cmd/fund/usage.md

@@ -24,6 +24,13 @@ $ polycli fund --number 20 --hd-derivation=false
 3:58PM INF Wallets' address(es) and private key(s) saved to file fileName=wallets.json
 3:58PM INF Wallet(s) funded! 💸
 3:58PM INF Total execution time: 1.027506s
+
+# Fund wallets from a key file (one private key in hex per line).
+$ polycli fund --key-file=keys.txt
+3:58PM INF Starting bulk funding wallets
+3:58PM INF Wallet(s) derived from key file count=3
+3:58PM INF Wallet(s) funded! 💸
+3:58PM INF Total execution time: 1.2s
 ```
 
 Extract from `wallets.json`.

cmd/loadtest/account.go

@@ -11,6 +11,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/0xPolygon/polygon-cli/util"
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -202,7 +203,7 @@ func (ap *AccountPool) Add(ctx context.Context, privateKey *ecdsa.PrivateKey, st
 		return fmt.Errorf("failed to create account: %w", err)
 	}
 
-	addressHex, privateKeyHex := getAddressAndPrivateKeyHex(ctx, privateKey)
+	addressHex, privateKeyHex := util.GetAddressAndPrivateKeyHex(ctx, privateKey)
 	log.Debug().
 		Str("address", addressHex).
 		Str("privateKey", privateKeyHex).
@@ -504,7 +505,7 @@ func (ap *AccountPool) ReturnFunds(ctx context.Context) error {
 	txCh := make(chan *types.Transaction, len(ap.accounts))
 	errCh := make(chan error, len(ap.accounts))
 
-	fundingAddressHex, _ := getAddressAndPrivateKeyHex(ctx, ap.fundingPrivateKey)
+	fundingAddressHex, _ := util.GetAddressAndPrivateKeyHex(ctx, ap.fundingPrivateKey)
 	fundingAddress := common.HexToAddress(fundingAddressHex)
 
 	err = ap.clientRateLimiter.Wait(ctx)
@@ -934,14 +935,3 @@ func (ap *AccountPool) createEOATransferTx(ctx context.Context, sender *ecdsa.Pr
 
 	return signedTx, nil
 }
-
-// Returns the address and private key of the given private key
-func getAddressAndPrivateKeyHex(ctx context.Context, privateKey *ecdsa.PrivateKey) (string, string) {
-	privateKeyBytes := crypto.FromECDSA(privateKey)
-	privateKeyHex := fmt.Sprintf("0x%x", privateKeyBytes)
-
-	publicKey := privateKey.Public().(*ecdsa.PublicKey)
-	address := crypto.PubkeyToAddress(*publicKey)
-
-	return address.String(), privateKeyHex
-}

cmd/loadtest/loadtest.go

@@ -1,9 +1,7 @@
 package loadtest
 
 import (
-	"bufio"
 	"context"
-	"crypto/ecdsa"
 	_ "embed"
 	"encoding/hex"
 	"encoding/json"
@@ -317,7 +315,7 @@ func initializeLoadTestParams(ctx context.Context, c *ethclient.Client) error {
 			Str("sendingAccountsFile", sendingAccountsFile).
 			Msg("Adding accounts from file to the account pool")
 
-		privateKeys, iErr := readPrivateKeysFromFile(sendingAccountsFile)
+		privateKeys, iErr := util.ReadPrivateKeysFromFile(sendingAccountsFile)
 		if iErr != nil {
 			log.Error().
 				Err(iErr).
@@ -388,35 +386,6 @@ func initializeLoadTestParams(ctx context.Context, c *ethclient.Client) error {
 	return nil
 }
 
-func readPrivateKeysFromFile(sendingAccountsFile string) ([]*ecdsa.PrivateKey, error) {
-	file, err := os.Open(sendingAccountsFile)
-	if err != nil {
-		return nil, fmt.Errorf("unable to open sending accounts file: %w", err)
-	}
-	defer file.Close()
-
-	var privateKeys []*ecdsa.PrivateKey
-	scanner := bufio.NewScanner(file)
-	for scanner.Scan() {
-		line := strings.TrimSpace(scanner.Text())
-		if len(line) == 0 {
-			continue
-		}
-		privateKey, err := ethcrypto.HexToECDSA(strings.TrimPrefix(line, "0x"))
-		if err != nil {
-			log.Error().Err(err).Str("key", line).Msg("Unable to parse private key")
-			return nil, fmt.Errorf("unable to parse private key: %w", err)
-		}
-		privateKeys = append(privateKeys, privateKey)
-	}
-
-	if err := scanner.Err(); err != nil {
-		return nil, fmt.Errorf("error reading sending accounts file: %w", err)
-	}
-
-	return privateKeys, nil
-}
-
 func completeLoadTest(ctx context.Context, c *ethclient.Client, rpc *ethrpc.Client) error {
 	if *inputLoadTestParams.FireAndForget {
 		log.Info().

doc/polycli_fund.md

@@ -45,6 +45,13 @@ $ polycli fund --number 20 --hd-derivation=false
 3:58PM INF Wallets' address(es) and private key(s) saved to file fileName=wallets.json
 3:58PM INF Wallet(s) funded! 💸
 3:58PM INF Total execution time: 1.027506s
+
+# Fund wallets from a key file (one private key in hex per line).
+$ polycli fund --key-file=keys.txt
+3:58PM INF Starting bulk funding wallets
+3:58PM INF Wallet(s) derived from key file count=3
+3:58PM INF Wallet(s) funded! 💸
+3:58PM INF Total execution time: 1.2s
 ```
 
 Extract from `wallets.json`.
@@ -83,6 +90,7 @@ $ cast balance 0x5D8121cf716B70d3e345adB58157752304eED5C3
   -f, --file string               The output JSON file path for storing the addresses and private keys of funded wallets (default "wallets.json")
       --hd-derivation             Derive wallets to fund from the private key in a deterministic way (default true)
   -h, --help                      help for fund
+      --key-file string           The file containing the accounts private keys, one per line.
   -n, --number uint               The number of wallets to fund (default 10)
       --private-key string        The hex encoded private key that we'll use to send transactions (default "0x42b6e34dc21598a807dc19d7784c71b2a7a01f6480dc6f58258f78e539f1a1fa")
   -r, --rpc-url string            The RPC endpoint url (default "http://localhost:8545")