Skip to content

Commit

Permalink
Use command path when registering flagsets
Browse files Browse the repository at this point in the history
Avoids conflicts when different sub-commands have the same command name.
(Eg: hubble status vs hubble foobar status).

Since commands paths rely on having parent-commands, we must initialize
the map containing the flagsets after all sub-commands have been added.

Signed-off-by: Chance Zibolski <chance.zibolski@gmail.com>
  • Loading branch information
chancez authored and rolinh committed Aug 2, 2022
1 parent af4703f commit a5e8697
Show file tree
Hide file tree
Showing 12 changed files with 47 additions and 22 deletions.
34 changes: 26 additions & 8 deletions cmd/common/template/usage.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,29 +17,47 @@ package template
import (
"strings"

"github.com/cilium/hubble/cmd/common/config"
"github.com/spf13/cobra"
"github.com/spf13/pflag"

"github.com/cilium/hubble/cmd/common/config"
)

var commandFlagSets = map[string][]*pflag.FlagSet{}
var (
commandFlagSets = map[*cobra.Command][]*pflag.FlagSet{}
commandPathFlagSets = map[string][]*pflag.FlagSet{}
)

func init() {
cobra.AddTemplateFunc("title", strings.Title)
cobra.AddTemplateFunc("getFlagSets", getFlagSets)
}

// Initialize goes through the registered commands, and their flagsets and
// initializes the help template command registry.
//
// This must be called after all commands are added as sub-commands, because
// cmd.CommandPath relies on the commands having parents.
func Initialize() {
for cmd, fs := range commandFlagSets {
commandPathFlagSets[cmd.CommandPath()] = fs
}
}

// RegisterFlagSets registers flags to be included in a commands usage text (--help).
func RegisterFlagSets(cmdName string, flagsets ...*pflag.FlagSet) {
commandFlagSets[cmdName] = append(commandFlagSets[cmdName], flagsets...)
func RegisterFlagSets(cmd *cobra.Command, flagsets ...*pflag.FlagSet) {
commandFlagSets[cmd] = append(commandFlagSets[cmd], flagsets...)
}

func getFlagSets(cmdName string) []*pflag.FlagSet {
flagsets, ok := commandFlagSets[cmdName]
func getFlagSets(cmd *cobra.Command) []*pflag.FlagSet {
flagsets, ok := commandPathFlagSets[cmd.CommandPath()]
if !ok {
return []*pflag.FlagSet{config.GlobalFlags}
}
return append(flagsets, config.GlobalFlags)
newFlagSet := make([]*pflag.FlagSet, len(flagsets))
copy(newFlagSet, flagsets)
newFlagSet = append(newFlagSet, config.GlobalFlags)
return newFlagSet
}

const (
Expand All @@ -57,7 +75,7 @@ Examples:
Available Commands:{{range .Commands}}{{if (or .IsAvailableCommand (eq .Name "help"))}}
{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}
{{range getFlagSets .Name }}{{ title .Name}} Flags:
{{range getFlagSets . }}{{ title .Name}} Flags:
{{ .FlagUsages }}
{{end}}Get help:
-h, --help Help for any command or subcommand
Expand Down
4 changes: 3 additions & 1 deletion cmd/common/template/usage_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ func TestUsage(t *testing.T) {
flags.String("baz", "", "baz usage")
cmd.Flags().AddFlagSet(flags)

RegisterFlagSets(cmd.Name(), flags)
RegisterFlagSets(cmd, flags)
cmd.SetUsageTemplate(Usage)

subCmd := &cobra.Command{
Expand All @@ -50,6 +50,8 @@ func TestUsage(t *testing.T) {
}
cmd.AddCommand(subCmd)

Initialize()

var out strings.Builder
cmd.SetOut(&out)
cmd.Usage()
Expand Down
2 changes: 1 addition & 1 deletion cmd/list/list.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ func New(vp *viper.Viper) *cobra.Command {

// add config.ServerFlags to the help template as these flags are used by
// this command
template.RegisterFlagSets(listCmd.Name(), config.ServerFlags)
template.RegisterFlagSets(listCmd, config.ServerFlags)

listCmd.AddCommand(
newNodeCommand(vp),
Expand Down
2 changes: 1 addition & 1 deletion cmd/list/node.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ func newNodeCommand(vp *viper.Viper) *cobra.Command {
}, cobra.ShellCompDirectiveDefault
})

template.RegisterFlagSets(listCmd.Name(), formattingFlags, config.ServerFlags)
template.RegisterFlagSets(listCmd, formattingFlags, config.ServerFlags)
return listCmd
}

Expand Down
2 changes: 1 addition & 1 deletion cmd/observe/agent_events.go
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ func newAgentEventsCommand(vp *viper.Viper, flagSets ...*pflag.FlagSet) *cobra.C
},
}

template.RegisterFlagSets(agentEventsCmd.Name(), flagSets...)
template.RegisterFlagSets(agentEventsCmd, flagSets...)

return agentEventsCmd
}
Expand Down
2 changes: 1 addition & 1 deletion cmd/observe/debug_events.go
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ func newDebugEventsCommand(vp *viper.Viper, flagSets ...*pflag.FlagSet) *cobra.C
},
}

template.RegisterFlagSets(debugEventsCmd.Name(), flagSets...)
template.RegisterFlagSets(debugEventsCmd, flagSets...)

return debugEventsCmd
}
Expand Down
2 changes: 1 addition & 1 deletion cmd/observe/flows.go
Original file line number Diff line number Diff line change
Expand Up @@ -578,7 +578,7 @@ more.`,

formattingFlags.AddFlagSet(observeFormattingFlags)

template.RegisterFlagSets(observeCmd.Name(), selectorFlags, filterFlags, rawFilterFlags, formattingFlags, config.ServerFlags, otherFlags)
template.RegisterFlagSets(observeCmd, selectorFlags, filterFlags, rawFilterFlags, formattingFlags, config.ServerFlags, otherFlags)

return observeCmd
}
Expand Down
9 changes: 5 additions & 4 deletions cmd/record/record.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,15 @@ import (
"time"

recorderpb "github.com/cilium/cilium/api/v1/recorder"
"github.com/cilium/hubble/cmd/common/config"
"github.com/cilium/hubble/cmd/common/conn"
"github.com/cilium/hubble/cmd/common/template"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"github.com/spf13/viper"
"google.golang.org/grpc"
"google.golang.org/protobuf/types/known/durationpb"

"github.com/cilium/hubble/cmd/common/config"
"github.com/cilium/hubble/cmd/common/conn"
"github.com/cilium/hubble/cmd/common/template"
)

const (
Expand Down Expand Up @@ -91,7 +92,7 @@ protocols are TCP, UDP and ANY.`,
recorderFlags.DurationVar(&timeLimit, "time-limit", 0, "Sets a limit on how long to capture on each node")

recordCmd.Flags().AddFlagSet(recorderFlags)
template.RegisterFlagSets(recordCmd.Name(), config.ServerFlags, recorderFlags)
template.RegisterFlagSets(recordCmd, config.ServerFlags, recorderFlags)

return recordCmd
}
Expand Down
2 changes: 1 addition & 1 deletion cmd/reflect/reflect.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ func New(vp *viper.Viper) *cobra.Command {

// add config.ServerFlags to the help template as these flags are used by
// this command
template.RegisterFlagSets(reflectCmd.Name(), config.ServerFlags)
template.RegisterFlagSets(reflectCmd, config.ServerFlags)

return reflectCmd
}
Expand Down
6 changes: 5 additions & 1 deletion cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,9 @@ func New() *cobra.Command {

// NewWithViper creates a new root command with the given viper.
func NewWithViper(vp *viper.Viper) *cobra.Command {
// Initialize must be called after the sub-commands are all added
defer template.Initialize()

rootCmd := &cobra.Command{
Use: "hubble",
Short: "CLI",
Expand Down Expand Up @@ -83,7 +86,7 @@ func NewWithViper(vp *viper.Viper) *cobra.Command {
// add it by default in the help template
// config.GlobalFlags is always added to the help template as it's global
// to all commands
template.RegisterFlagSets(rootCmd.Name())
template.RegisterFlagSets(rootCmd)
rootCmd.SetUsageTemplate(template.Usage)

rootCmd.SetErr(os.Stderr)
Expand All @@ -99,6 +102,7 @@ func NewWithViper(vp *viper.Viper) *cobra.Command {
version.New(),
watch.New(vp),
)

return rootCmd
}

Expand Down
2 changes: 1 addition & 1 deletion cmd/status/status.go
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ connectivity health check.`,

// add config.ServerFlags to the help template as these flags are used by
// this command
template.RegisterFlagSets(statusCmd.Name(), config.ServerFlags)
template.RegisterFlagSets(statusCmd, config.ServerFlags)

return statusCmd
}
Expand Down
2 changes: 1 addition & 1 deletion cmd/watch/watch.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ func New(vp *viper.Viper) *cobra.Command {

// add config.ServerFlags to the help template as these flags are used by
// this command
template.RegisterFlagSets(peerCmd.Name(), config.ServerFlags)
template.RegisterFlagSets(peerCmd, config.ServerFlags)

peerCmd.AddCommand(
newPeerCommand(vp),
Expand Down

0 comments on commit a5e8697

Please sign in to comment.