From 332fd73d46e7c1a95097b6e62c3dd314865b199d Mon Sep 17 00:00:00 2001 From: Chris Kim <30601846+Oats87@users.noreply.github.com> Date: Wed, 16 Dec 2020 09:27:57 -0800 Subject: [PATCH] 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 --- cmd/k3s/main.go | 54 +++++++++++++++++++++--------- pkg/cli/cmds/root.go | 4 +++ pkg/configfilearg/defaultparser.go | 14 ++++++++ pkg/configfilearg/parser.go | 26 ++++++++++++++ 4 files changed, 82 insertions(+), 16 deletions(-) diff --git a/cmd/k3s/main.go b/cmd/k3s/main.go index 9b7f468a1e..893afe0dcd 100644 --- a/cmd/k3s/main.go +++ b/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 } diff --git a/pkg/cli/cmds/root.go b/pkg/cli/cmds/root.go index f517971651..ebd536685c 100644 --- a/pkg/cli/cmds/root.go +++ b/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) diff --git a/pkg/configfilearg/defaultparser.go b/pkg/configfilearg/defaultparser.go index 7a7f10c0fc..2eb1bf536c 100644 --- a/pkg/configfilearg/defaultparser.go +++ b/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 +} diff --git a/pkg/configfilearg/parser.go b/pkg/configfilearg/parser.go index 4533fcac8a..fade98d5ba 100644 --- a/pkg/configfilearg/parser.go +++ b/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