fix: plugin and container context cancellation

Signed-off-by: Alano Terblanche <18033717+Benehiko@users.noreply.github.com>

Author: Benehiko

cli/command/container/attach.go

@@ -105,7 +105,13 @@ func RunAttach(ctx context.Context, dockerCLI command.Cli, containerID string, o
 
 	if opts.Proxy && !c.Config.Tty {
 		sigc := notifyAllSignals()
-		go ForwardAllSignals(ctx, apiClient, containerID, sigc)
+		// since the ForwardAllSignals already registers its own signal handler
+		// and should notify the daemon independently of the clients ctx cancellation
+		// we should use a background context to avoid the ForwardAllSignals from
+		// returning before all signals are forwarded to the daemon
+		bgCtx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		go ForwardAllSignals(bgCtx, apiClient, containerID, sigc)
 		defer signal.StopCatch(sigc)
 	}
 

cli/command/container/run.go

@@ -150,7 +150,13 @@ func runContainer(ctx context.Context, dockerCli command.Cli, runOpts *runOption
 	}
 	if runOpts.sigProxy {
 		sigc := notifyAllSignals()
-		go ForwardAllSignals(ctx, apiClient, containerID, sigc)
+		// since the ForwardAllSignals already registers its own signal handler
+		// and should notify the daemon independently of the clients ctx cancellation
+		// we should use a background context to avoid the ForwardAllSignals from
+		// returning before all signals are forwarded to the daemon
+		bgCtx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		go ForwardAllSignals(bgCtx, apiClient, containerID, sigc)
 		defer signal.StopCatch(sigc)
 	}
 

cmd/docker/docker.go

@@ -37,7 +37,8 @@ func main() {
 	dockerCli, err := command.NewDockerCli(command.WithBaseContext(ctx))
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
-		os.Exit(1)
+		defer os.Exit(1)
+		return
 	}
 	logrus.SetOutput(dockerCli.Err())
 	otel.SetErrorHandler(debug.OTELErrorHandler)
@@ -50,15 +51,19 @@ func main() {
 			// StatusError should only be used for errors, and all errors should
 			// have a non-zero exit status, so never exit with 0
 			if sterr.StatusCode == 0 {
-				os.Exit(1)
+				defer os.Exit(1)
+				return
 			}
-			os.Exit(sterr.StatusCode)
+			defer os.Exit(sterr.StatusCode)
+			return
 		}
 		if errdefs.IsCancelled(err) {
-			os.Exit(0)
+			defer os.Exit(0)
+			return
 		}
 		fmt.Fprintln(dockerCli.Err(), err)
-		os.Exit(1)
+		defer os.Exit(1)
+		return
 	}
 }
 
@@ -248,42 +253,56 @@ func tryPluginRun(ctx context.Context, dockerCli command.Cli, cmd *cobra.Command
 	// notify the plugin via the PluginServer (or signal) as appropriate.
 	const exitLimit = 2
 
+	tryTerminatePlugin := func(force bool) {
+		// If stdin is a TTY, the kernel will forward
+		// signals to the subprocess because the shared
+		// pgid makes the TTY a controlling terminal.
+		//
+		// The plugin should have it's own copy of this
+		// termination logic, and exit after 3 retries
+		// on it's own.
+		if dockerCli.Out().IsTerminal() {
+			return
+		}
+
+		// Terminate the plugin server, which will
+		// close all connections with plugin
+		// subprocesses, and signal them to exit.
+		//
+		// Repeated invocations will result in EINVAL,
+		// or EBADF; but that is fine for our purposes.
+		_ = srv.Close()
+
+		// If we're still running after 3 interruptions
+		// (SIGINT/SIGTERM), send a SIGKILL to the plugin as a
+		// final attempt to terminate, and exit.
+		if force {
+			_, _ = fmt.Fprint(dockerCli.Err(), "got 3 SIGTERM/SIGINTs, forcefully exiting\n")
+			_ = plugincmd.Process.Kill()
+			os.Exit(1)
+		}
+	}
+
 	go func() {
 		retries := 0
+		force := false
+		// catch the first signal through context cancellation
 		<-ctx.Done()
+		tryTerminatePlugin(force)
 
+		// register subsequent signals
 		signals := make(chan os.Signal, exitLimit)
 		signal.Notify(signals, platformsignals.TerminationSignals...)
 
 		for range signals {
-			// If stdin is a TTY, the kernel will forward
-			// signals to the subprocess because the shared
-			// pgid makes the TTY a controlling terminal.
-			//
-			// The plugin should have it's own copy of this
-			// termination logic, and exit after 3 retries
-			// on it's own.
-			if dockerCli.Out().IsTerminal() {
-				continue
-			}
-
-			// Terminate the plugin server, which will
-			// close all connections with plugin
-			// subprocesses, and signal them to exit.
-			//
-			// Repeated invocations will result in EINVAL,
-			// or EBADF; but that is fine for our purposes.
-			_ = srv.Close()
-
+			retries++
 			// If we're still running after 3 interruptions
 			// (SIGINT/SIGTERM), send a SIGKILL to the plugin as a
 			// final attempt to terminate, and exit.
-			retries++
 			if retries >= exitLimit {
-				_, _ = fmt.Fprint(dockerCli.Err(), "got 3 SIGTERM/SIGINTs, forcefully exiting\n")
-				_ = plugincmd.Process.Kill()
-				os.Exit(1)
+				force = true
 			}
+			tryTerminatePlugin(force)
 		}
 	}()