From 1d666b410526b485958b92ec9bdd3e0d5e0acd56 Mon Sep 17 00:00:00 2001
From: Alano Terblanche <18033717+Benehiko@users.noreply.github.com>
Date: Fri, 26 Apr 2024 13:03:56 +0200
Subject: [PATCH] feat: wire ctx into plugin hooks

Signed-off-by: Alano Terblanche <18033717+Benehiko@users.noreply.github.com>
---
 cli-plugins/manager/hooks.go  | 24 ++++++++++++++++--------
 cli-plugins/manager/plugin.go |  5 +++--
 cmd/docker/docker.go          |  4 ++--
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/cli-plugins/manager/hooks.go b/cli-plugins/manager/hooks.go
index 4d99e679b0..a0f4d5ee46 100644
--- a/cli-plugins/manager/hooks.go
+++ b/cli-plugins/manager/hooks.go
@@ -1,6 +1,7 @@
 package manager
 
 import (
+	"context"
 	"encoding/json"
 	"strings"
 
@@ -27,29 +28,36 @@ type HookPluginData struct {
 // a main CLI command was executed. It calls the hook subcommand for all
 // present CLI plugins that declare support for hooks in their metadata and
 // parses/prints their responses.
-func RunCLICommandHooks(dockerCli command.Cli, rootCmd, subCommand *cobra.Command, cmdErrorMessage string) {
+func RunCLICommandHooks(ctx context.Context, dockerCli command.Cli, rootCmd, subCommand *cobra.Command, cmdErrorMessage string) {
 	commandName := strings.TrimPrefix(subCommand.CommandPath(), rootCmd.Name()+" ")
 	flags := getCommandFlags(subCommand)
 
-	runHooks(dockerCli, rootCmd, subCommand, commandName, flags, cmdErrorMessage)
+	runHooks(ctx, dockerCli, rootCmd, subCommand, commandName, flags, cmdErrorMessage)
 }
 
 // RunPluginHooks is the entrypoint for the hooks execution flow
 // after a plugin command was just executed by the CLI.
-func RunPluginHooks(dockerCli command.Cli, rootCmd, subCommand *cobra.Command, args []string) {
+func RunPluginHooks(ctx context.Context, dockerCli command.Cli, rootCmd, subCommand *cobra.Command, args []string) {
 	commandName := strings.Join(args, " ")
 	flags := getNaiveFlags(args)
 
-	runHooks(dockerCli, rootCmd, subCommand, commandName, flags, "")
+	runHooks(ctx, dockerCli, rootCmd, subCommand, commandName, flags, "")
 }
 
-func runHooks(dockerCli command.Cli, rootCmd, subCommand *cobra.Command, invokedCommand string, flags map[string]string, cmdErrorMessage string) {
-	nextSteps := invokeAndCollectHooks(dockerCli, rootCmd, subCommand, invokedCommand, flags, cmdErrorMessage)
+func runHooks(ctx context.Context, dockerCli command.Cli, rootCmd, subCommand *cobra.Command, invokedCommand string, flags map[string]string, cmdErrorMessage string) {
+	nextSteps := invokeAndCollectHooks(ctx, dockerCli, rootCmd, subCommand, invokedCommand, flags, cmdErrorMessage)
 
 	hooks.PrintNextSteps(dockerCli.Err(), nextSteps)
 }
 
-func invokeAndCollectHooks(dockerCli command.Cli, rootCmd, subCmd *cobra.Command, subCmdStr string, flags map[string]string, cmdErrorMessage string) []string {
+func invokeAndCollectHooks(ctx context.Context, dockerCli command.Cli, rootCmd, subCmd *cobra.Command, subCmdStr string, flags map[string]string, cmdErrorMessage string) []string {
+	// check if the context was cancelled before invoking hooks
+	select {
+	case <-ctx.Done():
+		return nil
+	default:
+	}
+
 	pluginsCfg := dockerCli.ConfigFile().Plugins
 	if pluginsCfg == nil {
 		return nil
@@ -67,7 +75,7 @@ func invokeAndCollectHooks(dockerCli command.Cli, rootCmd, subCmd *cobra.Command
 			continue
 		}
 
-		hookReturn, err := p.RunHook(HookPluginData{
+		hookReturn, err := p.RunHook(ctx, HookPluginData{
 			RootCmd:      match,
 			Flags:        flags,
 			CommandError: cmdErrorMessage,
diff --git a/cli-plugins/manager/plugin.go b/cli-plugins/manager/plugin.go
index 2cffafacca..877241e0b8 100644
--- a/cli-plugins/manager/plugin.go
+++ b/cli-plugins/manager/plugin.go
@@ -1,6 +1,7 @@
 package manager
 
 import (
+	"context"
 	"encoding/json"
 	"os"
 	"os/exec"
@@ -105,13 +106,13 @@ func newPlugin(c Candidate, cmds []*cobra.Command) (Plugin, error) {
 
 // RunHook executes the plugin's hooks command
 // and returns its unprocessed output.
-func (p *Plugin) RunHook(hookData HookPluginData) ([]byte, error) {
+func (p *Plugin) RunHook(ctx context.Context, hookData HookPluginData) ([]byte, error) {
 	hDataBytes, err := json.Marshal(hookData)
 	if err != nil {
 		return nil, wrapAsPluginError(err, "failed to marshall hook data")
 	}
 
-	pCmd := exec.Command(p.Path, p.Name, HookSubcommandName, string(hDataBytes))
+	pCmd := exec.CommandContext(ctx, p.Path, p.Name, HookSubcommandName, string(hDataBytes))
 	pCmd.Env = os.Environ()
 	pCmd.Env = append(pCmd.Env, ReexecEnvvar+"="+os.Args[0])
 	hookCmdOutput, err := pCmd.Output()
diff --git a/cmd/docker/docker.go b/cmd/docker/docker.go
index f68bda6fe7..64990573d6 100644
--- a/cmd/docker/docker.go
+++ b/cmd/docker/docker.go
@@ -337,7 +337,7 @@ func runDocker(ctx context.Context, dockerCli *command.DockerCli) error {
 			err := tryPluginRun(dockerCli, cmd, args[0], envs)
 			if err == nil {
 				if dockerCli.HooksEnabled() && dockerCli.Out().IsTerminal() && ccmd != nil {
-					pluginmanager.RunPluginHooks(dockerCli, cmd, ccmd, args)
+					pluginmanager.RunPluginHooks(ctx, dockerCli, cmd, ccmd, args)
 				}
 				return nil
 			}
@@ -362,7 +362,7 @@ func runDocker(ctx context.Context, dockerCli *command.DockerCli) error {
 		if err != nil {
 			errMessage = err.Error()
 		}
-		pluginmanager.RunCLICommandHooks(dockerCli, cmd, subCommand, errMessage)
+		pluginmanager.RunCLICommandHooks(ctx, dockerCli, cmd, subCommand, errMessage)
 	}
 
 	return err