From e29e11047668a5e5557be074c3c55383651aa1ff Mon Sep 17 00:00:00 2001 From: Jesse Duffield Date: Sun, 9 Jan 2022 17:58:51 +1100 Subject: [PATCH] refactor and add test --- pkg/commands/docker.go | 132 +---------------------------- pkg/commands/ssh/ssh.go | 156 +++++++++++++++++++++++++++++++++++ pkg/commands/ssh/ssh_test.go | 102 +++++++++++++++++++++++ pkg/utils/utils.go | 1 - 4 files changed, 262 insertions(+), 129 deletions(-) create mode 100644 pkg/commands/ssh/ssh.go create mode 100644 pkg/commands/ssh/ssh_test.go diff --git a/pkg/commands/docker.go b/pkg/commands/docker.go index 5702a80a..9b4e825f 100644 --- a/pkg/commands/docker.go +++ b/pkg/commands/docker.go @@ -6,23 +6,18 @@ import ( "encoding/json" "fmt" "io" - "io/ioutil" ogLog "log" - "net" - "net/url" - "os" "os/exec" - "path" "sort" "strings" "sync" - "syscall" "time" "github.com/acarl005/stripansi" "github.com/docker/docker/api/types" "github.com/docker/docker/client" "github.com/imdario/mergo" + "github.com/jesseduffield/lazydocker/pkg/commands/ssh" "github.com/jesseduffield/lazydocker/pkg/config" "github.com/jesseduffield/lazydocker/pkg/i18n" "github.com/jesseduffield/lazydocker/pkg/utils" @@ -54,6 +49,8 @@ type DockerCommand struct { Closers []io.Closer } +var _ io.Closer = &DockerCommand{} + // LimitedDockerCommand is a stripped-down DockerCommand with just the methods the container/service/image might need type LimitedDockerCommand interface { NewCommandObject(CommandObject) CommandObject @@ -75,130 +72,9 @@ func (c *DockerCommand) NewCommandObject(obj CommandObject) CommandObject { return defaultObj } -// handleSSHDockerHost overrides the DOCKER_HOST environment variable -// to point towards a local unix socket tunneled over SSH to the specified ssh host. -func handleSSHDockerHost() (io.Closer, error) { - const key = "DOCKER_HOST" - ctx := context.Background() - u, err := url.Parse(os.Getenv(key)) - if err != nil { - // if no or an invalid docker host is specified, continue nominally - return noopCloser{}, nil - } - - // if the docker host scheme is "ssh", forward the docker socket before creating the client - if u.Scheme == "ssh" { - tunnel, err := createDockerHostTunnel(ctx, u.Host) - if err != nil { - return noopCloser{}, fmt.Errorf("tunnel ssh docker host: %w", err) - } - err = os.Setenv(key, tunnel.SocketPath) - if err != nil { - return noopCloser{}, fmt.Errorf("override DOCKER_HOST to tunneled socket: %w", err) - } - - return tunnel, nil - } - return noopCloser{}, nil -} - -type noopCloser struct{} - -func (noopCloser) Close() error { return nil } - -type TunneledDockerHost struct { - SocketPath string - cmd *exec.Cmd -} - -var _ io.Closer = (*TunneledDockerHost)(nil) - -func (t *TunneledDockerHost) Close() error { - return syscall.Kill(-t.cmd.Process.Pid, syscall.SIGKILL) -} - -func createDockerHostTunnel(ctx context.Context, remoteHost string) (*TunneledDockerHost, error) { - socketDir, err := ioutil.TempDir("/tmp", "lazydocker-sshtunnel-") - if err != nil { - return nil, fmt.Errorf("create ssh tunnel tmp file: %w", err) - } - localSocket := path.Join(socketDir, "dockerhost.sock") - - cmd, err := tunnelSSH(ctx, remoteHost, localSocket) - if err != nil { - return nil, fmt.Errorf("tunnel docker host over ssh: %w", err) - } - - // set a reasonable timeout, then wait for the socket to dial successfully - // before attempting to create a new docker client - const socketTunnelTimeout = 8 * time.Second - ctx, cancel := context.WithTimeout(ctx, socketTunnelTimeout) - defer cancel() - - err = retrySocketDial(ctx, localSocket) - if err != nil { - return nil, fmt.Errorf("ssh tunneled socket never became available: %w", err) - } - - // construct the new DOCKER_HOST url with the proper scheme - newDockerHostURL := url.URL{Scheme: "unix", Path: localSocket} - return &TunneledDockerHost{ - SocketPath: newDockerHostURL.String(), - cmd: cmd, - }, nil -} - -// Attempt to dial the socket until it becomes available. -// The retry loop will continue until the parent context is canceled. -func retrySocketDial(ctx context.Context, socketPath string) error { - t := time.NewTicker(1 * time.Second) - defer t.Stop() - - for { - select { - case <-ctx.Done(): - return ctx.Err() - case <-t.C: - } - // attempt to dial the socket, exit on success - err := tryDial(ctx, socketPath) - if err != nil { - continue - } - return nil - } -} - -// Try to dial the specified unix socket, immediately close the connection if successfully created. -func tryDial(ctx context.Context, socketPath string) error { - var dialer net.Dialer - conn, err := dialer.DialContext(ctx, "unix", socketPath) - if err != nil { - return err - } - defer conn.Close() - return nil -} -func tunnelSSH(ctx context.Context, host, localSocket string) (*exec.Cmd, error) { - cmd := exec.CommandContext(ctx, "ssh", "-L", localSocket+":/var/run/docker.sock", host, "-N") - cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} - err := cmd.Start() - if err != nil { - return nil, err - } - return cmd, nil -} - -// Build a new docker client from the environment. -// -// Handle special cases including `ssh://` host schemes. -func clientBuilder(c *client.Client) error { - return nil -} - // NewDockerCommand it runs docker commands func NewDockerCommand(log *logrus.Entry, osCommand *OSCommand, tr *i18n.TranslationSet, config *config.AppConfig, errorChan chan error) (*DockerCommand, error) { - tunnelCloser, err := handleSSHDockerHost() + tunnelCloser, err := ssh.NewSSHHandler().HandleSSHDockerHost() if err != nil { ogLog.Fatal(err) } diff --git a/pkg/commands/ssh/ssh.go b/pkg/commands/ssh/ssh.go new file mode 100644 index 00000000..474affa6 --- /dev/null +++ b/pkg/commands/ssh/ssh.go @@ -0,0 +1,156 @@ +package ssh + +import ( + "context" + "fmt" + "io" + "io/ioutil" + "net" + "net/url" + "os" + "os/exec" + "path" + "syscall" + "time" +) + +type dependencies struct { + // storing all these dependencies as fields for the sake of testing + dialContext func(ctx context.Context, network, addr string) (io.Closer, error) + startCmd func(*exec.Cmd) error + tempDir func(dir string, pattern string) (name string, err error) + getenv func(key string) string + setenv func(key, value string) error +} + +type SSHHandler struct { + deps dependencies +} + +func NewSSHHandler() *SSHHandler { + return &SSHHandler{ + deps: dependencies{ + dialContext: func(ctx context.Context, network, addr string) (io.Closer, error) { + return (&net.Dialer{}).DialContext(ctx, network, addr) + }, + startCmd: func(cmd *exec.Cmd) error { return cmd.Start() }, + tempDir: ioutil.TempDir, + getenv: os.Getenv, + setenv: os.Setenv, + }, + } +} + +// HandleSSHDockerHost overrides the DOCKER_HOST environment variable +// to point towards a local unix socket tunneled over SSH to the specified ssh host. +func (self *SSHHandler) HandleSSHDockerHost() (io.Closer, error) { + const key = "DOCKER_HOST" + ctx := context.Background() + u, err := url.Parse(self.deps.getenv(key)) + if err != nil { + // if no or an invalid docker host is specified, continue nominally + return noopCloser{}, nil + } + + // if the docker host scheme is "ssh", forward the docker socket before creating the client + if u.Scheme == "ssh" { + tunnel, err := self.createDockerHostTunnel(ctx, u.Host) + if err != nil { + return noopCloser{}, fmt.Errorf("tunnel ssh docker host: %w", err) + } + err = self.deps.setenv(key, tunnel.socketPath) + if err != nil { + return noopCloser{}, fmt.Errorf("override DOCKER_HOST to tunneled socket: %w", err) + } + + return tunnel, nil + } + return noopCloser{}, nil +} + +type noopCloser struct{} + +func (noopCloser) Close() error { return nil } + +type tunneledDockerHost struct { + socketPath string + cmd *exec.Cmd +} + +var _ io.Closer = (*tunneledDockerHost)(nil) + +func (t *tunneledDockerHost) Close() error { + return syscall.Kill(-t.cmd.Process.Pid, syscall.SIGKILL) +} + +func (self *SSHHandler) createDockerHostTunnel(ctx context.Context, remoteHost string) (*tunneledDockerHost, error) { + socketDir, err := self.deps.tempDir("/tmp", "lazydocker-sshtunnel-") + if err != nil { + return nil, fmt.Errorf("create ssh tunnel tmp file: %w", err) + } + localSocket := path.Join(socketDir, "dockerhost.sock") + + cmd, err := self.tunnelSSH(ctx, remoteHost, localSocket) + if err != nil { + return nil, fmt.Errorf("tunnel docker host over ssh: %w", err) + } + + // set a reasonable timeout, then wait for the socket to dial successfully + // before attempting to create a new docker client + const socketTunnelTimeout = 8 * time.Second + ctx, cancel := context.WithTimeout(ctx, socketTunnelTimeout) + defer cancel() + + err = self.retrySocketDial(ctx, localSocket) + if err != nil { + return nil, fmt.Errorf("ssh tunneled socket never became available: %w", err) + } + + // construct the new DOCKER_HOST url with the proper scheme + newDockerHostURL := url.URL{Scheme: "unix", Path: localSocket} + return &tunneledDockerHost{ + socketPath: newDockerHostURL.String(), + cmd: cmd, + }, nil +} + +// Attempt to dial the socket until it becomes available. +// The retry loop will continue until the parent context is canceled. +func (self *SSHHandler) retrySocketDial(ctx context.Context, socketPath string) error { + t := time.NewTicker(1 * time.Second) + defer t.Stop() + + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.C: + } + // attempt to dial the socket, exit on success + err := self.tryDial(ctx, socketPath) + if err != nil { + continue + } + return nil + } +} + +// Try to dial the specified unix socket, immediately close the connection if successfully created. +func (self *SSHHandler) tryDial(ctx context.Context, socketPath string) error { + conn, err := self.deps.dialContext(ctx, "unix", socketPath) + if err != nil { + return err + } + defer conn.Close() + return nil +} + +func (self *SSHHandler) tunnelSSH(ctx context.Context, host, localSocket string) (*exec.Cmd, error) { + cmd := exec.CommandContext(ctx, "ssh", "-L", localSocket+":/var/run/docker.sock", host, "-N") + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + err := self.deps.startCmd(cmd) + if err != nil { + return nil, err + } + return cmd, nil +} diff --git a/pkg/commands/ssh/ssh_test.go b/pkg/commands/ssh/ssh_test.go new file mode 100644 index 00000000..5c16e224 --- /dev/null +++ b/pkg/commands/ssh/ssh_test.go @@ -0,0 +1,102 @@ +package ssh + +import ( + "context" + "io" + "os/exec" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestSSHHandlerHandleSSHDockerHost(t *testing.T) { + type scenario struct { + testName string + envVarValue string + expectedDialContextCount int + expectedStartCmdCount int + } + + scenarios := []scenario{ + { + testName: "No env var set", + envVarValue: "", + expectedDialContextCount: 0, + expectedStartCmdCount: 0, + }, + { + testName: "Env var set with https scheme", + envVarValue: "https://myhost.com", + expectedStartCmdCount: 0, + expectedDialContextCount: 0, + }, + { + testName: "Env var set with ssh scheme", + envVarValue: "ssh://myhost@192.168.5.178", + expectedStartCmdCount: 1, + expectedDialContextCount: 1, + }, + } + + for _, s := range scenarios { + s := s + t.Run(s.testName, func(t *testing.T) { + getenv := func(key string) string { + if key != "DOCKER_HOST" { + t.Errorf("Expected key to be DOCKER_HOST, got %s", key) + } + + return s.envVarValue + } + + tempDir := func(dir string, pattern string) (string, error) { + assert.Equal(t, "/tmp", dir) + assert.Equal(t, "lazydocker-sshtunnel-", pattern) + + return "/tmp/lazydocker-ssh-tunnel-12345", nil + } + + setenv := func(key, value string) error { + assert.Equal(t, "DOCKER_HOST", key) + assert.Equal(t, "unix:///tmp/lazydocker-ssh-tunnel-12345/dockerhost.sock", value) + return nil + } + + startCmdCount := 0 + startCmd := func(cmd *exec.Cmd) error { + assert.EqualValues(t, []string{"ssh", "-L", "/tmp/lazydocker-ssh-tunnel-12345/dockerhost.sock:/var/run/docker.sock", "192.168.5.178", "-N"}, cmd.Args) + assert.Equal(t, true, cmd.SysProcAttr.Setpgid) + + startCmdCount++ + + return nil + } + + dialContextCount := 0 + dialContext := func(ctx context.Context, network string, address string) (io.Closer, error) { + assert.Equal(t, "unix", network) + assert.Equal(t, "/tmp/lazydocker-ssh-tunnel-12345/dockerhost.sock", address) + + dialContextCount++ + + return noopCloser{}, nil + } + + handler := &SSHHandler{ + deps: dependencies{ + dialContext: dialContext, + startCmd: startCmd, + tempDir: tempDir, + getenv: getenv, + setenv: setenv, + }, + } + + _, err := handler.HandleSSHDockerHost() + assert.NoError(t, err) + + assert.Equal(t, s.expectedDialContextCount, dialContextCount) + assert.Equal(t, s.expectedStartCmdCount, startCmdCount) + }) + } +} diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go index 4a33f999..f005c5b9 100644 --- a/pkg/utils/utils.go +++ b/pkg/utils/utils.go @@ -373,4 +373,3 @@ func CloseMany(closers []io.Closer) error { } return multiErr(errs) } -