From 9111b1f77e620b35ed4deb7491b79ea0deceae9f Mon Sep 17 00:00:00 2001
From: Brad Davidson <brad.davidson@rancher.com>
Date: Mon, 29 Jul 2024 19:22:50 +0000
Subject: [PATCH] Add K3S_DATA_DIR as env var for --data-dir flag

Signed-off-by: Brad Davidson <brad.davidson@rancher.com>
---
 cmd/k3s/main.go        | 7 +++++--
 pkg/cli/cmds/agent.go  | 3 +++
 pkg/cli/cmds/certs.go  | 6 +-----
 pkg/cli/cmds/root.go   | 5 +----
 pkg/cli/cmds/server.go | 1 +
 5 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/cmd/k3s/main.go b/cmd/k3s/main.go
index df2f0a4c61..6a17998505 100644
--- a/cmd/k3s/main.go
+++ b/cmd/k3s/main.go
@@ -86,11 +86,14 @@ func main() {
 	}
 }
 
-// findDataDir reads data-dir settings from the CLI args and config file.
+// findDataDir reads data-dir settings from the environment, 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(args []string) string {
-	var dataDir string
+	dataDir := os.Getenv(version.ProgramUpper + "_DATA_DIR")
+	if dataDir != "" {
+		return dataDir
+	}
 	fs := pflag.NewFlagSet("data-dir-set", pflag.ContinueOnError)
 	fs.ParseErrorsWhitelist.UnknownFlags = true
 	fs.SetOutput(io.Discard)
diff --git a/pkg/cli/cmds/agent.go b/pkg/cli/cmds/agent.go
index 16e0a196c1..751c6c31b1 100644
--- a/pkg/cli/cmds/agent.go
+++ b/pkg/cli/cmds/agent.go
@@ -266,11 +266,14 @@ func NewAgentCommand(action func(ctx *cli.Context) error) cli.Command {
 				EnvVar:      version.ProgramUpper + "_URL",
 				Destination: &AgentConfig.ServerURL,
 			},
+			// Note that this is different from DataDirFlag used elswhere in the CLI,
+			// as this is bound to AgentConfig instead of ServerConfig.
 			&cli.StringFlag{
 				Name:        "data-dir,d",
 				Usage:       "(agent/data) Folder to hold state",
 				Destination: &AgentConfig.DataDir,
 				Value:       "/var/lib/rancher/" + version.Program + "",
+				EnvVar:      version.ProgramUpper + "_DATA_DIR",
 			},
 			NodeNameFlag,
 			WithNodeIDFlag,
diff --git a/pkg/cli/cmds/certs.go b/pkg/cli/cmds/certs.go
index 50834e69fc..0c3ee9c702 100644
--- a/pkg/cli/cmds/certs.go
+++ b/pkg/cli/cmds/certs.go
@@ -28,6 +28,7 @@ var (
 		},
 	}
 	CertRotateCACommandFlags = []cli.Flag{
+		DataDirFlag,
 		cli.StringFlag{
 			Name:        "server,s",
 			Usage:       "(cluster) Server to connect to",
@@ -35,11 +36,6 @@ var (
 			Value:       "https://127.0.0.1:6443",
 			Destination: &ServerConfig.ServerURL,
 		},
-		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",
-			Destination: &ServerConfig.DataDir,
-		},
 		cli.StringFlag{
 			Name:        "path",
 			Usage:       "Path to directory containing new CA certificates",
diff --git a/pkg/cli/cmds/root.go b/pkg/cli/cmds/root.go
index 6ab66400fd..e57133baae 100644
--- a/pkg/cli/cmds/root.go
+++ b/pkg/cli/cmds/root.go
@@ -41,10 +41,7 @@ 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)",
-		},
+		DataDirFlag,
 	}
 
 	return app
diff --git a/pkg/cli/cmds/server.go b/pkg/cli/cmds/server.go
index c7fec4f541..91dafa3099 100644
--- a/pkg/cli/cmds/server.go
+++ b/pkg/cli/cmds/server.go
@@ -117,6 +117,7 @@ var (
 		Name:        "data-dir,d",
 		Usage:       "(data) Folder to hold state default /var/lib/rancher/" + version.Program + " or ${HOME}/.rancher/" + version.Program + " if not root",
 		Destination: &ServerConfig.DataDir,
+		EnvVar:      version.ProgramUpper + "_DATA_DIR",
 	}
 	ServerToken = &cli.StringFlag{
 		Name:        "token,t",