From 043ae74622c3cb68f1332bc83d067c32778cbeb7 Mon Sep 17 00:00:00 2001
From: Kimmo Lehto <klehto@mirantis.com>
Date: Wed, 5 Feb 2025 11:06:53 +0200
Subject: [PATCH] Attempt a graceful exit on interrupt

Signed-off-by: Kimmo Lehto <klehto@mirantis.com>
---
 cmd/root.go          | 30 ++++++++++++++++++++++++++++++
 phase/uploadfiles.go |  5 ++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/cmd/root.go b/cmd/root.go
index 5918357fa..03ccfd192 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -1,9 +1,27 @@
 package cmd
 
 import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+
+	log "github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
 )
 
+func trapSignals(ctx context.Context, cancel context.CancelFunc) {
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+	select {
+	case <-ch:
+		log.Warnf("received an interrupt signal, aborting operation")
+		cancel()
+	case <-ctx.Done():
+		return
+	}
+}
+
 // App is the main urfave/cli.App for k0sctl
 var App = &cli.App{
 	Name:  "k0sctl",
@@ -31,4 +49,16 @@ var App = &cli.App{
 		completionCommand,
 	},
 	EnableBashCompletion: true,
+	Before: func(ctx *cli.Context) error {
+		if globalCancel == nil {
+			cancelCtx, cancel := context.WithCancel(ctx.Context)
+			ctx.Context = cancelCtx
+			globalCancel = cancel
+		}
+		go trapSignals(ctx.Context, globalCancel)
+		return nil
+	},
+	After: func(ctx *cli.Context) error {
+		return cancelTimeout(ctx)
+	},
 }
diff --git a/phase/uploadfiles.go b/phase/uploadfiles.go
index 486096c38..0c06a54f4 100644
--- a/phase/uploadfiles.go
+++ b/phase/uploadfiles.go
@@ -46,8 +46,11 @@ func (p *UploadFiles) Run(ctx context.Context) error {
 	return p.parallelDoUpload(ctx, p.Config.Spec.Hosts, p.uploadFiles)
 }
 
-func (p *UploadFiles) uploadFiles(_ context.Context, h *cluster.Host) error {
+func (p *UploadFiles) uploadFiles(ctx context.Context, h *cluster.Host) error {
 	for _, f := range h.Files {
+		if ctx.Err() != nil {
+			return fmt.Errorf("upload canceled: %w", ctx.Err())
+		}
 		var err error
 		if f.IsURL() {
 			err = p.uploadURL(h, f)