From 8ba575e6947ac701ca0acbcb12a8f468d298d411 Mon Sep 17 00:00:00 2001 From: Ionut Nicula Date: Sat, 14 Sep 2024 17:08:27 +0300 Subject: [PATCH] Fix shorthand combination edge case in c.Find() code path Fixes: #2188 --- command.go | 68 +++++++++++++++++++++++++++++++++++++------------ command_test.go | 60 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 111 insertions(+), 17 deletions(-) diff --git a/command.go b/command.go index 2df6975f..cc20afed 100644 --- a/command.go +++ b/command.go @@ -643,13 +643,14 @@ func shortHasNoOptDefVal(name string, fs *flag.FlagSet) bool { return flag.NoOptDefVal != "" } -func stripFlags(args []string, c *Command) []string { +func stripFlags(args []string, c *Command) ([]string, []string) { if len(args) == 0 { - return args + return args, nil } c.mergePersistentFlags() commands := []string{} + flagsThatConsumeNextArg := []string{} // We use this to avoid repeating the same lengthy logic for parsing shorthand combinations in argsMinusFirstX flags := c.Flags() Loop: @@ -665,31 +666,70 @@ Loop: // delete arg from args. fallthrough // (do the same as below) case strings.HasPrefix(s, "-") && !strings.Contains(s, "=") && len(s) == 2 && !shortHasNoOptDefVal(s[1:], flags): + flagsThatConsumeNextArg = append(flagsThatConsumeNextArg, s) // If '-f arg' then // delete 'arg' from args or break the loop if len(args) <= 1. if len(args) <= 1 { break Loop } else { args = args[1:] - continue + } + case strings.HasPrefix(s, "-") && !strings.HasPrefix(s, "--") && !strings.Contains(s, "=") && len(s) > 2: + shorthandCombination := s[1:] // Skip the leading "-" + lastPos := len(shorthandCombination) - 1 + for i, shorthand := range shorthandCombination { + if shortHasNoOptDefVal(string(shorthand), flags) { + continue + } + + // We found a shorthand that needs a value. + + if i == lastPos { + // Since we're at the end of the shorthand combination, this means that the + // value for the shorthand is given in the next argument. (e.g. '-xyzf arg', + // where -x, -y, -z are boolean flags, and -f is a flag that needs a value). + + // The whole combination will take a value. + flagsThatConsumeNextArg = append(flagsThatConsumeNextArg, s) + + if len(args) <= 1 { + break Loop + } else { + args = args[1:] + } + } else { + // Since the shorthand combination doesn't end here, this means that the + // value for the shorthand is given in the same argument, meaning we don't + // have to consume the next one. (e.g. '-xyzfarg', where -x, -y, -z are + // boolean flags, and -f is a flag that needs a value). + } + + break } case s != "" && !strings.HasPrefix(s, "-"): commands = append(commands, s) } } - return commands + return commands, flagsThatConsumeNextArg } // argsMinusFirstX removes only the first x from args. Otherwise, commands that look like // openshift admin policy add-role-to-user admin my-user, lose the admin argument (arg[4]). // Special care needs to be taken not to remove a flag value. -func (c *Command) argsMinusFirstX(args []string, x string) []string { +func (c *Command) argsMinusFirstX(args, flagsThatConsumeNextArg []string, x string) []string { if len(args) == 0 { return args } - c.mergePersistentFlags() - flags := c.Flags() + + consumesNextArg := func(flag string) bool { + for _, f := range flagsThatConsumeNextArg { + if flag == f { + return true + } + } + return false + } Loop: for pos := 0; pos < len(args); pos++ { @@ -698,13 +738,8 @@ Loop: case s == "--": // -- means we have reached the end of the parseable args. Break out of the loop now. break Loop - case strings.HasPrefix(s, "--") && !strings.Contains(s, "=") && !hasNoOptDefVal(s[2:], flags): - fallthrough - case strings.HasPrefix(s, "-") && !strings.Contains(s, "=") && len(s) == 2 && !shortHasNoOptDefVal(s[1:], flags): - // This is a flag without a default value, and an equal sign is not used. Increment pos in order to skip - // over the next arg, because that is the value of this flag. + case consumesNextArg(s): pos++ - continue case !strings.HasPrefix(s, "-"): // This is not a flag or a flag value. Check to see if it matches what we're looking for, and if so, // return the args, excluding the one at this position. @@ -730,7 +765,7 @@ func (c *Command) Find(args []string) (*Command, []string, error) { var innerfind func(*Command, []string) (*Command, []string) innerfind = func(c *Command, innerArgs []string) (*Command, []string) { - argsWOflags := stripFlags(innerArgs, c) + argsWOflags, flagsThatConsumeNextArg := stripFlags(innerArgs, c) if len(argsWOflags) == 0 { return c, innerArgs } @@ -738,14 +773,15 @@ func (c *Command) Find(args []string) (*Command, []string, error) { cmd := c.findNext(nextSubCmd) if cmd != nil { - return innerfind(cmd, c.argsMinusFirstX(innerArgs, nextSubCmd)) + return innerfind(cmd, c.argsMinusFirstX(innerArgs, flagsThatConsumeNextArg, nextSubCmd)) } return c, innerArgs } commandFound, a := innerfind(c, args) if commandFound.Args == nil { - return commandFound, a, legacyArgs(commandFound, stripFlags(a, commandFound)) + argsWOflags, _ := stripFlags(a, commandFound) + return commandFound, a, legacyArgs(commandFound, argsWOflags) } return commandFound, a, nil } diff --git a/command_test.go b/command_test.go index 9ce7a529..2f0f4b16 100644 --- a/command_test.go +++ b/command_test.go @@ -693,6 +693,30 @@ func TestStripFlags(t *testing.T) { []string{"-p", "bar"}, []string{"bar"}, }, + { + []string{"-s", "value", "bar"}, + []string{"bar"}, + }, + { + []string{"-s=value", "bar"}, + []string{"bar"}, + }, + { + []string{"-svalue", "bar"}, + []string{"bar"}, + }, + { + []string{"-ps", "value", "bar"}, + []string{"bar"}, + }, + { + []string{"-ps=value", "bar"}, + []string{"bar"}, + }, + { + []string{"-psvalue", "bar"}, + []string{"bar"}, + }, } c := &Command{Use: "c", Run: emptyRun} @@ -702,7 +726,7 @@ func TestStripFlags(t *testing.T) { c.Flags().BoolP("bool", "b", false, "") for i, test := range tests { - got := stripFlags(test.input, c) + got, _ := stripFlags(test.input, c) if !reflect.DeepEqual(test.output, got) { t.Errorf("(%v) Expected: %v, got: %v", i, test.output, got) } @@ -2688,11 +2712,13 @@ func TestHelpflagCommandExecutedWithoutVersionSet(t *testing.T) { func TestFind(t *testing.T) { var foo, bar string + var persist bool root := &Command{ Use: "root", } root.PersistentFlags().StringVarP(&foo, "foo", "f", "", "") root.PersistentFlags().StringVarP(&bar, "bar", "b", "something", "") + root.PersistentFlags().BoolVarP(&persist, "persist", "p", false, "") child := &Command{ Use: "child", @@ -2755,6 +2781,38 @@ func TestFind(t *testing.T) { []string{"--foo", "child", "--bar", "something", "child"}, []string{"--foo", "child", "--bar", "something"}, }, + { + []string{"-f", "value", "child"}, + []string{"-f", "value"}, + }, + { + []string{"-f=value", "child"}, + []string{"-f=value"}, + }, + { + []string{"-fvalue", "child"}, + []string{"-fvalue"}, + }, + { + []string{"-pf", "value", "child"}, + []string{"-pf", "value"}, + }, + { + []string{"-pf=value", "child"}, + []string{"-pf=value"}, + }, + { + []string{"-pfvalue", "child"}, + []string{"-pfvalue"}, + }, + { + []string{"-pf", "child", "child"}, + []string{"-pf", "child"}, + }, + { + []string{"-pf", "child", "-pb", "something", "child"}, + []string{"-pf", "child", "-pb", "something"}, + }, } for _, tc := range testCases {