Add data-dir flag (ava-labs#1386)

Co-authored-by: Stephen <[email protected]>

Author: aaronbuchwald

config/config.go

@@ -74,7 +74,7 @@ var (
 func GetRunnerConfig(v *viper.Viper) (runner.Config, error) {
 	config := runner.Config{
 		DisplayVersionAndExit: v.GetBool(VersionKey),
-		BuildDir:              os.ExpandEnv(v.GetString(BuildDirKey)),
+		BuildDir:              GetExpandedArg(v, BuildDirKey),
 		PluginMode:            v.GetBool(PluginModeKey),
 	}
 
@@ -100,6 +100,7 @@ func GetRunnerConfig(v *viper.Viper) (runner.Config, error) {
 
 	foundBuildDir := false
 	for _, dir := range defaultBuildDirs {
+		dir = GetExpandedString(v, dir)
 		if validBuildDir(dir) {
 			config.BuildDir = dir
 			foundBuildDir = true
@@ -136,7 +137,7 @@ func getConsensusConfig(v *viper.Viper) avalanche.Parameters {
 
 func getLoggingConfig(v *viper.Viper) (logging.Config, error) {
 	loggingConfig := logging.Config{}
-	loggingConfig.Directory = os.ExpandEnv(v.GetString(LogsDirKey))
+	loggingConfig.Directory = GetExpandedArg(v, LogsDirKey)
 	var err error
 	loggingConfig.LogLevel, err = logging.ToLevel(v.GetString(LogLevelKey))
 	if err != nil {
@@ -194,7 +195,7 @@ func getIPCConfig(v *viper.Viper) node.IPCConfig {
 		config.IPCDefaultChainIDs = strings.Split(v.GetString(IpcsChainIDsKey), ",")
 	}
 	if v.IsSet(IpcsPathKey) {
-		config.IPCPath = os.ExpandEnv(v.GetString(IpcsPathKey))
+		config.IPCPath = GetExpandedArg(v, IpcsPathKey)
 	}
 	return config
 }
@@ -213,7 +214,7 @@ func getHTTPConfig(v *viper.Viper) (node.HTTPConfig, error) {
 			return node.HTTPConfig{}, fmt.Errorf("unable to decode base64 content: %w", err)
 		}
 	case v.IsSet(HTTPSKeyFileKey):
-		httpsKeyFilepath := os.ExpandEnv(v.GetString(HTTPSKeyFileKey))
+		httpsKeyFilepath := GetExpandedArg(v, HTTPSKeyFileKey)
 		if httpsKey, err = os.ReadFile(filepath.Clean(httpsKeyFilepath)); err != nil {
 			return node.HTTPConfig{}, err
 		}
@@ -227,7 +228,7 @@ func getHTTPConfig(v *viper.Viper) (node.HTTPConfig, error) {
 			return node.HTTPConfig{}, fmt.Errorf("unable to decode base64 content: %w", err)
 		}
 	case v.IsSet(HTTPSCertFileKey):
-		httpsCertFilepath := os.ExpandEnv(v.GetString(HTTPSCertFileKey))
+		httpsCertFilepath := GetExpandedArg(v, HTTPSCertFileKey)
 		if httpsCert, err = os.ReadFile(filepath.Clean(httpsCertFilepath)); err != nil {
 			return node.HTTPConfig{}, err
 		}
@@ -599,7 +600,7 @@ func getIPConfig(v *viper.Viper) (node.IPConfig, error) {
 
 func getProfilerConfig(v *viper.Viper) (profiler.Config, error) {
 	config := profiler.Config{
-		Dir:         os.ExpandEnv(v.GetString(ProfileDirKey)),
+		Dir:         GetExpandedArg(v, ProfileDirKey),
 		Enabled:     v.GetBool(ProfileContinuousEnabledKey),
 		Freq:        v.GetDuration(ProfileContinuousFreqKey),
 		MaxNumFiles: v.GetInt(ProfileContinuousMaxFilesKey),
@@ -633,8 +634,8 @@ func getStakingTLSCertFromFlag(v *viper.Viper) (tls.Certificate, error) {
 
 func getStakingTLSCertFromFile(v *viper.Viper) (tls.Certificate, error) {
 	// Parse the staking key/cert paths and expand environment variables
-	stakingKeyPath := os.ExpandEnv(v.GetString(StakingKeyPathKey))
-	stakingCertPath := os.ExpandEnv(v.GetString(StakingCertPathKey))
+	stakingKeyPath := GetExpandedArg(v, StakingKeyPathKey)
+	stakingCertPath := GetExpandedArg(v, StakingCertPathKey)
 
 	// If staking key/cert locations are specified but not found, error
 	if v.IsSet(StakingKeyPathKey) || v.IsSet(StakingCertPathKey) {
@@ -684,8 +685,8 @@ func getStakingConfig(v *viper.Viper, networkID uint32) (node.StakingConfig, err
 	config := node.StakingConfig{
 		EnableStaking:         v.GetBool(StakingEnabledKey),
 		DisabledStakingWeight: v.GetUint64(StakingDisabledWeightKey),
-		StakingKeyPath:        os.ExpandEnv(v.GetString(StakingKeyPathKey)),
-		StakingCertPath:       os.ExpandEnv(v.GetString(StakingCertPathKey)),
+		StakingKeyPath:        GetExpandedArg(v, StakingKeyPathKey),
+		StakingCertPath:       GetExpandedArg(v, StakingCertPathKey),
 	}
 	if !config.EnableStaking && config.DisabledStakingWeight == 0 {
 		return node.StakingConfig{}, errInvalidStakerWeights
@@ -755,7 +756,7 @@ func getGenesisData(v *viper.Viper, networkID uint32) ([]byte, ids.ID, error) {
 
 	// if content is not specified go for the file
 	if v.IsSet(GenesisConfigFileKey) {
-		genesisFileName := os.ExpandEnv(v.GetString(GenesisConfigFileKey))
+		genesisFileName := GetExpandedArg(v, GenesisConfigFileKey)
 		return genesis.FromFile(networkID, genesisFileName)
 	}
 
@@ -794,7 +795,7 @@ func getDatabaseConfig(v *viper.Viper, networkID uint32) (node.DatabaseConfig, e
 			return node.DatabaseConfig{}, fmt.Errorf("unable to decode base64 content: %w", err)
 		}
 	} else if v.IsSet(DBConfigFileKey) {
-		path := os.ExpandEnv(v.GetString(DBConfigFileKey))
+		path := GetExpandedArg(v, DBConfigFileKey)
 		configBytes, err = os.ReadFile(path)
 		if err != nil {
 			return node.DatabaseConfig{}, err
@@ -804,7 +805,7 @@ func getDatabaseConfig(v *viper.Viper, networkID uint32) (node.DatabaseConfig, e
 	return node.DatabaseConfig{
 		Name: v.GetString(DBTypeKey),
 		Path: filepath.Join(
-			os.ExpandEnv(v.GetString(DBPathKey)),
+			GetExpandedArg(v, DBPathKey),
 			constants.NetworkName(networkID),
 		),
 		Config: configBytes,
@@ -866,7 +867,7 @@ func getVMManager(v *viper.Viper) (vms.Manager, error) {
 
 // getPathFromDirKey reads flag value from viper instance and then checks the folder existence
 func getPathFromDirKey(v *viper.Viper, configKey string) (string, error) {
-	configDir := os.ExpandEnv(v.GetString(configKey))
+	configDir := GetExpandedArg(v, configKey)
 	cleanPath := filepath.Clean(configDir)
 	ok, err := storage.FolderExists(cleanPath)
 	if err != nil {

config/flags.go

@@ -14,6 +14,8 @@ import (
 
 	"github.com/kardianos/osext"
 
+	"github.com/spf13/viper"
+
 	"github.com/ava-labs/avalanchego/database/leveldb"
 	"github.com/ava-labs/avalanchego/database/memdb"
 	"github.com/ava-labs/avalanchego/database/rocksdb"
@@ -26,21 +28,21 @@ import (
 const (
 	DefaultHTTPPort    = 9650
 	DefaultStakingPort = 9651
+
+	AvalancheGoDataDirVar    = "AVALANCHEGO_DATA_DIR"
+	defaultUnexpandedDataDir = "$" + AvalancheGoDataDirVar
 )
 
-// Results of parsing the CLI
 var (
-	defaultNetworkName     = constants.MainnetName
-	homeDir                = os.ExpandEnv("$HOME")
-	prefixedAppName        = fmt.Sprintf(".%s", constants.AppName)
-	defaultDataDir         = filepath.Join(homeDir, prefixedAppName)
-	defaultDBDir           = filepath.Join(defaultDataDir, "db")
-	defaultProfileDir      = filepath.Join(defaultDataDir, "profiles")
-	defaultStakingPath     = filepath.Join(defaultDataDir, "staking")
+	// [defaultUnexpandedDataDir] will be expanded when reading the flags
+	defaultDataDir         = filepath.Join("$HOME", ".avalanchego")
+	defaultDBDir           = filepath.Join(defaultUnexpandedDataDir, "db")
+	defaultLogDir          = filepath.Join(defaultUnexpandedDataDir, "logs")
+	defaultProfileDir      = filepath.Join(defaultUnexpandedDataDir, "profiles")
+	defaultStakingPath     = filepath.Join(defaultUnexpandedDataDir, "staking")
 	defaultStakingKeyPath  = filepath.Join(defaultStakingPath, "staker.key")
 	defaultStakingCertPath = filepath.Join(defaultStakingPath, "staker.crt")
-	defaultLogDirectory    = filepath.Join(defaultDataDir, "logs")
-	defaultConfigDir       = filepath.Join(defaultDataDir, "configs")
+	defaultConfigDir       = filepath.Join(defaultUnexpandedDataDir, "configs")
 	defaultChainConfigDir  = filepath.Join(defaultConfigDir, "chains")
 	defaultVMConfigDir     = filepath.Join(defaultConfigDir, "vms")
 	defaultVMAliasFilePath = filepath.Join(defaultVMConfigDir, "aliases.json")
@@ -63,7 +65,7 @@ func init() {
 	defaultBuildDirs = append(defaultBuildDirs,
 		wd,
 		filepath.Join("/", "usr", "local", "lib", constants.AppName),
-		defaultDataDir,
+		defaultUnexpandedDataDir,
 	)
 }
 
@@ -79,6 +81,8 @@ func addProcessFlags(fs *flag.FlagSet) {
 }
 
 func addNodeFlags(fs *flag.FlagSet) {
+	// Home directory
+	fs.String(DataDirKey, defaultDataDir, "Sets the base data directory where default sub-directories will be placed unless otherwise specified.")
 	// System
 	fs.Uint64(FdLimitKey, ulimit.DefaultFDLimit, "Attempts to raise the process file descriptor limit to at least this value and error if the value is above the system max")
 
@@ -93,7 +97,7 @@ func addNodeFlags(fs *flag.FlagSet) {
 	fs.String(GenesisConfigContentKey, "", "Specifies base64 encoded genesis content")
 
 	// Network ID
-	fs.String(NetworkNameKey, defaultNetworkName, "Network ID this node will connect to")
+	fs.String(NetworkNameKey, constants.MainnetName, "Network ID this node will connect to")
 
 	// AVAX fees
 	fs.Uint64(TxFeeKey, genesis.LocalParams.TxFee, "Transaction fee, in nAVAX")
@@ -108,7 +112,7 @@ func addNodeFlags(fs *flag.FlagSet) {
 	fs.String(DBConfigContentKey, "", "Specifies base64 encoded database config content")
 
 	// Logging
-	fs.String(LogsDirKey, defaultLogDirectory, "Logging directory for Avalanche")
+	fs.String(LogsDirKey, defaultLogDir, "Logging directory for Avalanche")
 	fs.String(LogLevelKey, "info", "The log level. Should be one of {verbo, debug, trace, info, warn, error, fatal, off}")
 	fs.String(LogDisplayLevelKey, "", "The log display level. If left blank, will inherit the value of log-level. Otherwise, should be one of {verbo, debug, trace, info, warn, error, fatal, off}")
 	fs.String(LogFormatKey, "auto", "The structure of log format. Defaults to 'auto' which formats terminal-like logs, when the output is a terminal. Otherwise, should be one of {auto, plain, colors, json}")
@@ -355,3 +359,26 @@ func BuildFlagSet() *flag.FlagSet {
 	addNodeFlags(fs)
 	return fs
 }
+
+// GetExpandedArg gets the string in viper corresponding to [key] and expands
+// any variables using the OS env. If the [AvalancheGoDataDirVar] var is used,
+// we expand the value of the variable with the string in viper corresponding to
+// [DataDirKey].
+func GetExpandedArg(v *viper.Viper, key string) string {
+	return GetExpandedString(v, v.GetString(key))
+}
+
+// GetExpandedString expands [s] with any variables using the OS env. If the
+// [AvalancheGoDataDirVar] var is used, we expand the value of the variable with
+// the string in viper corresponding to [DataDirKey].
+func GetExpandedString(v *viper.Viper, s string) string {
+	return os.Expand(
+		s,
+		func(strVar string) string {
+			if strVar == AvalancheGoDataDirVar {
+				return os.ExpandEnv(v.GetString(DataDirKey))
+			}
+			return os.Getenv(strVar)
+		},
+	)
+}

config/keys.go

@@ -5,6 +5,7 @@ package config
 
 // #nosec G101
 const (
+	DataDirKey                                         = "data-dir"
 	ConfigFileKey                                      = "config-file"
 	ConfigContentKey                                   = "config-file-content"
 	ConfigContentTypeKey                               = "config-file-content-type"

config/viper.go

@@ -9,7 +9,6 @@ import (
 	"flag"
 	"fmt"
 	"io"
-	"os"
 	"strings"
 
 	"github.com/spf13/viper"
@@ -49,7 +48,7 @@ func BuildViper(fs *flag.FlagSet, args []string) (*viper.Viper, error) {
 		}
 
 	case v.IsSet(ConfigFileKey):
-		filename := os.ExpandEnv(v.GetString(ConfigFileKey))
+		filename := GetExpandedArg(v, ConfigFileKey)
 		v.SetConfigFile(filename)
 		if err := v.ReadInConfig(); err != nil {
 			return nil, err