Add --start and --token-env flags to install command

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

Author: s3rj1k

cmd/controller/controller.go

@@ -182,10 +182,31 @@ func (c *command) start(ctx context.Context, flags *config.ControllerOptions, de
 
 	var joinClient *token.JoinClient
 
-	if (c.TokenArg != "" || c.TokenFile != "") && c.needToJoin(nodeConfig) {
+	if (c.TokenArg != "" || c.TokenFile != "" || c.TokenEnv != "") && c.needToJoin(nodeConfig) {
+		tokenSources := 0
+		if c.TokenArg != "" {
+			tokenSources++
+		}
+		if c.TokenFile != "" {
+			tokenSources++
+		}
+		if c.TokenEnv != "" {
+			tokenSources++
+		}
+
+		if tokenSources > 1 {
+			return errors.New("you can only pass one token source: either as a CLI argument 'k0s controller [token]', via '--token-file [path]', or via '--token-env [var]'")
+		}
+
 		var tokenData string
 		if c.TokenArg != "" {
 			tokenData = c.TokenArg
+		} else if c.TokenEnv != "" {
+			tokenValue := os.Getenv(c.TokenEnv)
+			if tokenValue == "" {
+				return fmt.Errorf("environment variable %q is not set or is empty", c.TokenEnv)
+			}
+			tokenData = tokenValue
 		} else {
 			data, err := os.ReadFile(c.TokenFile)
 			if err != nil {

cmd/controller/controller_test.go

@@ -75,6 +75,7 @@ Flags:
       --single                                         enable single node (implies --enable-worker, default false)
       --status-socket string                           Full file path to the socket file. (default: <rundir>/status.sock)
       --taints strings                                 Node taints, list of key=value:effect strings
+      --token-env string                               Environment variable name containing the join-token.
       --token-file string                              Path to the file containing join-token.
   -v, --verbose                                        Verbose logging (default true)
 `, out.String())

cmd/install/controller.go

@@ -49,11 +49,21 @@ With the controller subcommand you can setup a single node cluster by running:
 				return fmt.Errorf("invalid node config: %w", errors.Join(errs...))
 			}
 
+			// Convert --token-env to --token-file
+			tokenFilePath, err := handleTokenEnv(cmd, k0sVars.DataDir)
+			if err != nil {
+				return err
+			}
+
 			flagsAndVals, err := cmdFlagsToArgs(cmd)
 			if err != nil {
 				return err
 			}
 
+			if tokenFilePath != "" {
+				flagsAndVals = append(flagsAndVals, "--token-file="+tokenFilePath)
+			}
+
 			systemUsers := nodeConfig.Spec.Install.SystemUsers
 			homeDir := k0sVars.DataDir
 			if err := install.EnsureControllerUsers(systemUsers, homeDir); err != nil {
@@ -65,6 +75,12 @@ With the controller subcommand you can setup a single node cluster by running:
 				return fmt.Errorf("failed to install controller service: %w", err)
 			}
 
+			if installFlags.start {
+				if err := startInstalledService(installFlags.force); err != nil {
+					return fmt.Errorf("failed to start controller service: %w", err)
+				}
+			}
+
 			return nil
 		},
 	}

cmd/install/controller_test.go

@@ -68,13 +68,15 @@ Flags:
       --single                                         enable single node (implies --enable-worker, default false)
       --status-socket string                           Full file path to the socket file. (default: <rundir>/status.sock)
       --taints strings                                 Node taints, list of key=value:effect strings
+      --token-env string                               Environment variable name containing the join-token.
       --token-file string                              Path to the file containing join-token.
 
 Global Flags:
   -d, --debug                  Debug logging (implies verbose logging)
       --debugListenOn string   Http listenOn for Debug pprof handler (default ":6060")
   -e, --env stringArray        set environment variable
       --force                  force init script creation
+      --start                  start the service immediately after installation
   -v, --verbose                Verbose logging
 `, out.String())
 }

cmd/install/install.go

@@ -6,13 +6,16 @@ package install
 import (
 	"github.com/k0sproject/k0s/cmd/internal"
 	"github.com/k0sproject/k0s/pkg/config"
+	"github.com/k0sproject/k0s/pkg/install"
+
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 )
 
 type installFlags struct {
 	force   bool
 	envVars []string
+	start   bool
 }
 
 func NewInstallCmd() *cobra.Command {
@@ -38,9 +41,15 @@ func NewInstallCmd() *cobra.Command {
 	})
 	pflags.BoolVar(&installFlags.force, "force", false, "force init script creation")
 	pflags.StringArrayVarP(&installFlags.envVars, "env", "e", nil, "set environment variable")
+	pflags.BoolVar(&installFlags.start, "start", false, "start the service immediately after installation")
 
 	cmd.AddCommand(installWorkerCmd(&installFlags))
 	addPlatformSpecificCommands(cmd, &installFlags)
 
 	return cmd
 }
+
+// startInstalledService starts (or restarts with force) the installed k0s service.
+func startInstalledService(force bool) error {
+	return install.StartInstalledService(force)
+}

cmd/install/util.go

@@ -6,6 +6,7 @@ package install
 import (
 	"errors"
 	"fmt"
+	"os"
 	"path/filepath"
 	"strings"
 
@@ -24,7 +25,7 @@ func cmdFlagsToArgs(cmd *cobra.Command) ([]string, error) {
 			flagsAndVals = append(flagsAndVals, fmt.Sprintf(`--%s=%s`, f.Name, strings.Trim(val, "[]")))
 		default:
 			switch f.Name {
-			case "env", "force":
+			case "env", "force", "start", "token-env":
 				return
 			case "data-dir", "kubelet-root-dir", "token-file", "config":
 				if absVal, err := filepath.Abs(val); err != nil {
@@ -44,3 +45,28 @@ func cmdFlagsToArgs(cmd *cobra.Command) ([]string, error) {
 
 	return flagsAndVals, nil
 }
+
+// handleTokenEnv converts --token-env to a token file and returns its path.
+func handleTokenEnv(cmd *cobra.Command, dataDir string) (string, error) {
+	tokenEnvFlag := cmd.Flags().Lookup("token-env")
+	if tokenEnvFlag == nil || !tokenEnvFlag.Changed {
+		return "", nil
+	}
+
+	envVarName := tokenEnvFlag.Value.String()
+	tokenValue := os.Getenv(envVarName)
+	if tokenValue == "" {
+		return "", fmt.Errorf("environment variable %q is not set or is empty", envVarName)
+	}
+
+	tokenFilePath := filepath.Join(dataDir, ".token")
+	if err := os.MkdirAll(dataDir, 0755); err != nil {
+		return "", fmt.Errorf("failed to create data directory: %w", err)
+	}
+
+	if err := os.WriteFile(tokenFilePath, []byte(tokenValue), 0600); err != nil {
+		return "", fmt.Errorf("failed to write token file: %w", err)
+	}
+
+	return tokenFilePath, nil
+}

cmd/install/worker.go

@@ -27,16 +27,37 @@ All default values of worker command will be passed to the service stub unless o
 				return errors.New("this command must be run as root")
 			}
 
+			k0sVars, err := config.NewCfgVars(cmd)
+			if err != nil {
+				return fmt.Errorf("failed to initialize configuration variables: %w", err)
+			}
+
+			// Convert --token-env to --token-file
+			tokenFilePath, err := handleTokenEnv(cmd, k0sVars.DataDir)
+			if err != nil {
+				return err
+			}
+
 			flagsAndVals, err := cmdFlagsToArgs(cmd)
 			if err != nil {
 				return err
 			}
 
+			if tokenFilePath != "" {
+				flagsAndVals = append(flagsAndVals, "--token-file="+tokenFilePath)
+			}
+
 			args := append([]string{"worker"}, flagsAndVals...)
 			if err := install.InstallService(args, installFlags.envVars, installFlags.force); err != nil {
 				return fmt.Errorf("failed to install worker service: %w", err)
 			}
 
+			if installFlags.start {
+				if err := startInstalledService(installFlags.force); err != nil {
+					return fmt.Errorf("failed to start worker service: %w", err)
+				}
+			}
+
 			return nil
 		},
 	}

cmd/worker/worker.go

@@ -80,7 +80,7 @@ func NewWorkerCmd() *cobra.Command {
 				c.TokenArg = args[0]
 			}
 
-			getBootstrapKubeconfig, err := kubeconfigGetterFromJoinToken(c.TokenFile, c.TokenArg)
+			getBootstrapKubeconfig, err := kubeconfigGetterFromJoinToken(c.TokenFile, c.TokenEnv, c.TokenArg)
 			if err != nil {
 				return err
 			}
@@ -150,12 +150,23 @@ func GetNodeName(opts *config.WorkerOptions) (apitypes.NodeName, stringmap.Strin
 	return nodeName, kubeletExtraArgs, nil
 }
 
-func kubeconfigGetterFromJoinToken(tokenFile, tokenArg string) (clientcmd.KubeconfigGetter, error) {
+func kubeconfigGetterFromJoinToken(tokenFile, tokenEnv, tokenArg string) (clientcmd.KubeconfigGetter, error) {
+	tokenSources := 0
 	if tokenArg != "" {
-		if tokenFile != "" {
-			return nil, errors.New("you can only pass one token argument either as a CLI argument 'k0s worker [token]' or as a flag 'k0s worker --token-file [path]'")
-		}
+		tokenSources++
+	}
+	if tokenFile != "" {
+		tokenSources++
+	}
+	if tokenEnv != "" {
+		tokenSources++
+	}
 
+	if tokenSources > 1 {
+		return nil, errors.New("you can only pass one token source: either as a CLI argument 'k0s worker [token]', via '--token-file [path]', or via '--token-env [var]'")
+	}
+
+	if tokenArg != "" {
 		kubeconfig, err := loadKubeconfigFromJoinToken(tokenArg)
 		if err != nil {
 			return nil, err
@@ -166,12 +177,28 @@ func kubeconfigGetterFromJoinToken(tokenFile, tokenArg string) (clientcmd.Kubeco
 		}, nil
 	}
 
+	if tokenEnv != "" {
+		tokenValue := os.Getenv(tokenEnv)
+		if tokenValue == "" {
+			return nil, fmt.Errorf("environment variable %q is not set or is empty", tokenEnv)
+		}
+
+		kubeconfig, err := loadKubeconfigFromJoinToken(tokenValue)
+		if err != nil {
+			return nil, err
+		}
+
+		return func() (*clientcmdapi.Config, error) {
+			return kubeconfig, nil
+		}, nil
+	}
+
 	if tokenFile == "" {
 		return nil, nil
 	}
 
 	return func() (*clientcmdapi.Config, error) {
-		return loadKubeconfigFromTokenFile(tokenFile)
+		return loadKubeconfigFromToken(tokenFile)
 	}, nil
 }
 
@@ -193,7 +220,7 @@ func loadKubeconfigFromJoinToken(tokenData string) (*clientcmdapi.Config, error)
 	return kubeconfig, nil
 }
 
-func loadKubeconfigFromTokenFile(path string) (*clientcmdapi.Config, error) {
+func loadKubeconfigFromToken(path string) (*clientcmdapi.Config, error) {
 	var problem string
 	tokenBytes, err := os.ReadFile(path)
 	if errors.Is(err, os.ErrNotExist) {

pkg/config/cli.go

@@ -70,6 +70,7 @@ type WorkerOptions struct {
 	Labels           map[string]string
 	Taints           []string
 	TokenFile        string
+	TokenEnv         string
 	TokenArg         string
 	WorkerProfile    string
 	IPTablesMode     string
@@ -251,6 +252,7 @@ func GetWorkerFlags() *pflag.FlagSet {
 	flagset.StringVar(&workerOpts.WorkerProfile, "profile", defaultWorkerProfile, "worker profile to use on the node")
 	flagset.BoolVar(&workerOpts.CloudProvider, "enable-cloud-provider", false, "Whether or not to enable cloud provider support in kubelet")
 	flagset.StringVar(&workerOpts.TokenFile, "token-file", "", "Path to the file containing join-token.")
+	flagset.StringVar(&workerOpts.TokenEnv, "token-env", "", "Environment variable name containing the join-token.")
 	flagset.VarP((*logLevelsFlag)(&workerOpts.LogLevels), "logging", "l", "Logging Levels for the different components")
 	flagset.Var((*cliflag.ConfigurationMap)(&workerOpts.Labels), "labels", "Node labels, list of key=value pairs")
 	flagset.StringSliceVarP(&workerOpts.Taints, "taints", "", []string{}, "Node taints, list of key=value:effect strings")

pkg/install/service.go

@@ -5,6 +5,7 @@ package install
 
 import (
 	"errors"
+	"fmt"
 	"runtime"
 
 	"github.com/kardianos/service"
@@ -122,3 +123,22 @@ func GetServiceConfig(role string) *service.Config {
 		Description: k0sDescription,
 	}
 }
+
+// StartInstalledService starts (or restarts with force) the installed k0s service.
+func StartInstalledService(force bool) error {
+	svc, err := InstalledService()
+	if err != nil {
+		return err
+	}
+	status, _ := svc.Status()
+	if status == service.StatusRunning {
+		if force {
+			if err := svc.Restart(); err != nil {
+				return fmt.Errorf("failed to restart service: %w", err)
+			}
+			return nil
+		}
+		return errors.New("already running")
+	}
+	return svc.Start()
+}