diff --git a/cmds/audittail/cmd/root.go b/cmds/audittail/cmd/root.go index 45498ca9..ac13eb71 100644 --- a/cmds/audittail/cmd/root.go +++ b/cmds/audittail/cmd/root.go @@ -35,6 +35,10 @@ const ( // rootCmd represents the base command when called without any subcommands. var rootCmd = NewRootCmd() +func GetCmd() *cobra.Command { + return rootCmd +} + func NewRootCmd() *cobra.Command { c := &cobra.Command{ Use: "audittail", diff --git a/cmds/audittail/cmd/root_test.go b/cmds/audittail/cmd/root_test.go index 78a89227..22bf8217 100644 --- a/cmds/audittail/cmd/root_test.go +++ b/cmds/audittail/cmd/root_test.go @@ -237,3 +237,10 @@ func TestFileTrailerUnknownErrorWhenTailing(t *testing.T) { err := ft.tailFile(context.Background()) require.Error(t, err, "unexpected success") } + +func TestRootCmdSingletonGet(t *testing.T) { + t.Parallel() + + c := GetCmd() + require.Equal(t, &rootCmd, &c, "GetCmd() should return the rootCmd singleton") +} diff --git a/cmds/audittail/main.go b/cmds/audittail/main.go index 016aab98..72bc9cd3 100644 --- a/cmds/audittail/main.go +++ b/cmds/audittail/main.go @@ -16,14 +16,42 @@ limitations under the License. package main import ( + "context" "os" + "os/signal" "github.com/metal-toolbox/auditevent/cmds/audittail/cmd" ) -func main() { - c := cmd.NewRootCmd() - if err := c.Execute(); err != nil { - os.Exit(1) +func runCommand() int { + ctx := context.Background() + + // trap Ctrl+C and call cancel on the context + ctx, cancel := context.WithCancel(ctx) + + c := make(chan os.Signal, 1) + signal.Notify(c, os.Interrupt) + + defer func() { + signal.Stop(c) + cancel() + }() + + go func() { + select { + case <-c: + cancel() + case <-ctx.Done(): + } + }() + + if err := cmd.GetCmd().ExecuteContext(ctx); err != nil { + return 1 } + + return 0 +} + +func main() { + os.Exit(runCommand()) }