From 63a2225fcfe2c0abb4f129e7bdf0ed255e5863ef Mon Sep 17 00:00:00 2001
From: Alano Terblanche <18033717+Benehiko@users.noreply.github.com>
Date: Mon, 8 Apr 2024 10:11:09 +0200
Subject: [PATCH] fix: hoist signal handling from prompt confirmation

Signed-off-by: Alano Terblanche <18033717+Benehiko@users.noreply.github.com>
---
 cli/command/utils.go      | 10 +---------
 cli/command/utils_test.go |  6 +++++-
 cmd/docker/docker.go      |  7 ++++++-
 internal/test/cmd.go      |  8 +++++---
 4 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/cli/command/utils.go b/cli/command/utils.go
index d7184f8c4e6a..5b2cb9721569 100644
--- a/cli/command/utils.go
+++ b/cli/command/utils.go
@@ -9,11 +9,9 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"os/signal"
 	"path/filepath"
 	"runtime"
 	"strings"
-	"syscall"
 
 	"github.com/docker/cli/cli/streams"
 	"github.com/docker/docker/api/types/filters"
@@ -103,11 +101,6 @@ func PromptForConfirmation(ctx context.Context, ins io.Reader, outs io.Writer, m
 
 	result := make(chan bool)
 
-	// Catch the termination signal and exit the prompt gracefully.
-	// The caller is responsible for properly handling the termination.
-	notifyCtx, notifyCancel := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
-	defer notifyCancel()
-
 	go func() {
 		var res bool
 		scanner := bufio.NewScanner(ins)
@@ -121,8 +114,7 @@ func PromptForConfirmation(ctx context.Context, ins io.Reader, outs io.Writer, m
 	}()
 
 	select {
-	case <-notifyCtx.Done():
-		// print a newline on termination
+	case <-ctx.Done():
 		_, _ = fmt.Fprintln(outs, "")
 		return false, ErrPromptTerminated
 	case r := <-result:
diff --git a/cli/command/utils_test.go b/cli/command/utils_test.go
index b1ea2dd74c51..1566067f3b38 100644
--- a/cli/command/utils_test.go
+++ b/cli/command/utils_test.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
 	"syscall"
@@ -135,6 +136,9 @@ func TestPromptForConfirmation(t *testing.T) {
 		}, promptResult{false, nil}},
 	} {
 		t.Run("case="+tc.desc, func(t *testing.T) {
+			notifyCtx, notifyCancel := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
+			t.Cleanup(notifyCancel)
+
 			buf.Reset()
 			promptReader, promptWriter = io.Pipe()
 
@@ -145,7 +149,7 @@ func TestPromptForConfirmation(t *testing.T) {
 
 			result := make(chan promptResult, 1)
 			go func() {
-				r, err := command.PromptForConfirmation(ctx, promptReader, promptOut, "")
+				r, err := command.PromptForConfirmation(notifyCtx, promptReader, promptOut, "")
 				result <- promptResult{r, err}
 			}()
 
diff --git a/cmd/docker/docker.go b/cmd/docker/docker.go
index 3b3b6cb5c90f..4ac8c52d17d5 100644
--- a/cmd/docker/docker.go
+++ b/cmd/docker/docker.go
@@ -349,10 +349,15 @@ func runDocker(ctx context.Context, dockerCli *command.DockerCli) error {
 		}
 	}
 
+	// Catch the termination signal and exit the prompt gracefully.
+	// The caller is responsible for properly handling the termination.
+	notifyCtx, notifyCancel := signal.NotifyContext(ctx, platformsignals.TerminationSignals...)
+	defer notifyCancel()
+
 	// We've parsed global args already, so reset args to those
 	// which remain.
 	cmd.SetArgs(args)
-	err = cmd.Execute()
+	err = cmd.ExecuteContext(notifyCtx)
 	if err != nil {
 		return err
 	}
diff --git a/internal/test/cmd.go b/internal/test/cmd.go
index 04df496a833b..52b44d66f1d4 100644
--- a/internal/test/cmd.go
+++ b/internal/test/cmd.go
@@ -3,7 +3,6 @@ package test
 import (
 	"context"
 	"os"
-	"syscall"
 	"testing"
 	"time"
 
@@ -32,8 +31,11 @@ func TerminatePrompt(ctx context.Context, t *testing.T, cmd *cobra.Command, cli
 	assert.NilError(t, err)
 	cli.SetIn(streams.NewIn(r))
 
+	notifyCtx, notifyCancel := context.WithCancel(ctx)
+	t.Cleanup(notifyCancel)
+
 	go func() {
-		errChan <- cmd.ExecuteContext(ctx)
+		errChan <- cmd.ExecuteContext(notifyCtx)
 	}()
 
 	writeCtx, writeCancel := context.WithTimeout(ctx, 100*time.Millisecond)
@@ -66,7 +68,7 @@ func TerminatePrompt(ctx context.Context, t *testing.T, cmd *cobra.Command, cli
 
 	// sigint and sigterm are caught by the prompt
 	// this allows us to gracefully exit the prompt with a 0 exit code
-	syscall.Kill(syscall.Getpid(), syscall.SIGINT)
+	notifyCancel()
 
 	select {
 	case <-errCtx.Done():