cli-plugins: move socket code into common package

Signed-off-by: Bjorn Neergaard <bjorn.neergaard@docker.com>

cli-plugins/plugin/plugin.go

@@ -3,26 +3,19 @@ package plugin
 import (
 	"context"
 	"encoding/json"
-	"errors"
 	"fmt"
-	"io"
-	"net"
 	"os"
 	"sync"
 
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli-plugins/manager"
+	"github.com/docker/cli/cli-plugins/socket"
 	"github.com/docker/cli/cli/command"
 	"github.com/docker/cli/cli/connhelper"
 	"github.com/docker/docker/client"
 	"github.com/spf13/cobra"
 )
 
-// CLIPluginSocketEnvKey is used to pass the plugin being
-// executed the abstract socket name it should listen on to know
-// when the CLI has exited.
-const CLIPluginSocketEnvKey = "DOCKER_CLI_PLUGIN_SOCKET"
-
 // PersistentPreRunE must be called by any plugin command (or
 // subcommand) which uses the cobra `PersistentPreRun*` hook. Plugins
 // which do not make use of `PersistentPreRun*` do not need to call
@@ -33,38 +26,6 @@ const CLIPluginSocketEnvKey = "DOCKER_CLI_PLUGIN_SOCKET"
 // called.
 var PersistentPreRunE func(*cobra.Command, []string) error
 
-// closeOnCLISocketClose connects to the socket specified
-// by the DOCKER_CLI_PLUGIN_SOCKET env var, if present, and attempts
-// to read from it until it receives an EOF, which signals that
-// the CLI is going to exit and the plugin should also exit.
-func closeOnCLISocketClose(cancel func()) {
-	socketAddr, ok := os.LookupEnv(CLIPluginSocketEnvKey)
-	if !ok {
-		// if a plugin compiled against a more recent version of docker/cli
-		// is executed by an older CLI binary, ignore missing environment
-		// variable and behave as usual
-		return
-	}
-	addr, err := net.ResolveUnixAddr("unix", socketAddr)
-	if err != nil {
-		return
-	}
-	cliCloseConn, err := net.DialUnix("unix", nil, addr)
-	if err != nil {
-		return
-	}
-
-	go func() {
-		b := make([]byte, 1)
-		for {
-			_, err := cliCloseConn.Read(b)
-			if errors.Is(err, io.EOF) {
-				cancel()
-			}
-		}
-	}()
-}
-
 // RunPlugin executes the specified plugin command
 func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta manager.Metadata) error {
 	tcmd := newPluginCommand(dockerCli, plugin, meta)
@@ -81,7 +42,8 @@ func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta manager
 			}
 			ctx, cancel := context.WithCancel(cmdContext)
 			cmd.SetContext(ctx)
-			closeOnCLISocketClose(cancel)
+			// Set up the context to cancel based on signalling via CLI socket.
+			socket.ConnectAndWait(cancel)
 
 			var opts []command.CLIOption
 			if os.Getenv("DOCKER_CLI_PLUGIN_USE_DIAL_STDIO") != "" {

cli-plugins/socket/socket.go

@@ -0,0 +1,77 @@
+package socket
+
+import (
+	"errors"
+	"io"
+	"net"
+	"os"
+
+	"github.com/docker/distribution/uuid"
+)
+
+// EnvKey represents the well-known environment variable used to pass the plugin being
+// executed the socket name it should listen on to coordinate with the host CLI.
+const EnvKey = "DOCKER_CLI_PLUGIN_SOCKET"
+
+// SetupConn sets up a Unix socket listener, establishes a goroutine to handle connections
+// and update the conn pointer, and returns the environment variable to pass to the plugin.
+func SetupConn(conn **net.UnixConn) (string, error) {
+	listener, err := listen()
+	if err != nil {
+		return "", err
+	}
+
+	accept(listener, conn)
+
+	return EnvKey + "=" + listener.Addr().String(), nil
+}
+
+func listen() (*net.UnixListener, error) {
+	return net.ListenUnix("unix", &net.UnixAddr{
+		Name: "@docker_cli_" + uuid.Generate().String(),
+		Net:  "unix",
+	})
+}
+
+func accept(listener *net.UnixListener, conn **net.UnixConn) {
+	defer listener.Close()
+
+	go func() {
+		for {
+			// ignore error here, if we failed to accept a connection,
+			// conn is nil and we fallback to previous behavior
+			*conn, _ = listener.AcceptUnix()
+		}
+	}()
+}
+
+// ConnectAndWait connects to the socket passed via well-known env var,
+// if present, and attempts to read from it until it receives an EOF, at which
+// point cb is called.
+func ConnectAndWait(cb func()) {
+	socketAddr, ok := os.LookupEnv(EnvKey)
+	if !ok {
+		// if a plugin compiled against a more recent version of docker/cli
+		// is executed by an older CLI binary, ignore missing environment
+		// variable and behave as usual
+		return
+	}
+	addr, err := net.ResolveUnixAddr("unix", socketAddr)
+	if err != nil {
+		return
+	}
+	conn, err := net.DialUnix("unix", nil, addr)
+	if err != nil {
+		return
+	}
+
+	go func() {
+		b := make([]byte, 1)
+		for {
+			_, err := conn.Read(b)
+			if errors.Is(err, io.EOF) {
+				cb()
+			}
+		}
+	}()
+}

cmd/docker/docker.go

@@ -11,13 +11,12 @@ import (
 
 	"github.com/docker/cli/cli"
 	pluginmanager "github.com/docker/cli/cli-plugins/manager"
-	"github.com/docker/cli/cli-plugins/plugin"
+	"github.com/docker/cli/cli-plugins/socket"
 	"github.com/docker/cli/cli/command"
 	"github.com/docker/cli/cli/command/commands"
 	cliflags "github.com/docker/cli/cli/flags"
 	"github.com/docker/cli/cli/version"
 	platformsignals "github.com/docker/cli/cmd/docker/internal/signals"
-	"github.com/docker/distribution/uuid"
 	"github.com/docker/docker/api/types/versions"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
@@ -216,34 +215,20 @@ func setValidateArgs(dockerCli command.Cli, cmd *cobra.Command) {
 	})
 }
 
-func setupPluginSocket() (*net.UnixListener, error) {
-	return net.ListenUnix("unix", &net.UnixAddr{
-		Name: "@docker_cli_" + uuid.Generate().String(),
-		Net:  "unix",
-	})
-}
-
 func tryPluginRun(dockerCli command.Cli, cmd *cobra.Command, subcommand string, envs []string) error {
 	plugincmd, err := pluginmanager.PluginRunCommand(dockerCli, subcommand, cmd)
 	if err != nil {
 		return err
 	}
-	plugincmd.Env = append(envs, plugincmd.Env...)
 
+	// Establish the plugin socket, adding it to the environment under a well-known key if successful.
 	var conn *net.UnixConn
-	listener, err := setupPluginSocket()
+	socketenv, err := socket.SetupConn(&conn)
 	if err == nil {
-		defer listener.Close()
+		envs = append(envs, socketenv)
-		plugincmd.Env = append(plugincmd.Env, plugin.CLIPluginSocketEnvKey+"="+listener.Addr().String())
+	}
 
-		go func() {
+	plugincmd.Env = append(envs, plugincmd.Env...)
-			for {
-				// ignore error here, if we failed to accept a connection,
-				// conn is nil and we fallback to previous behavior
-				conn, _ = listener.AcceptUnix()
-			}
-		}()
-	}
 
 	const exitLimit = 3