fix: cancel Cobra parent context on interrupt (#2567)

## Description

Right now different functions handle interrupts in different ways. Some
will register their own signal listeners while others will use the
global signal handler that exits the program. This means that contexts
are never cancelled giving the process time to shut down and clean up.
This change adds a signal handler to the Cobra parent context and
removes the other signal handlers. Now all commands will use the same
signal handler.

## Related Issue

Fixes #2505 
Relates to #2520 

## Checklist before merging

- [x] Test, docs, adr added or updated as needed
- [x] [Contributor Guide
Steps](https://github.com/defenseunicorns/zarf/blob/main/.github/CONTRIBUTING.md#developer-workflow)
followed

Co-authored-by: Austin Abro <37223396+AustinAbro321@users.noreply.github.com>

main.go

@@ -5,7 +5,10 @@
 package main
 
 import (
+	"context"
 	"embed"
+	"os/signal"
+	"syscall"
 
 	"github.com/defenseunicorns/zarf/src/cmd"
 	"github.com/defenseunicorns/zarf/src/config"
@@ -19,7 +22,10 @@ var cosignPublicKey string
 var zarfSchema embed.FS
 
 func main() {
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
+
 	config.CosignPublicKey = cosignPublicKey
 	lint.ZarfSchema = zarfSchema
-	cmd.Execute()
+	cmd.Execute(ctx)
 }

src/cmd/common/setup.go

@@ -21,8 +21,6 @@ var LogLevelCLI string
 
 // SetupCLI sets up the CLI logging, interrupt functions, and more
 func SetupCLI() {
-	ExitOnInterrupt()
-
 	match := map[string]message.LogLevel{
 		"warn":  message.WarnLevel,
 		"info":  message.InfoLevel,

src/cmd/common/utils.go

@@ -6,18 +6,11 @@ package common
 
 import (
 	"context"
-	"os"
-	"os/signal"
-	"syscall"
 
-	"github.com/defenseunicorns/zarf/src/config/lang"
 	"github.com/defenseunicorns/zarf/src/pkg/cluster"
 	"github.com/defenseunicorns/zarf/src/pkg/message"
 )
 
-// SuppressGlobalInterrupt suppresses the global error on an interrupt
-var SuppressGlobalInterrupt = false
-
 // SetBaseDirectory sets the base directory. This is a directory with a zarf.yaml.
 func SetBaseDirectory(args []string) string {
 	if len(args) > 0 {
@@ -26,18 +19,6 @@ func SetBaseDirectory(args []string) string {
 	return "."
 }
 
-// ExitOnInterrupt catches an interrupt and exits with fatal error
-func ExitOnInterrupt() {
-	c := make(chan os.Signal, 1)
-	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
-	go func() {
-		<-c
-		if !SuppressGlobalInterrupt {
-			message.Fatal(lang.ErrInterrupt, lang.ErrInterrupt.Error())
-		}
-	}()
-}
-
 // NewClusterOrDie creates a new Cluster instance and waits for the cluster to be ready or throws a fatal error.
 func NewClusterOrDie(ctx context.Context) *cluster.Cluster {
 	timeoutCtx, cancel := context.WithTimeout(ctx, cluster.DefaultTimeout)

src/cmd/connect.go

@@ -7,8 +7,6 @@ package cmd
 import (
 	"fmt"
 	"os"
-	"os/signal"
-	"syscall"
 
 	"github.com/defenseunicorns/zarf/src/cmd/common"
 	"github.com/defenseunicorns/zarf/src/config/lang"
@@ -72,17 +70,12 @@ var (
 				}
 			}
 
-			// Keep this open until an interrupt signal is received.
-			interruptChan := make(chan os.Signal, 1)
-			signal.Notify(interruptChan, os.Interrupt, syscall.SIGTERM)
-			common.SuppressGlobalInterrupt = true
-
 			// Wait for the interrupt signal or an error.
 			select {
+			case <-ctx.Done():
+				spinner.Successf(lang.CmdConnectTunnelClosed, url)
 			case err = <-tunnel.ErrChan():
 				spinner.Fatalf(err, lang.CmdConnectErrService, err.Error())
-			case <-interruptChan:
-				spinner.Successf(lang.CmdConnectTunnelClosed, url)
 			}
 			os.Exit(0)
 		},

src/cmd/internal.go

@@ -38,17 +38,17 @@ var agentCmd = &cobra.Command{
 	Use:   "agent",
 	Short: lang.CmdInternalAgentShort,
 	Long:  lang.CmdInternalAgentLong,
-	Run: func(_ *cobra.Command, _ []string) {
-		agent.StartWebhook()
+	RunE: func(cmd *cobra.Command, _ []string) error {
+		return agent.StartWebhook(cmd.Context())
 	},
 }
 
 var httpProxyCmd = &cobra.Command{
 	Use:   "http-proxy",
 	Short: lang.CmdInternalProxyShort,
 	Long:  lang.CmdInternalProxyLong,
-	Run: func(_ *cobra.Command, _ []string) {
-		agent.StartHTTPProxy()
+	RunE: func(cmd *cobra.Command, _ []string) error {
+		return agent.StartHTTPProxy(cmd.Context())
 	},
 }
 

src/cmd/root.go

@@ -10,15 +10,16 @@ import (
 	"os"
 	"strings"
 
+	"github.com/pterm/pterm"
+	"github.com/spf13/cobra"
+
 	"github.com/defenseunicorns/zarf/src/cmd/common"
 	"github.com/defenseunicorns/zarf/src/cmd/tools"
 	"github.com/defenseunicorns/zarf/src/config"
 	"github.com/defenseunicorns/zarf/src/config/lang"
 	"github.com/defenseunicorns/zarf/src/pkg/layout"
 	"github.com/defenseunicorns/zarf/src/pkg/message"
 	"github.com/defenseunicorns/zarf/src/types"
-	"github.com/pterm/pterm"
-	"github.com/spf13/cobra"
 )
 
 var (
@@ -33,16 +34,10 @@ var rootCmd = &cobra.Command{
 		if common.CheckVendorOnlyFromPath(cmd) {
 			return
 		}
-
 		// Don't log the help command
 		if cmd.Parent() == nil {
 			config.SkipLogFile = true
 		}
-
-		// Set the global context for the root command and all child commands
-		ctx := context.Background()
-		cmd.SetContext(ctx)
-
 		common.SetupCLI()
 	},
 	Short:         lang.RootCmdShort,
@@ -67,8 +62,8 @@ var rootCmd = &cobra.Command{
 }
 
 // Execute is the entrypoint for the CLI.
-func Execute() {
-	cmd, err := rootCmd.ExecuteC()
+func Execute(ctx context.Context) {
+	cmd, err := rootCmd.ExecuteContextC(ctx)
 	if err == nil {
 		return
 	}

src/cmd/tools/crane.go

@@ -10,7 +10,6 @@ import (
 	"strings"
 
 	"github.com/AlecAivazis/survey/v2"
-	"github.com/defenseunicorns/zarf/src/cmd/common"
 	"github.com/defenseunicorns/zarf/src/config"
 	"github.com/defenseunicorns/zarf/src/config/lang"
 	"github.com/defenseunicorns/zarf/src/internal/packager/images"
@@ -39,9 +38,6 @@ func init() {
 		Aliases: []string{"r", "crane"},
 		Short:   lang.CmdToolsRegistryShort,
 		PersistentPreRun: func(cmd *cobra.Command, _ []string) {
-
-			common.ExitOnInterrupt()
-
 			// The crane options loading here comes from the rootCmd of crane
 			craneOptions = append(craneOptions, crane.WithContext(cmd.Context()))
 			// TODO(jonjohnsonjr): crane.Verbose option?

src/internal/agent/start.go

@@ -6,10 +6,11 @@ package agent
 
 import (
 	"context"
+	"errors"
 	"net/http"
-	"os"
-	"os/signal"
-	"syscall"
+	"time"
+
+	"golang.org/x/sync/errgroup"
 
 	"github.com/defenseunicorns/zarf/src/config/lang"
 	agentHttp "github.com/defenseunicorns/zarf/src/internal/agent/http"
@@ -27,35 +28,40 @@ const (
 )
 
 // StartWebhook launches the Zarf agent mutating webhook in the cluster.
-func StartWebhook() {
+func StartWebhook(ctx context.Context) error {
 	message.Debug("agent.StartWebhook()")
-
-	startServer(agentHttp.NewAdmissionServer(httpPort))
+	return startServer(ctx, agentHttp.NewAdmissionServer(httpPort))
 }
 
 // StartHTTPProxy launches the zarf agent proxy in the cluster.
-func StartHTTPProxy() {
+func StartHTTPProxy(ctx context.Context) error {
 	message.Debug("agent.StartHttpProxy()")
-
-	startServer(agentHttp.NewProxyServer(httpPort))
+	return startServer(ctx, agentHttp.NewProxyServer(httpPort))
 }
 
-func startServer(server *http.Server) {
-	go func() {
-		if err := server.ListenAndServeTLS(tlsCert, tlsKey); err != nil && err != http.ErrServerClosed {
-			message.Fatal(err, lang.AgentErrStart)
+func startServer(ctx context.Context, srv *http.Server) error {
+	g, gCtx := errgroup.WithContext(ctx)
+	g.Go(func() error {
+		err := srv.ListenAndServeTLS(tlsCert, tlsKey)
+		if err != nil && !errors.Is(err, http.ErrServerClosed) {
+			return err
 		}
-	}()
-
+		return nil
+	})
+	g.Go(func() error {
+		<-gCtx.Done()
+		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cancel()
+		err := srv.Shutdown(ctx)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
 	message.Infof(lang.AgentInfoPort, httpPort)
-
-	// listen shutdown signal
-	signalChan := make(chan os.Signal, 1)
-	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
-	<-signalChan
-
-	message.Infof(lang.AgentInfoShutdown)
-	if err := server.Shutdown(context.Background()); err != nil {
-		message.Fatal(err, lang.AgentErrShutdown)
+	err := g.Wait()
+	if err != nil {
+		return err
 	}
+	return nil
 }