Merge pull request #369 from pyroscope-io/fix-exec-flags

Fix argument processing

cmd/pyroscope/command/agent.go

@@ -12,8 +12,9 @@ func newAgentCmd(cfg *config.Agent) *cobra.Command {
 	agentCmd := &cobra.Command{
 		Use:   "agent [flags]",
 		Short: "Start pyroscope agent",
-		Args:  cobra.NoArgs,
-		RunE: createCmdRunFn(cfg, vpr, func(cmd *cobra.Command, args []string) error {
+
+		DisableFlagParsing: true,
+		RunE: createCmdRunFn(cfg, vpr, func(_ *cobra.Command, _ []string) error {
 			return cli.StartAgent(cfg)
 		}),
 	}

cmd/pyroscope/command/command.go

@@ -5,14 +5,18 @@ import (
 	"os"
 	"strings"
 
-	"github.com/pyroscope-io/pyroscope/pkg/cli"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
 
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
 	"github.com/pyroscope-io/pyroscope/pkg/config"
+	"github.com/pyroscope-io/pyroscope/pkg/util/slices"
 )
 
+// https://pubs.opengroup.org/onlinepubs/9699919799/basedefs/V1_chap12.html#tag_12_02
+const optionsEnd = "--"
+
 type cmdRunFn func(cmd *cobra.Command, args []string) error
 
 func createCmdRunFn(cfg interface{}, vpr *viper.Viper, fn cmdRunFn) cmdRunFn {
@@ -29,13 +33,69 @@ func createCmdRunFn(cfg interface{}, vpr *viper.Viper, fn cmdRunFn) cmdRunFn {
 		if err = cli.Unmarshal(vpr, cfg); err != nil {
 			return err
 		}
+
+		var xargs []string
+		x := firstArgumentIndex(cmd.Flags(), prependDash(args))
+		if x >= 0 {
+			xargs = args[:x]
+			args = args[x:]
+		} else {
+			xargs = args
+			args = nil
+		}
+		if err = cmd.Flags().Parse(prependDash(xargs)); err != nil {
+			return err
+		}
+		if slices.StringContains(xargs, "--help") {
+			_ = cmd.Help()
+			return nil
+		}
+
 		if err = fn(cmd, args); err != nil {
 			cmd.SilenceUsage = true
 		}
 		return err
 	}
 }
 
+func prependDash(args []string) []string {
+	for i, arg := range args {
+		if len(arg) > 2 && strings.HasPrefix(arg, "-") && !strings.HasPrefix(arg, "--") {
+			args[i] = "-" + arg
+		}
+	}
+	return args
+}
+
+// firstArgumentIndex returns index of the first encountered argument.
+// If args does not contain arguments, or contains undefined flags,
+// the call returns -1.
+func firstArgumentIndex(flags *pflag.FlagSet, args []string) int {
+	for i := 0; i < len(args); i++ {
+		a := args[i]
+		switch {
+		default:
+			return i
+		case a == optionsEnd:
+			return i + 1
+		case strings.HasPrefix(a, optionsEnd) && len(a) > 2:
+			x := strings.SplitN(a[2:], "=", 2)
+			f := flags.Lookup(x[0])
+			if f == nil {
+				return -1
+			}
+			if f.Value.Type() == "bool" {
+				continue
+			}
+			if len(x) == 1 {
+				i++
+			}
+		}
+	}
+	// Should have returned earlier.
+	return -1
+}
+
 func newViper() *viper.Viper {
 	v := viper.New()
 	v.SetEnvPrefix("PYROSCOPE")

cmd/pyroscope/command/connect.go

@@ -13,8 +13,9 @@ func newConnectCmd(cfg *config.Exec) *cobra.Command {
 	connectCmd := &cobra.Command{
 		Use:   "connect [flags]",
 		Short: "Connect to an existing process and profile it",
-		Args:  cobra.NoArgs,
-		RunE: createCmdRunFn(cfg, vpr, func(cmd *cobra.Command, args []string) error {
+
+		DisableFlagParsing: true,
+		RunE: createCmdRunFn(cfg, vpr, func(_ *cobra.Command, args []string) error {
 			return exec.Cli(cfg, args)
 		}),
 	}

cmd/pyroscope/command/convert.go

@@ -19,8 +19,9 @@ func newConvertCmd(cfg *config.Convert) *cobra.Command {
 	convertCmd := &cobra.Command{
 		Use:    "convert [flags]",
 		Short:  "Convert between different profiling formats",
-		Args:   cobra.NoArgs,
 		Hidden: true,
+
+		DisableFlagParsing: true,
 		RunE: createCmdRunFn(cfg, vpr, func(_ *cobra.Command, _ []string) error {
 			return parse(os.Stdin, cfg.Format)
 		}),

cmd/pyroscope/command/dbmanager.go

@@ -15,7 +15,9 @@ func newDbManagerCmd(cfg *config.CombinedDbManager) *cobra.Command {
 		Short:  "Manage database",
 		Args:   cobra.ExactArgs(1), // TODO: should be implemented as subcommands.
 		Hidden: true,
-		RunE: createCmdRunFn(cfg, vpr, func(cmd *cobra.Command, args []string) error {
+
+		DisableFlagParsing: true,
+		RunE: createCmdRunFn(cfg, vpr, func(_ *cobra.Command, args []string) error {
 			return dbmanager.Cli(cfg.DbManager, cfg.Server, args)
 		}),
 	}

cmd/pyroscope/command/exec.go

@@ -17,7 +17,9 @@ func newExecCmd(cfg *config.Exec) *cobra.Command {
 		Use:   "exec [flags] <args>",
 		Short: "Start a new process from arguments and profile it",
 		Args:  cobra.MinimumNArgs(1),
-		RunE: createCmdRunFn(cfg, vpr, func(cmd *cobra.Command, args []string) error {
+
+		DisableFlagParsing: true,
+		RunE: createCmdRunFn(cfg, vpr, func(_ *cobra.Command, args []string) error {
 			err := exec.Cli(cfg, args)
 			// Normally, if the program ran, the call should return ExitError and
 			// the exit code must be preserved. Otherwise, the error originates from

cmd/pyroscope/command/exec_test.go

@@ -0,0 +1,104 @@
+package command
+
+import (
+	"errors"
+	"io"
+	"testing"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+	"github.com/spf13/cobra"
+
+	"github.com/pyroscope-io/pyroscope/pkg/cli"
+	"github.com/pyroscope-io/pyroscope/pkg/config"
+)
+
+type cmdArgsTestCase struct {
+	description  string
+	inputArgs    []string
+	expectedArgs []string
+	err          error
+}
+
+var appArgs = []string{"app", "-f", "1", "-b", "arg"}
+
+func execArgs(args ...string) []string { return append(args, appArgs...) }
+
+func TestExecCommand(t *testing.T) {
+	RegisterFailHandler(Fail)
+
+	testCases := []cmdArgsTestCase{
+		{
+			description: "no_arguments",
+			inputArgs:   []string{},
+		},
+		{
+			description: "help_flag",
+			inputArgs:   []string{"--help"},
+		},
+		{
+			description:  "delimiter",
+			inputArgs:    []string{optionsEnd},
+			expectedArgs: []string{},
+		},
+		{
+			description: "unknown_flag",
+			inputArgs:   []string{"--non-existing_flag"},
+			err:         errors.New("unknown flag: --non-existing_flag"),
+		},
+		{
+			description:  "exec_no_arguments",
+			inputArgs:    appArgs,
+			expectedArgs: appArgs,
+		},
+		{
+			description:  "exec_separated",
+			inputArgs:    execArgs("-spy-name", "debugspy", optionsEnd),
+			expectedArgs: appArgs,
+		},
+		{
+			description:  "exec_flags_mixed",
+			expectedArgs: appArgs,
+			inputArgs: execArgs(
+				"--spy-name=debugspy",
+				"--application-name", "app",
+				"--no-logging",
+				"-server-address=http://localhost:4040",
+				"-log-level", "debug",
+				"-no-logging",
+			),
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.description, func(t *testing.T) {
+			g := NewGomegaWithT(t)
+			cfg := new(config.Exec)
+			vpr := newViper()
+			var cmdArgs []string
+			cmd := &cobra.Command{
+				SilenceErrors:      true,
+				DisableFlagParsing: true,
+				RunE: createCmdRunFn(cfg, vpr, func(_ *cobra.Command, args []string) error {
+					cmdArgs = args
+					return nil
+				}),
+			}
+
+			cmd.SetUsageFunc(printUsageMessage)
+			cmd.SetHelpFunc(printHelpMessage)
+			cmd.SetArgs(testCase.inputArgs)
+			cmd.SetOut(io.Discard)
+			cli.PopulateFlagSet(cfg, cmd.Flags(), vpr)
+
+			err := cmd.Execute()
+			if testCase.err == nil {
+				g.Expect(err).To(BeNil())
+				g.Expect(cmdArgs).To(Equal(testCase.expectedArgs))
+				return
+			}
+
+			g.Expect(err).To(Equal(testCase.err))
+		})
+	}
+}

cmd/pyroscope/command/root.go

@@ -2,9 +2,7 @@ package command
 
 import (
 	"fmt"
-	"os"
 	"runtime"
-	"strings"
 
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -47,14 +45,6 @@ func Execute() error {
 		},
 	})
 
-	args := os.Args[1:]
-	for i, arg := range args {
-		if len(arg) > 2 && strings.HasPrefix(arg, "-") && !strings.HasPrefix(arg, "--") {
-			args[i] = fmt.Sprintf("-%s", arg)
-		}
-	}
-
-	rootCmd.SetArgs(args)
 	return rootCmd.Execute()
 }
 

cmd/pyroscope/command/server.go

@@ -11,9 +11,10 @@ func newServerCmd(cfg *config.Server) *cobra.Command {
 	vpr := newViper()
 	serverCmd := &cobra.Command{
 		Use:   "server [flags]",
-		Args:  cobra.NoArgs,
 		Short: "Start pyroscope server. This is the database + web-based user interface",
-		RunE: createCmdRunFn(cfg, vpr, func(cmd *cobra.Command, args []string) error {
+
+		DisableFlagParsing: true,
+		RunE: createCmdRunFn(cfg, vpr, func(_ *cobra.Command, _ []string) error {
 			return cli.StartServer(cfg)
 		}),
 	}