feature: support init from environment variables for dfdaemon and sup…

…ernode

Signed-off-by: SataQiu <[email protected]>

cmd/dfdaemon/app/root.go

@@ -39,6 +39,12 @@ import (
 	"gopkg.in/yaml.v2"
 )
 
+const (
+	// DFDaemonEnvPrefix is the default environment prefix for Viper.
+	// Both BindEnv and AutomaticEnv will use this prefix.
+	DFDaemonEnvPrefix = "dfdaemon"
+)
+
 var rootCmd = &cobra.Command{
 	Use:               "dfdaemon",
 	Short:             "The dfdaemon is a proxy that intercepts image download requests.",
@@ -106,7 +112,13 @@ func bindRootFlags(v *viper.Viper) error {
 	if err := v.BindPFlags(rootCmd.Flags()); err != nil {
 		return err
 	}
-	return v.BindPFlag("registry_mirror.remote", rootCmd.Flag("registry"))
+	if err := v.BindPFlag("registry_mirror.remote", rootCmd.Flag("registry")); err != nil {
+		return err
+	}
+	v.SetEnvPrefix(DFDaemonEnvPrefix)
+	v.AutomaticEnv()
+
+	return nil
 }
 
 // readConfigFile reads config file into the given viper instance. If we're

cmd/dfdaemon/app/root_test.go

@@ -21,6 +21,7 @@ import (
 	"io/ioutil"
 	"math/rand"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -136,6 +137,28 @@ func (ts *rootTestSuite) TestBindRootFlags() {
 	r.Equal(v.GetString("registry"), v.GetString("registry_mirror.remote"))
 }
 
+func (ts *rootTestSuite) TestAutomaticEnv() {
+	r := ts.Require()
+	v := viper.New()
+	r.Nil(bindRootFlags(v))
+
+	maxprocsEnvKey := strings.ToUpper(DFDaemonEnvPrefix + "_maxprocs")
+	workHomeEnvKey := strings.ToUpper(DFDaemonEnvPrefix + "_workHome")
+	registryEnvKey := strings.ToUpper(DFDaemonEnvPrefix + "_registry_mirror.remote")
+
+	os.Setenv(maxprocsEnvKey, "17")
+	os.Setenv(workHomeEnvKey, "/dragonfly/home")
+	os.Setenv(registryEnvKey, "https://dragonfly.io")
+
+	r.Equal(17, v.GetInt("maxprocs"))
+	r.Equal("/dragonfly/home", v.GetString("workHome"))
+	r.Equal("https://dragonfly.io", v.GetString("registry_mirror.remote"))
+
+	os.Unsetenv(maxprocsEnvKey)
+	os.Unsetenv(workHomeEnvKey)
+	os.Unsetenv(registryEnvKey)
+}
+
 var testCrt = `-----BEGIN CERTIFICATE-----
 MIICKzCCAZQCCQDZrCsm2rX81DANBgkqhkiG9w0BAQUFADBaMQswCQYDVQQGEwJD
 TjERMA8GA1UECAwIWmhlamlhbmcxETAPBgNVBAcMCEhhbmd6aG91MQ4wDAYDVQQK

cmd/supernode/app/root.go

@@ -41,6 +41,12 @@ import (
 	"gopkg.in/yaml.v2"
 )
 
+const (
+	// SupernodeEnvPrefix is the default environment prefix for Viper.
+	// Both BindEnv and AutomaticEnv will use this prefix.
+	SupernodeEnvPrefix = "supernode"
+)
+
 var (
 	supernodeViper = viper.GetViper()
 )
@@ -259,6 +265,10 @@ func bindRootFlags(v *viper.Viper) error {
 			return err
 		}
 	}
+
+	v.SetEnvPrefix(SupernodeEnvPrefix)
+	v.AutomaticEnv()
+
 	return nil
 }
 

cmd/supernode/app/root_test.go

@@ -21,6 +21,7 @@ import (
 	"math/rand"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"time"
 
@@ -140,6 +141,33 @@ func (ts *rootTestSuite) TestHomeDirFlag() {
 	}
 }
 
+func (ts *rootTestSuite) TestAutomaticEnv() {
+	r := ts.Require()
+	v := viper.New()
+	r.Nil(bindRootFlags(v))
+
+	listenPortEnvKey := strings.ToUpper(SupernodeEnvPrefix + "_base.listenPort")
+	homeDirEnvKey := strings.ToUpper(SupernodeEnvPrefix + "_base.homeDir")
+	debugEnvKey := strings.ToUpper(SupernodeEnvPrefix + "_base.debug")
+	failAccessIntervalEnvKey := strings.ToUpper(SupernodeEnvPrefix + "_base.failAccessInterval")
+
+	os.Setenv(listenPortEnvKey, "2019")
+	os.Setenv(homeDirEnvKey, "/dragonfly/home")
+	os.Setenv(debugEnvKey, "true")
+	os.Setenv(failAccessIntervalEnvKey, "10m30s")
+	expectFailAccessInterval, _ := time.ParseDuration("10m30s")
+
+	r.Equal(2019, v.GetInt("base.listenPort"))
+	r.Equal("/dragonfly/home", v.GetString("base.homeDir"))
+	r.Equal(true, v.GetBool("base.debug"))
+	r.Equal(expectFailAccessInterval, v.GetDuration("base.failAccessInterval"))
+
+	os.Unsetenv(listenPortEnvKey)
+	os.Unsetenv(homeDirEnvKey)
+	os.Unsetenv(debugEnvKey)
+	os.Unsetenv(failAccessIntervalEnvKey)
+}
+
 func generateFakeFilename(fs afero.Fs) string {
 	for i := 0; i < 100; i++ {
 		d := fmt.Sprintf("/supernode-test-%d-%d", time.Now().UnixNano(), rand.Int())