Fix multiple issues with CLI wrapper data-dir handling

cmd/k3s/main.go

@@ -17,25 +17,32 @@ import (
 	"github.com/rancher/k3s/pkg/flock"
 	"github.com/rancher/k3s/pkg/untar"
 	"github.com/rancher/k3s/pkg/version"
+	"github.com/rancher/wrangler/pkg/resolvehome"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli"
 )
 
+var criDefaultConfigPath = "/etc/crictl.yaml"
+
+// main entrypoint for the k3s multicall binary
 func main() {
 	dataDir := findDataDir()
+
+	// Handle direct invocation via symlink alias (multicall binary behavior)
 	if runCLIs(dataDir) {
 		return
 	}
 
+	// Handle subcommand invocation (k3s server, k3s crictl, etc)
 	app := cmds.NewApp()
 	app.Commands = []cli.Command{
-		cmds.NewServerCommand(wrap(version.Program+"-server", dataDir, os.Args)),
-		cmds.NewAgentCommand(wrap(version.Program+"-agent", dataDir, os.Args)),
+		cmds.NewServerCommand(internalCLIAction(version.Program+"-server", dataDir, os.Args)),
+		cmds.NewAgentCommand(internalCLIAction(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)),
-		cmds.NewEtcdSnapshotCommand(wrap(version.Program+"-"+cmds.EtcdSnapshotCommand, dataDir, os.Args)),
+		cmds.NewEtcdSnapshotCommand(internalCLIAction(version.Program+"-"+cmds.EtcdSnapshotCommand, dataDir, os.Args)),
 	}
 
 	if err := app.Run(os.Args); err != nil {
@@ -43,6 +50,9 @@ func main() {
 	}
 }
 
+// findDataDir reads data-dir settings from the CLI args and config file.
+// If not found, the default will be used, which varies depending on whether
+// k3s is being run as root or not.
 func findDataDir() string {
 	for i, arg := range os.Args {
 		for _, flagName := range []string{"--data-dir", "-d"} {
@@ -56,61 +66,61 @@ func findDataDir() string {
 		}
 	}
 	dataDir := configfilearg.MustFindString(os.Args, "data-dir")
-	if dataDir == "" {
-		if os.Getuid() == 0 {
-			dataDir = datadir.DefaultDataDir
-		} else {
-			dataDir = datadir.DefaultHomeDataDir
-		}
-		logrus.Debug("Using default data dir in self-extracting wrapper")
+	if d, err := datadir.Resolve(dataDir); err == nil {
+		dataDir = d
+	} else {
+		logrus.Warnf("Failed to resolve user home directory: %s", err)
 	}
 	return dataDir
 }
 
+// runCLIs handles the case where the binary is being executed as a symlink alias,
+// /usr/local/bin/crictl for example. If the executable name is one of the external
+// binaries, it calls it directly and returns true. If it's not an external binary,
+// it returns false so that standard CLI wrapping can occur.
 func runCLIs(dataDir string) bool {
-	if os.Getenv("CRI_CONFIG_FILE") == "" {
-		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, dataDir, os.Args[1:]); err != nil {
-				logrus.Fatal(err)
-			}
-			return true
+	progName := filepath.Base(os.Args[0])
+	switch progName {
+	case "crictl", "ctr", "kubectl":
+		if err := externalCLI(progName, dataDir, os.Args[1:]); err != nil {
+			logrus.Fatal(err)
 		}
+		return true
 	}
 	return false
 }
 
+// externalCLIAction returns a function that will call an external binary, be used as the Action of a cli.Command.
 func externalCLIAction(cmd, dataDir string) func(cli *cli.Context) error {
 	return func(cli *cli.Context) error {
 		return externalCLI(cmd, dataDir, cli.Args())
 	}
 }
 
+// externalCLI calls an external binary, fixing up argv[0] to the correct name.
+// crictl needs extra help to find its config file so we do that here too.
 func externalCLI(cli, dataDir string, args []string) error {
-	dataDir, err := datadir.Resolve(dataDir)
-	if err != nil {
-		return err
+	if cli == "crictl" {
+		if os.Getenv("CRI_CONFIG_FILE") == "" {
+			os.Setenv("CRI_CONFIG_FILE", findCriConfig(dataDir))
+		}
 	}
 	return stageAndRun(dataDir, cli, append([]string{cli}, args...))
 }
 
-func wrap(cmd, dataDir string, args []string) func(ctx *cli.Context) error {
+// internalCLIAction returns a function that will call a K3s internal command, be used as the Action of a cli.Command.
+func internalCLIAction(cmd, dataDir string, args []string) func(ctx *cli.Context) error {
 	return func(ctx *cli.Context) error {
 		return stageAndRunCLI(ctx, cmd, dataDir, args)
 	}
 }
 
+// stageAndRunCLI calls an external binary.
 func stageAndRunCLI(cli *cli.Context, cmd string, dataDir string, args []string) error {
-	dataDir, err := datadir.Resolve(dataDir)
-	if err != nil {
-		return err
-	}
-
 	return stageAndRun(dataDir, cmd, args)
 }
 
+// stageAndRun does the actual work of setting up and calling an external binary.
 func stageAndRun(dataDir, cmd string, args []string) error {
 	dir, err := extract(dataDir)
 	if err != nil {
@@ -135,22 +145,26 @@ func stageAndRun(dataDir, cmd string, args []string) error {
 	return syscall.Exec(cmd, args, os.Environ())
 }
 
+// getAssetAndDir returns the name of the bindata asset, along with a directory path
+// derived from the data-dir and bindata asset name.
 func getAssetAndDir(dataDir string) (string, string) {
 	asset := data.AssetNames()[0]
 	dir := filepath.Join(dataDir, "data", strings.SplitN(filepath.Base(asset), ".", 2)[0])
 	return asset, dir
 }
 
+// extract checks for and if necessary unpacks the bindata archive, returning the unique path
+// to the extracted bindata asset.
 func extract(dataDir string) (string, error) {
 	// first look for global asset folder so we don't create a HOME version if not needed
 	_, dir := getAssetAndDir(datadir.DefaultDataDir)
-	if _, err := os.Stat(dir); err == nil {
+	if _, err := os.Stat(filepath.Join(dir, "bin", "containerd")); err == nil {
 		return dir, nil
 	}
 
 	asset, dir := getAssetAndDir(dataDir)
-	// check if target directory already exists
-	if _, err := os.Stat(dir); err == nil {
+	// check if target content already exists
+	if _, err := os.Stat(filepath.Join(dir, "bin", "containerd")); err == nil {
 		return dir, nil
 	}
 
@@ -200,3 +214,35 @@ func extract(dataDir string) (string, error) {
 	}
 	return dir, os.Rename(tempDest, dir)
 }
+
+// findCriConfig returns the path to crictl.yaml
+// crictl won't search multiple locations for a config file. It will fall back to looking in
+// the same directory as the crictl binary, but that's it. We need to check the various possible
+// data-dir locations ourselves and then point it at the right one. We check:
+// - the configured data-dir
+// - the default user data-dir (assuming we can find the user's home directory)
+// - the default system data-dir
+// - the default path from upstream crictl
+func findCriConfig(dataDir string) string {
+	searchList := []string{filepath.Join(dataDir, "agent", criDefaultConfigPath)}
+
+	if homeDataDir, err := resolvehome.Resolve(datadir.DefaultHomeDataDir); err == nil {
+		searchList = append(searchList, filepath.Join(homeDataDir, "agent", criDefaultConfigPath))
+	} else {
+		logrus.Warnf("Failed to resolve user home directory: %s", err)
+	}
+
+	searchList = append(searchList, filepath.Join(datadir.DefaultDataDir, "agent", criDefaultConfigPath))
+	searchList = append(searchList, criDefaultConfigPath)
+
+	for _, path := range searchList {
+		_, err := os.Stat(path)
+		if err == nil {
+			return path
+		}
+		if !errors.Is(err, os.ErrNotExist) {
+			logrus.Warnf("Failed to %s", err)
+		}
+	}
+	return ""
+}