Reject unknown sub-commands

Cobra has the property of not reporting any unrecognized sub-commands and simply displaying the help page for any command that doesn't have a run function set. Add a unit test that probes all sub-commands to properly reject unknown sub-commands. Whenever a sub-command is added that accepts arguments and doesn't declare them as ValidArgs, it must be added to the ignore list at the beginning of that test. Signed-off-by: Tom Wieczorek <[email protected]>
k0sproject · Dec 20, 2024 · ca4856b · ca4856b
1 parent d8739ee
commit ca4856b
Show file tree

Hide file tree

Showing 39 changed files with 248 additions and 96 deletions.
diff --git a/cmd/airgap/airgap.go b/cmd/airgap/airgap.go
@@ -26,6 +26,8 @@ func NewAirgapCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "airgap",
 		Short: "Manage airgap setup",
+		Args:  cobra.NoArgs,
+		Run:   func(*cobra.Command, []string) { /* Enforce arg validation. */ },
 	}
 
 	cmd.AddCommand(NewAirgapListImagesCmd())

diff --git a/cmd/airgap/listimages.go b/cmd/airgap/listimages.go
@@ -32,7 +32,8 @@ func NewAirgapListImagesCmd() *cobra.Command {
 		Use:     "list-images",
 		Short:   "List image names and version needed for air-gap install",
 		Example: `k0s airgap list-images`,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		Args:    cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			opts, err := config.GetCmdOpts(cmd)
 			if err != nil {
 				return err

diff --git a/cmd/airgap/listimages_test.go b/cmd/airgap/listimages_test.go
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-package airgap
+package airgap_test
 
 import (
 	"errors"
@@ -26,6 +26,7 @@ import (
 	"testing"
 	"testing/iotest"
 
+	"github.com/k0sproject/k0s/cmd"
 	internalio "github.com/k0sproject/k0s/internal/io"
 	"github.com/k0sproject/k0s/pkg/apis/k0s/v1beta1"
 
@@ -46,15 +47,12 @@ func TestAirgapListImages(t *testing.T) {
 
 	t.Run("HonorsIOErrors", func(t *testing.T) {
 		var writes uint
-		underTest := NewAirgapListImagesCmd()
-		underTest.SetIn(iotest.ErrReader(errors.New("unexpected read from standard input")))
+		underTest, _, stderr := newAirgapListImagesCmdWithConfig(t, "")
 		underTest.SilenceUsage = true // Cobra writes usage to stdout on errors 🤔
 		underTest.SetOut(internalio.WriterFunc(func(p []byte) (int, error) {
 			writes++
 			return 0, assert.AnError
 		}))
-		var stderr strings.Builder
-		underTest.SetErr(&stderr)
 
 		assert.Same(t, assert.AnError, underTest.Execute())
 		assert.Equal(t, uint(1), writes, "Expected a single write to stdout")
@@ -127,8 +125,8 @@ func newAirgapListImagesCmdWithConfig(t *testing.T, config string, args ...strin
 	require.NoError(t, os.WriteFile(configFile, []byte(config), 0644))
 
 	out, err = new(strings.Builder), new(strings.Builder)
-	cmd := NewAirgapListImagesCmd()
-	cmd.SetArgs(append([]string{"--config=" + configFile}, args...))
+	cmd := cmd.NewRootCmd()
+	cmd.SetArgs(append([]string{"airgap", "--config=" + configFile, "list-images"}, args...))
 	cmd.SetIn(iotest.ErrReader(errors.New("unexpected read from standard input")))
 	cmd.SetOut(out)
 	cmd.SetErr(err)

diff --git a/cmd/api/api.go b/cmd/api/api.go
@@ -53,12 +53,13 @@ func NewAPICmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "api",
 		Short: "Run the controller API",
+		Args:  cobra.NoArgs,
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
 			logrus.SetOutput(cmd.OutOrStdout())
 			k0slog.SetInfoLevel()
 			return config.CallParentPersistentPreRun(cmd, args)
 		},
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			opts, err := config.GetCmdOpts(cmd)
 			if err != nil {
 				return err

diff --git a/cmd/backup/backup_unix.go b/cmd/backup/backup_unix.go
@@ -42,11 +42,12 @@ func NewBackupCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "backup",
 		Short: "Back-Up k0s configuration. Must be run as root (or with sudo)",
+		Args:  cobra.NoArgs,
 		PreRun: func(cmd *cobra.Command, args []string) {
 			// ensure logs don't mess up output
 			logrus.SetOutput(cmd.ErrOrStderr())
 		},
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			opts, err := config.GetCmdOpts(cmd)
 			if err != nil {
 				return err
@@ -69,7 +70,7 @@ func NewBackupCmd() *cobra.Command {
 
 func (c *command) backup(savePath string, out io.Writer) error {
 	if os.Geteuid() != 0 {
-		logrus.Fatal("this command must be run as root!")
+		return errors.New("this command must be run as root!")
 	}
 
 	if savePath != "-" && !dir.IsDirectory(savePath) {

diff --git a/cmd/backup/backup_windows.go b/cmd/backup/backup_windows.go
@@ -28,6 +28,7 @@ func NewBackupCmd() *cobra.Command {
 	return &cobra.Command{
 		Use:   "backup",
 		Short: "Back-Up k0s configuration. Not supported on Windows OS",
+		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			return errors.New("unsupported Operating System for this command")
 		},

diff --git a/cmd/config/config.go b/cmd/config/config.go
@@ -24,10 +24,29 @@ func NewConfigCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "config",
 		Short: "Configuration related sub-commands",
+		Args:  cobra.NoArgs,
+		Run:   func(*cobra.Command, []string) { /* Enforce arg validation. */ },
 	}
 	cmd.AddCommand(NewCreateCmd())
 	cmd.AddCommand(NewEditCmd())
 	cmd.AddCommand(NewStatusCmd())
 	cmd.AddCommand(NewValidateCmd())
 	return cmd
 }
+
+func reExecKubectl(cmd *cobra.Command, kubectlArgs ...string) error {
+	args := []string{"kubectl"}
+	if dataDir, err := cmd.Flags().GetString("data-dir"); err != nil {
+		return err
+	} else if dataDir != "" {
+		args = append(args, "--data-dir", dataDir)
+	}
+
+	root := cmd.Root()
+	root.SetArgs(append(args, kubectlArgs...))
+
+	silenceErrors := root.SilenceErrors
+	defer func() { root.SilenceErrors = silenceErrors }()
+	root.SilenceErrors = true // So that errors aren't printed twice.
+	return root.Execute()
+}
diff --git a/cmd/config/config_test.go b/cmd/config/config_test.go
@@ -0,0 +1,64 @@
+/*
+Copyright 2024 k0s authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package config_test
+
+import (
+	"errors"
+	"fmt"
+	"slices"
+	"strings"
+	"testing"
+	"testing/iotest"
+
+	"github.com/k0sproject/k0s/cmd"
+	"github.com/spf13/pflag"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestConfigCmd_RejectsUnknownCommands(t *testing.T) {
+	for _, cmds := range [][]string{
+		{"config"},
+		{"config", "create"},
+		{"config", "edit"},
+		{"config", "status"},
+		{"config", "validate"},
+	} {
+		subCommand := strings.Join(cmds, " ")
+		t.Run(subCommand, func(t *testing.T) {
+			var stdout, stderr strings.Builder
+			underTest := cmd.NewRootCmd()
+
+			// Reset any "required" annotations on flags
+			if cmd, _, err := underTest.Find(cmds); assert.NoError(t, err) {
+				flags := cmd.Flags()
+				flags.VisitAll(func(flag *pflag.Flag) {
+					flag.Annotations = nil
+				})
+			}
+
+			underTest.SetArgs(slices.Concat(cmds, []string{"bogus"}))
+			underTest.SetIn(iotest.ErrReader(errors.New("unexpected read from standard input")))
+			underTest.SetOut(&stdout)
+			underTest.SetErr(&stderr)
+
+			msg := fmt.Sprintf(`unknown command "bogus" for "k0s %s"`, subCommand)
+			assert.ErrorContains(t, underTest.Execute(), msg)
+			assert.Equal(t, "Error: "+msg+"\n", stderr.String())
+			assert.Empty(t, stdout.String())
+		})
+	}
+}
diff --git a/cmd/config/create.go b/cmd/config/create.go
@@ -32,7 +32,8 @@ func NewCreateCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "create",
 		Short: "Output the default k0s configuration yaml to stdout",
-		RunE: func(cmd *cobra.Command, args []string) error {
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			config := v1beta1.DefaultClusterConfig()
 			if !includeImages {
 				config.Spec.Images = nil

diff --git a/cmd/config/edit.go b/cmd/config/edit.go
@@ -17,8 +17,6 @@ limitations under the License.
 package config
 
 import (
-	"os"
-
 	"github.com/k0sproject/k0s/pkg/config"
 
 	"github.com/spf13/cobra"
@@ -28,13 +26,9 @@ func NewEditCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "edit",
 		Short: "Launch the editor configured in your shell to edit k0s configuration",
+		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			opts, err := config.GetCmdOpts(cmd)
-			if err != nil {
-				return err
-			}
-			os.Args = []string{os.Args[0], "kubectl", "--data-dir", opts.K0sVars.DataDir, "-n", "kube-system", "edit", "clusterconfig", "k0s"}
-			return cmd.Execute()
+			return reExecKubectl(cmd, "-n", "kube-system", "edit", "clusterconfig", "k0s")
 		},
 	}
 	cmd.PersistentFlags().AddFlagSet(config.GetKubeCtlFlagSet())

diff --git a/cmd/config/status.go b/cmd/config/status.go
@@ -17,8 +17,6 @@ limitations under the License.
 package config
 
 import (
-	"os"
-
 	"github.com/k0sproject/k0s/pkg/config"
 
 	"github.com/spf13/cobra"
@@ -30,16 +28,14 @@ func NewStatusCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "status",
 		Short: "Display dynamic configuration reconciliation status",
+		Args:  cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
-			opts, err := config.GetCmdOpts(cmd)
-			if err != nil {
-				return err
-			}
-			os.Args = []string{os.Args[0], "kubectl", "--data-dir", opts.K0sVars.DataDir, "-n", "kube-system", "get", "event", "--field-selector", "involvedObject.name=k0s"}
+			args := []string{"-n", "kube-system", "get", "event", "--field-selector", "involvedObject.name=k0s"}
 			if outputFormat != "" {
-				os.Args = append(os.Args, "-o", outputFormat)
+				args = append(args, "-o", outputFormat)
 			}
-			return cmd.Execute()
+
+			return reExecKubectl(cmd, args...)
 		},
 	}
 	cmd.PersistentFlags().AddFlagSet(config.GetKubeCtlFlagSet())

diff --git a/cmd/config/validate.go b/cmd/config/validate.go
@@ -34,7 +34,8 @@ func NewValidateCmd() *cobra.Command {
 		Short: "Validate k0s configuration",
 		Long: `Example:
    k0s config validate --config path_to_config.yaml`,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		Args: cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			var reader io.Reader
 
 			// config.CfgFile is the global value holder for --config flag, set by cobra/pflag

diff --git a/cmd/controller/controller.go b/cmd/controller/controller.go
@@ -80,6 +80,7 @@ func NewControllerCmd() *cobra.Command {
 	or CLI flag:
 	$ k0s controller --token-file [path_to_file]
 	Note: Token can be passed either as a CLI argument or as a flag`,
+		Args: cobra.MaximumNArgs(1),
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
 			logrus.SetOutput(cmd.OutOrStdout())
 			k0slog.SetInfoLevel()

diff --git a/cmd/etcd/etcd.go b/cmd/etcd/etcd.go
@@ -30,6 +30,7 @@ func NewEtcdCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "etcd",
 		Short: "Manage etcd cluster",
+		Args:  cobra.NoArgs,
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
 			if err := config.CallParentPersistentPreRun(cmd, args); err != nil {
 				return err
@@ -51,6 +52,7 @@ func NewEtcdCmd() *cobra.Command {
 			}
 			return nil
 		},
+		Run: func(*cobra.Command, []string) { /* Enforce arg validation. */ },
 	}
 	cmd.AddCommand(etcdLeaveCmd())
 	cmd.AddCommand(etcdListCmd())

diff --git a/cmd/etcd/leave.go b/cmd/etcd/leave.go
@@ -39,7 +39,7 @@ func etcdLeaveCmd() *cobra.Command {
 		Use:   "leave",
 		Short: "Leave the etcd cluster, or remove a specific peer",
 		Args:  cobra.NoArgs, // accept peer address via flag, not via arg
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			opts, err := config.GetCmdOpts(cmd)
 			if err != nil {
 				return err

diff --git a/cmd/etcd/list.go b/cmd/etcd/list.go
@@ -36,7 +36,7 @@ func etcdListCmd() *cobra.Command {
 			// ensure logs don't mess up the output
 			logrus.SetOutput(cmd.ErrOrStderr())
 		},
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			opts, err := config.GetCmdOpts(cmd)
 			if err != nil {
 				return err

diff --git a/cmd/etcd/list_test.go b/cmd/etcd/list_test.go
diff --git a/cmd/install/controller.go b/cmd/install/controller.go
@@ -36,7 +36,8 @@ With the controller subcommand you can setup a single node cluster by running:
 
 	k0s install controller --single
 	`,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		Args: cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			opts, err := config.GetCmdOpts(cmd)
 			if err != nil {
 				return err

diff --git a/cmd/install/install.go b/cmd/install/install.go
@@ -40,6 +40,8 @@ func NewInstallCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "install",
 		Short: "Install k0s on a brand-new system. Must be run as root (or with sudo)",
+		Args:  cobra.NoArgs,
+		Run:   func(*cobra.Command, []string) { /* Enforce arg validation. */ },
 	}
 
 	cmd.AddCommand(installControllerCmd(&installFlags))

diff --git a/cmd/install/worker.go b/cmd/install/worker.go
@@ -30,7 +30,8 @@ func installWorkerCmd(installFlags *installFlags) *cobra.Command {
 All default values of worker command will be passed to the service stub unless overridden.
 
 Windows flags like "--api-server", "--cidr-range" and "--cluster-dns" will be ignored since install command doesn't yet support Windows services`,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		Args: cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			opts, err := config.GetCmdOpts(cmd)
 			if err != nil {
 				return err

diff --git a/cmd/kubeconfig/admin.go b/cmd/kubeconfig/admin.go
@@ -38,6 +38,7 @@ func kubeConfigAdminCmd() *cobra.Command {
 		Example: `	$ k0s kubeconfig admin > ~/.kube/config
 	$ export KUBECONFIG=~/.kube/config
 	$ kubectl get nodes`,
+		Args: cobra.NoArgs,
 		PreRun: func(cmd *cobra.Command, args []string) {
 			// ensure logs don't mess up the output
 			logrus.SetOutput(cmd.ErrOrStderr())