Cancel Cobra parent context on interrupt

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

@@ -38,10 +38,6 @@ var rootCmd = &cobra.Command{
 			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,
@@ -64,8 +60,9 @@ var rootCmd = &cobra.Command{
 }
 
 // Execute is the entrypoint for the CLI.
-func Execute() {
-	cobra.CheckErr(rootCmd.Execute())
+func Execute(ctx context.Context) {
+	err := rootCmd.ExecuteContext(ctx)
+	cobra.CheckErr(err)
 }
 
 func init() {

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
 }