From dbf992f91f605b5fffa6b08172e6ed1cc7e84d80 Mon Sep 17 00:00:00 2001 From: Bjorn Neergaard Date: Fri, 12 Jan 2024 11:17:03 -0700 Subject: [PATCH] cli-plugins: move socket code into common package Signed-off-by: Bjorn Neergaard --- cli-plugins/plugin/plugin.go | 44 ++------------------- cli-plugins/socket/socket.go | 77 ++++++++++++++++++++++++++++++++++++ cmd/docker/docker.go | 27 +++---------- 3 files changed, 86 insertions(+), 62 deletions(-) create mode 100644 cli-plugins/socket/socket.go diff --git a/cli-plugins/plugin/plugin.go b/cli-plugins/plugin/plugin.go index 829e1baa8e..54471b41b0 100644 --- a/cli-plugins/plugin/plugin.go +++ b/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") != "" { diff --git a/cli-plugins/socket/socket.go b/cli-plugins/socket/socket.go new file mode 100644 index 0000000000..8f27c37439 --- /dev/null +++ b/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() + } + } + }() +} diff --git a/cmd/docker/docker.go b/cmd/docker/docker.go index b1961da6d9..2b1b951a50 100644 --- a/cmd/docker/docker.go +++ b/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,35 +215,21 @@ 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() - plugincmd.Env = append(plugincmd.Env, plugin.CLIPluginSocketEnvKey+"="+listener.Addr().String()) - - 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() - } - }() + envs = append(envs, socketenv) } + plugincmd.Env = append(envs, plugincmd.Env...) + const exitLimit = 3 signals := make(chan os.Signal, exitLimit)