Add support for both config-file and data-dir at a global level in the self-extracting wrapper for K3s (#2594)

* Add support for both config-file and data-dir at a global level in the self-extracting wrapper for K3s

Signed-off-by: Chris Kim <oats87g@gmail.com>

cmd/k3s/main.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/pkg/errors"
 	"github.com/rancher/k3s/pkg/cli/cmds"
+	"github.com/rancher/k3s/pkg/configfilearg"
 	"github.com/rancher/k3s/pkg/data"
 	"github.com/rancher/k3s/pkg/datadir"
 	"github.com/rancher/k3s/pkg/dataverify"
@@ -21,18 +22,19 @@ import (
 )
 
 func main() {
-	if runCLIs() {
+	dataDir := findDataDir()
+	if runCLIs(dataDir) {
 		return
 	}
 
 	app := cmds.NewApp()
 	app.Commands = []cli.Command{
-		cmds.NewServerCommand(wrap(version.Program+"-server", os.Args)),
-		cmds.NewAgentCommand(wrap(version.Program+"-agent", os.Args)),
-		cmds.NewKubectlCommand(externalCLIAction("kubectl")),
-		cmds.NewCRICTL(externalCLIAction("crictl")),
-		cmds.NewCtrCommand(externalCLIAction("ctr")),
-		cmds.NewCheckConfigCommand(externalCLIAction("check-config")),
+		cmds.NewServerCommand(wrap(version.Program+"-server", dataDir, os.Args)),
+		cmds.NewAgentCommand(wrap(version.Program+"-agent", dataDir, os.Args)),
+		cmds.NewKubectlCommand(externalCLIAction("kubectl", dataDir)),
+		cmds.NewCRICTL(externalCLIAction("crictl", dataDir)),
+		cmds.NewCtrCommand(externalCLIAction("ctr", dataDir)),
+		cmds.NewCheckConfigCommand(externalCLIAction("check-config", dataDir)),
 	}
 
 	err := app.Run(os.Args)
@@ -41,13 +43,33 @@ func main() {
 	}
 }
 
-func runCLIs() bool {
+func findDataDir() string {
+	for i, arg := range os.Args {
+		for _, flagName := range []string{"--data-dir", "-d"} {
+			if flagName == arg {
+				if len(os.Args) > i+1 {
+					return os.Args[i+1]
+				}
+			} else if strings.HasPrefix(arg, flagName+"=") {
+				return arg[len(flagName)+1:]
+			}
+		}
+	}
+	dataDir := configfilearg.MustFindString(os.Args, "data-dir")
+	if dataDir == "" {
+		dataDir = datadir.DefaultDataDir
+		logrus.Debug("Using default data dir in self-extracting wrapper")
+	}
+	return dataDir
+}
+
+func runCLIs(dataDir string) bool {
 	if os.Getenv("CRI_CONFIG_FILE") == "" {
-		os.Setenv("CRI_CONFIG_FILE", datadir.DefaultDataDir+"/agent/etc/crictl.yaml")
+		os.Setenv("CRI_CONFIG_FILE", dataDir+"/agent/etc/crictl.yaml")
 	}
 	for _, cmd := range []string{"kubectl", "ctr", "crictl"} {
 		if filepath.Base(os.Args[0]) == cmd {
-			if err := externalCLI(cmd, "", os.Args[1:]); err != nil {
+			if err := externalCLI(cmd, dataDir, os.Args[1:]); err != nil {
 				logrus.Fatal(err)
 			}
 			return true
@@ -56,9 +78,9 @@ func runCLIs() bool {
 	return false
 }
 
-func externalCLIAction(cmd string) func(cli *cli.Context) error {
+func externalCLIAction(cmd, dataDir string) func(cli *cli.Context) error {
 	return func(cli *cli.Context) error {
-		return externalCLI(cmd, cli.String("data-dir"), cli.Args())
+		return externalCLI(cmd, dataDir, cli.Args())
 	}
 }
 
@@ -70,14 +92,14 @@ func externalCLI(cli, dataDir string, args []string) error {
 	return stageAndRun(dataDir, cli, append([]string{cli}, args...))
 }
 
-func wrap(cmd string, args []string) func(ctx *cli.Context) error {
+func wrap(cmd string, dataDir string, args []string) func(ctx *cli.Context) error {
 	return func(ctx *cli.Context) error {
-		return stageAndRunCLI(ctx, cmd, args)
+		return stageAndRunCLI(ctx, cmd, dataDir, args)
 	}
 }
 
-func stageAndRunCLI(cli *cli.Context, cmd string, args []string) error {
-	dataDir, err := datadir.Resolve(cli.String("data-dir"))
+func stageAndRunCLI(cli *cli.Context, cmd string, dataDir string, args []string) error {
+	dataDir, err := datadir.Resolve(dataDir)
 	if err != nil {
 		return err
 	}

pkg/cli/cmds/root.go

@@ -38,6 +38,10 @@ func NewApp() *cli.App {
 	}
 	app.Flags = []cli.Flag{
 		DebugFlag,
+		cli.StringFlag{
+			Name:  "data-dir,d",
+			Usage: "(data) Folder to hold state default /var/lib/rancher/" + version.Program + " or ${HOME}/.rancher/" + version.Program + " if not root",
+		},
 	}
 	app.Before = SetupDebug(nil)
 

pkg/configfilearg/defaultparser.go

@@ -18,3 +18,17 @@ func MustParse(args []string) []string {
 	}
 	return result
 }
+
+func MustFindString(args []string, target string) string {
+	parser := &Parser{
+		After:         []string{},
+		FlagNames:     []string{},
+		EnvName:       version.ProgramUpper + "_CONFIG_FILE",
+		DefaultConfig: "/etc/rancher/" + version.Program + "/config.yaml",
+	}
+	result, err := parser.FindString(args, target)
+	if err != nil {
+		logrus.Fatal(err)
+	}
+	return result
+}

pkg/configfilearg/parser.go

@@ -46,6 +46,32 @@ func (p *Parser) Parse(args []string) ([]string, error) {
 	return args, nil
 }
 
+func (p *Parser) FindString(args []string, target string) (string, error) {
+	configFile, isSet := p.findConfigFileFlag(args)
+	if configFile != "" {
+		bytes, err := readConfigFileData(configFile)
+		if !isSet && os.IsNotExist(err) {
+			return "", nil
+		} else if err != nil {
+			return "", err
+		}
+
+		data := yaml.MapSlice{}
+		if err := yaml.Unmarshal(bytes, &data); err != nil {
+			return "", err
+		}
+
+		for _, i := range data {
+			k, v := convert.ToString(i.Key), convert.ToString(i.Value)
+			if k == target {
+				return v, nil
+			}
+		}
+	}
+
+	return "", nil
+}
+
 func (p *Parser) findConfigFileFlag(args []string) (string, bool) {
 	if envVal := os.Getenv(p.EnvName); p.EnvName != "" && envVal != "" {
 		return envVal, true