Merge pull request #524 from 0xPolygon/thiago/env-vars

tclemos · web-flow · commit 7d670ea00844 · 2025-02-28T13:31:33.000-03:00
add a flag loader to allow flags to be set also via env vars
diff --git a/cmd/cdk/cdk.go b/cmd/cdk/cdk.go
@@ -9,6 +9,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/0xPolygon/polygon-cli/cmd/flag_loader"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/ethclient"
 	"github.com/go-errors/errors"
@@ -86,7 +87,10 @@ var CDKCmd = &cobra.Command{
 	Use:   "cdk",
 	Short: "Utilities for interacting with CDK networks",
 	Long:  "Basic utility commands for interacting with the cdk contracts",
-	Args:  cobra.NoArgs,
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		cdkInputArgs.rpcURL = flag_loader.GetRpcUrlFlagValue(cmd)
+	},
+	Args: cobra.NoArgs,
 }
 
 type inputArgs struct {
diff --git a/cmd/dumpblocks/dumpblocks.go b/cmd/dumpblocks/dumpblocks.go
@@ -12,6 +12,7 @@ import (
 
 	_ "embed"
 
+	"github.com/0xPolygon/polygon-cli/cmd/flag_loader"
 	"github.com/0xPolygon/polygon-cli/proto/gen/pb"
 	"github.com/0xPolygon/polygon-cli/rpctypes"
 	"github.com/0xPolygon/polygon-cli/util"
@@ -54,6 +55,10 @@ var DumpblocksCmd = &cobra.Command{
 	Use:   "dumpblocks start end",
 	Short: "Export a range of blocks from a JSON-RPC endpoint.",
 	Long:  usage,
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		rpcUrlFlagValue := flag_loader.GetRpcUrlFlagValue(cmd)
+		inputDumpblocks.RpcUrl = *rpcUrlFlagValue
+	},
 	PreRunE: func(cmd *cobra.Command, args []string) error {
 		return checkFlags()
 	},
diff --git a/cmd/ecrecover/ecrecover.go b/cmd/ecrecover/ecrecover.go
@@ -7,6 +7,7 @@ import (
 	"math/big"
 	"os"
 
+	"github.com/0xPolygon/polygon-cli/cmd/flag_loader"
 	"github.com/0xPolygon/polygon-cli/util"
 	ethcommon "github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -31,6 +32,10 @@ var EcRecoverCmd = &cobra.Command{
 	Short: "Recovers and returns the public key of the signature",
 	Long:  usage,
 	Args:  cobra.NoArgs,
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		rpcUrlFlagValue := flag_loader.GetRpcUrlFlagValue(cmd)
+		rpcUrl = *rpcUrlFlagValue
+	},
 	PreRunE: func(cmd *cobra.Command, args []string) error {
 		return checkFlags()
 	},
diff --git a/cmd/fixnoncegap/fixnoncegap.go b/cmd/fixnoncegap/fixnoncegap.go
@@ -9,6 +9,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/0xPolygon/polygon-cli/cmd/flag_loader"
 	"github.com/ethereum/go-ethereum"
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/ethereum/go-ethereum/common"
@@ -20,10 +21,19 @@ import (
 )
 
 var FixNonceGapCmd = &cobra.Command{
-	Use:          "fix-nonce-gap",
-	Short:        "Send txs to fix the nonce gap for a specific account",
-	Long:         fixNonceGapUsage,
-	Args:         cobra.NoArgs,
+	Use:   "fix-nonce-gap",
+	Short: "Send txs to fix the nonce gap for a specific account",
+	Long:  fixNonceGapUsage,
+	Args:  cobra.NoArgs,
+	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+		var err error
+		inputFixNonceGapArgs.rpcURL = flag_loader.GetRpcUrlFlagValue(cmd)
+		inputFixNonceGapArgs.privateKey, err = flag_loader.GetRequiredPrivateKeyFlagValue(cmd)
+		if err != nil {
+			return err
+		}
+		return nil
+	},
 	PreRunE:      prepareRpcClient,
 	RunE:         fixNonceGap,
 	SilenceUsage: true,
@@ -221,7 +231,6 @@ func init() {
 	inputFixNonceGapArgs.privateKey = FixNonceGapCmd.PersistentFlags().String(ArgPrivateKey, "", "the private key to be used when sending the txs to fix the nonce gap")
 	inputFixNonceGapArgs.replace = FixNonceGapCmd.PersistentFlags().Bool(ArgReplace, false, "replace the existing txs in the pool")
 	inputFixNonceGapArgs.maxNonce = FixNonceGapCmd.PersistentFlags().Uint64(ArgMaxNonce, 0, "when set, the max nonce will be this value instead of trying to get it from the pool")
-	fatalIfError(FixNonceGapCmd.MarkPersistentFlagRequired(ArgPrivateKey))
 }
 
 // Wait for the transaction to be mined
@@ -255,13 +264,6 @@ func WaitMineTransaction(ctx context.Context, client *ethclient.Client, tx *type
 	}
 }
 
-func fatalIfError(err error) {
-	if err == nil {
-		return
-	}
-	log.Fatal().Err(err).Msg("Unexpected error occurred")
-}
-
 func getMaxNonceFromTxPool(addr common.Address) (uint64, error) {
 	var result PoolContent
 	err := rpcClient.Client().Call(&result, "txpool_content")
diff --git a/cmd/flag_loader/flag_loader.go b/cmd/flag_loader/flag_loader.go
@@ -0,0 +1,56 @@
+package flag_loader
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/spf13/cobra"
+)
+
+const (
+	rpcUrlFlagName, rpcUrlEnvVar         = "rpc-url", "ETH_RPC_URL"
+	privateKeyFlagName, privateKeyEnvVar = "private-key", "PRIVATE_KEY"
+)
+
+func GetRpcUrlFlagValue(cmd *cobra.Command) *string {
+	v, _ := getFlagValue(cmd, rpcUrlFlagName, rpcUrlEnvVar, false)
+	return v
+}
+
+func GetRequiredRpcUrlFlagValue(cmd *cobra.Command) (*string, error) {
+	return getFlagValue(cmd, rpcUrlFlagName, rpcUrlEnvVar, true)
+}
+
+func GetPrivateKeyFlagValue(cmd *cobra.Command) *string {
+	v, _ := getFlagValue(cmd, privateKeyFlagName, privateKeyEnvVar, false)
+	return v
+}
+
+func GetRequiredPrivateKeyFlagValue(cmd *cobra.Command) (*string, error) {
+	return getFlagValue(cmd, privateKeyFlagName, privateKeyEnvVar, true)
+}
+
+func getFlagValue(cmd *cobra.Command, flagName, envVarName string, required bool) (*string, error) {
+	flag := cmd.Flag(flagName)
+	var flagValue string
+	if flag.Changed {
+		flagValue = flag.Value.String()
+	}
+	flagDefaultValue := flag.DefValue
+
+	envVarValue := os.Getenv(envVarName)
+
+	value := flagDefaultValue
+	if envVarValue != "" {
+		value = envVarValue
+	}
+	if flag.Changed {
+		value = flagValue
+	}
+
+	if required && (!flag.Changed && envVarValue == "") {
+		return nil, fmt.Errorf("required flag(s) \"%s\" not set", flagName)
+	}
+
+	return &value, nil
+}
diff --git a/cmd/flag_loader/flag_loader_test.go b/cmd/flag_loader/flag_loader_test.go
@@ -0,0 +1,155 @@
+package flag_loader
+
+import (
+	"fmt"
+	"os"
+	"strconv"
+	"testing"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestValuePriority(t *testing.T) {
+	type testCase struct {
+		defaultValue *int
+		envVarValue  *int
+		flagValue    *int
+		required     bool
+
+		expectedValue *int
+		expectedError error
+	}
+
+	testCases := []testCase{
+		{
+			defaultValue:  ptr(1),
+			envVarValue:   ptr(2),
+			flagValue:     ptr(3),
+			expectedValue: ptr(3),
+			required:      true,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  ptr(1),
+			envVarValue:   ptr(2),
+			flagValue:     ptr(1),
+			expectedValue: ptr(1),
+			required:      true,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  ptr(1),
+			envVarValue:   ptr(2),
+			flagValue:     nil,
+			expectedValue: ptr(2),
+			required:      true,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  ptr(1),
+			envVarValue:   nil,
+			flagValue:     ptr(3),
+			expectedValue: ptr(3),
+			required:      true,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  nil,
+			envVarValue:   ptr(2),
+			flagValue:     ptr(3),
+			expectedValue: ptr(3),
+			required:      true,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  nil,
+			envVarValue:   nil,
+			flagValue:     ptr(3),
+			expectedValue: ptr(3),
+			required:      true,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  ptr(1),
+			envVarValue:   nil,
+			flagValue:     nil,
+			expectedValue: ptr(1),
+			required:      false,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  nil,
+			envVarValue:   ptr(2),
+			flagValue:     nil,
+			expectedValue: ptr(2),
+			required:      true,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  nil,
+			envVarValue:   nil,
+			flagValue:     nil,
+			expectedValue: ptr(0),
+			required:      false,
+			expectedError: nil,
+		},
+		{
+			defaultValue:  nil,
+			envVarValue:   nil,
+			flagValue:     nil,
+			expectedValue: nil,
+			required:      true,
+			expectedError: fmt.Errorf("required flag(s) \"flag\" not set"),
+		},
+	}
+
+	for _, tc := range testCases {
+		var value *int
+		cmd := &cobra.Command{
+			Use: "test",
+			PersistentPreRun: func(cmd *cobra.Command, args []string) {
+				valueStr, err := getFlagValue(cmd, "flag", "FLAG", tc.required)
+				if tc.expectedError != nil {
+					assert.EqualError(t, err, tc.expectedError.Error())
+					return
+				}
+				assert.NoError(t, err)
+				valueInt, err := strconv.Atoi(*valueStr)
+				assert.NoError(t, err)
+				value = &valueInt
+			},
+			Run: func(cmd *cobra.Command, args []string) {
+				if tc.expectedValue != nil {
+					assert.Equal(t, *tc.expectedValue, *value)
+				} else {
+					assert.Nil(t, value)
+				}
+			},
+		}
+		if tc.defaultValue != nil {
+			cmd.Flags().Int("flag", *tc.defaultValue, "flag")
+		} else {
+			cmd.Flags().Int("flag", 0, "flag")
+		}
+
+		os.Unsetenv("FLAG")
+		if tc.envVarValue != nil {
+			v := strconv.Itoa(*tc.envVarValue)
+			os.Setenv("FLAG", v)
+		}
+
+		if tc.flagValue != nil {
+			v := strconv.Itoa(*tc.flagValue)
+			cmd.SetArgs([]string{"--flag", v})
+		}
+
+		err := cmd.Execute()
+		assert.Nil(t, err)
+	}
+
+}
+
+func ptr[T any](v T) *T {
+	return &v
+}
diff --git a/cmd/fund/cmd.go b/cmd/fund/cmd.go
@@ -6,6 +6,7 @@ import (
 
 	_ "embed"
 
+	"github.com/0xPolygon/polygon-cli/cmd/flag_loader"
 	"github.com/0xPolygon/polygon-cli/util"
 	"github.com/spf13/cobra"
 )
@@ -44,6 +45,10 @@ var FundCmd = &cobra.Command{
 	Use:   "fund",
 	Short: "Bulk fund crypto wallets automatically.",
 	Long:  usage,
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		params.RpcUrl = flag_loader.GetRpcUrlFlagValue(cmd)
+		params.PrivateKey = flag_loader.GetPrivateKeyFlagValue(cmd)
+	},
 	PreRunE: func(cmd *cobra.Command, args []string) error {
 		return checkFlags()
 	},
diff --git a/cmd/loadtest/app.go b/cmd/loadtest/app.go
@@ -9,6 +9,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/0xPolygon/polygon-cli/cmd/flag_loader"
 	"github.com/0xPolygon/polygon-cli/rpctypes"
 	"github.com/0xPolygon/polygon-cli/util"
 	ethcommon "github.com/ethereum/go-ethereum/common"
@@ -166,6 +167,10 @@ var LoadtestCmd = &cobra.Command{
 	Short: "Run a generic load test against an Eth/EVM style JSON-RPC endpoint.",
 	Long:  loadtestUsage,
 	Args:  cobra.NoArgs,
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		inputLoadTestParams.RPCUrl = flag_loader.GetRpcUrlFlagValue(cmd)
+		inputLoadTestParams.PrivateKey = flag_loader.GetPrivateKeyFlagValue(cmd)
+	},
 	PreRunE: func(cmd *cobra.Command, args []string) error {
 		zerolog.DurationFieldUnit = time.Second
 		zerolog.DurationFieldInteger = true
diff --git a/cmd/monitor/cmd.go b/cmd/monitor/cmd.go
@@ -7,6 +7,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/0xPolygon/polygon-cli/cmd/flag_loader"
 	"github.com/0xPolygon/polygon-cli/util"
 	"github.com/spf13/cobra"
 )
@@ -57,6 +58,8 @@ var MonitorCmd = &cobra.Command{
 	Args:         cobra.NoArgs,
 	SilenceUsage: true,
 	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+		rpcUrlFlagValue := flag_loader.GetRpcUrlFlagValue(cmd)
+		rpcUrl = *rpcUrlFlagValue
 		// By default, hide logs from `polycli monitor`.
 		verbosityFlag := cmd.Flag("verbosity")
 		if verbosityFlag != nil && !verbosityFlag.Changed {
diff --git a/cmd/nodekey/nodekey.go b/cmd/nodekey/nodekey.go
@@ -15,6 +15,7 @@ import (
 
 	_ "embed"
 
+	"github.com/0xPolygon/polygon-cli/cmd/flag_loader"
 	gethcrypto "github.com/ethereum/go-ethereum/crypto"
 	gethenode "github.com/ethereum/go-ethereum/p2p/enode"
 	libp2pcrypto "github.com/libp2p/go-libp2p/core/crypto"
@@ -65,6 +66,9 @@ var NodekeyCmd = &cobra.Command{
 	Use:   "nodekey",
 	Short: "Generate node keys for different blockchain clients and protocols.",
 	Long:  usage,
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		inputNodeKeyPrivateKey = flag_loader.GetRpcUrlFlagValue(cmd)
+	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		var nko nodeKeyOut
 		var withSeed bool
diff --git a/cmd/rpcfuzz/cmd.go b/cmd/rpcfuzz/cmd.go
diff --git a/cmd/signer/signer.go b/cmd/signer/signer.go
diff --git a/cmd/ulxly/ulxly.go b/cmd/ulxly/ulxly.go
