docker_exec_session.go

package session

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/anmitsu/go-shlex"
	"github.com/borderzero/border0-cli/internal/api/models"
	"github.com/borderzero/border0-cli/internal/border0"
	sshConfig "github.com/borderzero/border0-cli/internal/ssh/config"
	"github.com/borderzero/border0-cli/internal/ssh/session/common"
	"github.com/borderzero/border0-go/lib/types/set"
	"github.com/borderzero/border0-go/lib/types/slice"
	"github.com/borderzero/border0-go/lib/types/wildcard"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/client"
	"github.com/manifoldco/promptui"
	"go.uber.org/zap"
	"golang.org/x/crypto/ssh"
)

type dockerExecSessionHandler struct {
	logger      *zap.Logger
	proxyConfig *sshConfig.ProxyConfig
}

// ensure dockerExecSessionHandler implements SessionHandler.
var _ SessionHandler = (*dockerExecSessionHandler)(nil)

type dockerExecSession struct {
	logger      *zap.Logger
	proxyConfig *sshConfig.ProxyConfig

	e2eeMetadata *border0.E2EEncryptionMetadata

	sshServerConfig *ssh.ServerConfig
	sshHeight       int
	sshWidth        int
	pty             bool
	// active channels
	downstreamSshConn  *ssh.ServerConn
	downstreamSshChans <-chan ssh.NewChannel
}

func NewDockerExecSessionHandler(
	logger *zap.Logger,
	proxyConfig *sshConfig.ProxyConfig,
) *dockerExecSessionHandler {
	return &dockerExecSessionHandler{
		logger:      logger,
		proxyConfig: proxyConfig,
	}
}

// Proxy runs the local proxying function between the connection to the
// remote Border0 proxy and the origin service (in this case the origin
// service is a connection to a remote docker executor / docker engine).
func (s *dockerExecSessionHandler) Proxy(conn net.Conn) {
	defer conn.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	dockerSess := &dockerExecSession{
		logger:          s.logger,
		proxyConfig:     s.proxyConfig,
		sshServerConfig: s.proxyConfig.SshServerConfig,
		sshWidth:        80,
		sshHeight:       24,
	}

	if s.proxyConfig.EndToEndEncryption {
		e2EEncryptionConn, ok := conn.(border0.E2EEncryptionConn)
		if !ok {
			conn.Close()
			s.logger.Error("failed to cast connection to e2eencryption")
			return
		}
		if e2EEncryptionConn.Metadata == nil {
			s.logger.Error("invalid e2e metadata")
			return
		}
		dockerSess.e2eeMetadata = e2EEncryptionConn.Metadata
		dockerSess.logger = dockerSess.logger.With(zap.String("session_key", dockerSess.e2eeMetadata.SessionKey))
		dockerSess.sshServerConfig.PublicKeyCallback = common.GetPublicKeyCallback(
			s.proxyConfig.OrgSshCA,
			s.proxyConfig.Border0API,
			s.proxyConfig.Socket,
			e2EEncryptionConn.Metadata,
			false,
		)
	}

	// accept SSH connection from Border0 proxy
	dsConn, dsChanns, dsReqs, err := ssh.NewServerConn(conn, dockerSess.proxyConfig.SshServerConfig)
	if err != nil {
		dockerSess.logger.Error("failed to accept ssh connection from upstream proxy", zap.Error(err))
		return
	}
	dockerSess.downstreamSshConn = dsConn
	dockerSess.downstreamSshChans = dsChanns

	// we don't support global requests (yet)
	// so we can disregard the reqs channel
	go ssh.DiscardRequests(dsReqs)

	if dockerSess.proxyConfig.EndToEndEncryption {
		username := dockerSess.downstreamSshConn.User()

		if s.proxyConfig.Username != "" {
			username = s.proxyConfig.Username
		}

		if err := dockerSess.proxyConfig.Border0API.UpdateSession(models.SessionUpdate{
			SessionKey: dockerSess.e2eeMetadata.SessionKey,
			Socket:     dockerSess.proxyConfig.Socket,
			UserData:   ",sshuser=" + username,
		}); err != nil {
			dockerSess.logger.Error("failed to update session", zap.Error(err))
			return
		}
	}

	if dockerSess.e2eeMetadata != nil {
		var max_session_duration int
		for _, action := range dockerSess.e2eeMetadata.AllowedActions {
			switch action := action.(type) {
			case models.ExtendedAction:
				if action.Ssh != nil {
					if action.Ssh.MaxSessionDurationSeconds > max_session_duration {
						max_session_duration = action.Ssh.MaxSessionDurationSeconds
					}
				}
			}
		}

		if max_session_duration > 0 {
			go func() {
				select {
				case <-time.After(time.Duration(max_session_duration) * time.Second):
					s.logger.Info("disconnecting session due to max session duration", zap.Duration("duration", time.Duration(max_session_duration)*time.Second))
					cancel()
				case <-ctx.Done():
				}
			}()
		}
	}

	if err := dockerSess.handleChannels(ctx); err != nil {
		s.logger.Error("failed to handle channels", zap.Error(err))
		return
	}
}

func (s *dockerExecSession) handleChannels(ctx context.Context) error {
	defer s.downstreamSshConn.Close()

	channelCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	for {
		select {
		case newChannel, ok := <-s.downstreamSshChans:
			if !ok {
				return nil
			}

			if newChannel.ChannelType() != "session" {
				newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
				continue
			}

			channel, requests, err := newChannel.Accept()
			if err != nil {
				return fmt.Errorf("failed to accept channel: %s", err)
			}

			go func(in <-chan *ssh.Request) {
				for req := range in {
					switch {
					case req == nil:
						continue
					// handled mostly for the benefit of session recordings
					case req.Type == "pty-req":
						termLen := req.Payload[3]
						w, h := common.ParseDims(req.Payload[termLen+4:])
						s.sshWidth = int(w)
						s.sshHeight = int(h)
						if req.WantReply {
							req.Reply(true, nil)
						}
						s.pty = true
					// handled mostly for the benefit of session recordings
					case req.Type == "window-change":
						w, h := common.ParseDims(req.Payload)
						s.sshWidth = int(w)
						s.sshHeight = int(h)
						if req.WantReply {
							req.Reply(true, nil)
						}
					case req.Type == "shell", req.Type == "exec":
						if req.WantReply {
							req.Reply(true, nil)
						}
						go s.handleChannel(channelCtx, channel, s.downstreamSshConn.User(), req)
					default:
						req.Reply(false, nil)
					}
				}
			}(requests)
		case <-ctx.Done():
			return nil
		}
	}
}

func (s *dockerExecSession) handleChannel(
	ctx context.Context,
	channel ssh.Channel,
	user string,
	req *ssh.Request,
) {
	defer channel.Close()

	var extendedActions []models.ExtendedAction

	if s.e2eeMetadata != nil {
		var allowed bool
		for _, allowedAction := range s.e2eeMetadata.AllowedActions {
			switch action := allowedAction.(type) {
			case string:
				allowed = true
			case models.ExtendedAction:
				if action.Ssh.DockerExec.Enabled {
					allowed = true
					extendedActions = append(extendedActions, action)
				}

			}
		}

		if !allowed {
			s.logger.Warn("docker exec access denied by policy")
			channel.Stderr().Write([]byte("docker exec access denied by policy\r\n"))
			channel.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{Status: 1}))
			return
		}
	}

	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
	if err != nil {
		channel.Write([]byte("An error occured. Try again later..."))
		s.logger.Error("failed to initialize new docker client", zap.Error(err))
		return
	}

	var payload = struct{ Value string }{}
	ssh.Unmarshal(req.Payload, &payload)
	command := payload.Value

	if s.proxyConfig.IsRecordingEnabled() {
		pwc := NewPipeWriteChannel(channel)
		channel = pwc
		r := common.NewRecording(s.logger, pwc.reader, s.proxyConfig.Socket.SocketID, s.e2eeMetadata.SessionKey, s.proxyConfig.Border0API, s.sshWidth, s.sshHeight)
		if err := r.Record(); err != nil {
			channel.Write([]byte("An error occured. Try again later..."))
			s.logger.Error("failed to record session", zap.Error(err))
			channel.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{Status: 1}))
			return
		}

		if command != "" {
			pwc.logWriter.Write([]byte(fmt.Sprintf("%s\n", command)))
		}
		defer r.Stop()
	}

	var clientTarget string
	userSlice := strings.SplitN(user, "/", 2)
	if len(userSlice) == 2 { // format is container/user
		clientTarget = userSlice[0]
		user = userSlice[1]
	}

	container, err := s.askForTarget(ctx, channel, cli, clientTarget, extendedActions)
	if err != nil {
		s.logger.Error("failed to determine target for docker exec", zap.Error(err))
		channel.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{Status: 1}))
		return
	}

	if container == "" {
		return // no containers available or user cancelled operation
	}

	switch req.Type {
	case "shell":
		// we iterate over the slice and not the set
		// because order is not maintained for the set
		shells := []string{"bash", "zsh", "ash", "sh"}
		shellSet := set.New(shells...)
		for _, shell := range shells {
			if shellSet.Size() == 0 {
				channel.Write([]byte("No shells available in the target container :("))
				s.logger.Error("no shells available in the target container", zap.Error(err))
				return
			}

			if err := s.exec(ctx, cli, channel, container, user, []string{shell}, true); err != nil {
				if err.Error() == "shell not found in container" {
					shellSet.Remove(shell)
					continue
				}

				channel.Stderr().Write([]byte("An error occured. Try again later...\r\n"))
				channel.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{Status: 1}))
				s.logger.Error(err.Error())
			}
			break
		}
	case "exec":
		commandArgs, err := shlex.Split(command, true)
		if err != nil {
			channel.Stderr().Write([]byte("failed to parse command\r\n"))
			s.logger.Error("failed to parse command", zap.Error(err))
			return
		}

		if err := s.exec(ctx, cli, channel, container, user, commandArgs, false); err != nil {
			channel.Stderr().Write([]byte("An error occured. Try again later...\r\n"))
			s.logger.Error(err.Error())
			channel.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{Status: 1}))
			return
		}
	}

	channel.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{Status: 0}))
}

func (s *dockerExecSession) exec(ctx context.Context, cli *client.Client, channel ssh.Channel, container, user string, command []string, returnOnNotFound bool) error {
	exec, err := cli.ContainerExecCreate(ctx, container, types.ExecConfig{
		User:       user, // User that will run the command
		Privileged: true, // Is the container in privileged mode
		Tty:        s.pty,
		// ConsoleSize  *[2]uint `json:",omitempty"` // Initial console size [height, width]
		AttachStdin:  true,
		AttachStderr: true,
		AttachStdout: true,
		Detach:       false, // Execute in detach mode
		// DetachKeys   string   // Escape keys for detach
		// Env          []string // Environment variables
		// WorkingDir   string   // Working directory
		Cmd: command, // Execution commands and args
	})
	if err != nil {
		return fmt.Errorf("failed to perform ContainerExecCreate operation against docker container %s: %w", container, err)
	}

	hijackedResponse, err := cli.ContainerExecAttach(ctx, exec.ID, types.ExecStartCheck{
		Detach: false, // ExecStart will first check if it's detached
		Tty:    true,  // Check if there's a tty
		// ConsoleSize  *[2]uint `json:",omitempty"`// Terminal size [height, width], unused if Tty == false
	})
	if err != nil {
		return fmt.Errorf("failed to perform ContainerExecAttach operation against docker container %s: %w", container, err)
	}
	defer hijackedResponse.Close()

	if returnOnNotFound {
		// perform one read to check whether the container
		// errored in finding shell in binaries path
		buf := make([]byte, 256)
		n, err := hijackedResponse.Conn.Read(buf)
		if err != nil {
			channel.Write([]byte("An error occured. Try again later..."))
			s.logger.Error(
				"failed to perform first-read on connection to docker remote executor",
				zap.String("container", container),
				zap.Error(err),
			)
		}

		if strings.Contains(string(buf[:n]), "executable file not found") ||
			strings.Contains(string(buf[:n]), "command terminated with exit code 127") {
			return fmt.Errorf("shell not found in container")
		}

		// if no error, we need to write that first read back to the ssh channel
		channel.Write(buf[:n])
	}

	// stitch together connection to proxy and connection to executor
	errs := make(chan error, 2)
	var wg sync.WaitGroup // wait group is for channel closure

	go func() { wg.Wait(); close(errs) }()
	wg.Add(2)
	go func() { defer wg.Done(); _, e := io.Copy(channel, hijackedResponse.Conn); errs <- e }()
	go func() { defer wg.Done(); _, e := io.Copy(hijackedResponse.Conn, channel); errs <- e }()

	if err = <-errs; err != nil {
		if !errors.Is(err, io.EOF) {
			channel.SendRequest("exit-status", false, []byte{1, 0, 0, 0})
			return fmt.Errorf("failed to proxy between border0 proxy and docker remote executor: %w", err)
		}
	}

	channel.SendRequest("exit-status", false, []byte{0, 0, 0, 0})
	return nil
}

func (s *dockerExecSession) askForTarget(ctx context.Context, channel ssh.Channel, cli *client.Client, clientTarget string, extendedActions []models.ExtendedAction) (string, error) {
	containerListCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	containers, err := cli.ContainerList(containerListCtx, container.ListOptions{})
	if err != nil {
		return "", fmt.Errorf("failed to list Docker containers: %v", err)
	}

	if len(containers) == 0 {
		channel.Stderr().Write([]byte("No containers available to you - sorry :(\r\n"))
		return "", nil
	}

	type containerPrompt struct {
		ID               string
		PromptIdentifier string
	}
	containerPrompts := []containerPrompt{}

	for _, container := range containers {
		// if there's an allowlist present, process it
		if len(s.proxyConfig.DockerExecProxy.ContainerNameAllowlist) > 0 {
			include := false
			for _, name := range container.Names {
				for _, template := range s.proxyConfig.DockerExecProxy.ContainerNameAllowlist {
					// match, ignoring leading slash
					if wildcard.Match(template, name) || wildcard.Match(template, strings.TrimPrefix(name, "/")) {
						include = true
						break
					}
				}
			}
			if !include {
				continue // won't include this one, move on to next container
			}
		}

		if len(extendedActions) > 0 {
			include := false
			for _, name := range container.Names {
				for _, action := range extendedActions {
					if len(action.Ssh.DockerExec.AllowedContainers) > 0 {
						for _, allowedContainer := range action.Ssh.DockerExec.AllowedContainers {
							if wildcard.Match(allowedContainer, name) || wildcard.Match(allowedContainer, strings.TrimPrefix(name, "/")) {
								include = true
								break
							}
						}
					} else {
						include = true
					}
				}
			}

			if !include {
				continue // won't include this one, move on to next container
			}
		}

		include := false
		if clientTarget != "" {
			for _, name := range container.Names {
				if strings.EqualFold(strings.TrimPrefix(name, "/"), clientTarget) {
					include = true
					break
				}
			}

			if !include {
				continue // won't include this one, move on to next container
			}
		}

		friendlyName := container.ID
		if len(container.Names) > 0 {
			friendlyName = strings.Join(
				// strip leading slashes
				slice.Transform(
					container.Names,
					func(name string) string { return strings.TrimPrefix(name, "/") },
				),
				// use comma as join delimeter
				", ",
			)
		}

		containerPrompts = append(containerPrompts, containerPrompt{
			ID:               container.ID,
			PromptIdentifier: fmt.Sprintf("%s (%s)", friendlyName, container.Image),
		})
	}

	if len(containerPrompts) == 0 {
		channel.Stderr().Write([]byte("No containers available to you - sorry :(\r\n"))
		return "", errors.New("no allowlisted containers available")
	}
	if len(containerPrompts) == 1 {
		return containerPrompts[0].ID, nil
	}

	// sorting by alphanumeric order of prompt identifiers
	sort.Slice(containerPrompts, func(i, j int) bool {
		return containerPrompts[i].PromptIdentifier < containerPrompts[j].PromptIdentifier
	})

	conatinerPrompt := promptui.Select{
		Label:             "Choose a container",
		Items:             slice.Transform(containerPrompts, func(c containerPrompt) string { return c.PromptIdentifier }),
		Stdout:            channel,
		Stdin:             channel,
		StartInSearchMode: true,
		Searcher: func(input string, index int) bool {
			return strings.Contains(strings.ToLower(containerPrompts[index].PromptIdentifier), strings.ToLower(input))
		},
		Size: 10, // default is 5
	}
	index, _, err := conatinerPrompt.Run()
	if err != nil {
		// handle cancellation by user (e.g. ^C or ^D)
		if errors.Is(err, promptui.ErrInterrupt) || errors.Is(err, promptui.ErrEOF) {
			return "", nil
		}
		return "", fmt.Errorf("unable to select container: %v", err)
	}

	return containerPrompts[index].ID, nil
}