From 10f56bb5e3b64b09234b092a235388464246d533 Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Tue, 21 Mar 2023 22:58:24 +0400 Subject: [PATCH] Introduce "orchard ssh" and "orchard vnc" commands (#36) * proxy.Connections(): handle "use of closed network connection" error * Controller: less strict timeouts that work nicely for WebSockets * Worker: only attempt connect to the gRPC once our UID is known * Introduce "orchard ssh" and "orchard vnc" commands * Worker: prevent context leak by moving logic into a separate function * Fix linter errors * Port forwarding integration test * Check for "uname -mo" output --- go.mod | 2 + go.sum | 3 + internal/command/root.go | 4 + internal/command/ssh/ssh.go | 179 +++++++++++++++++++++++++++++ internal/command/vnc/vnc.go | 97 ++++++++++++++++ internal/controller/controller.go | 4 +- internal/proxy/proxy.go | 43 ++++--- internal/tests/integration_test.go | 64 +++++++++++ internal/worker/worker.go | 76 ++++++------ 9 files changed, 411 insertions(+), 61 deletions(-) create mode 100644 internal/command/ssh/ssh.go create mode 100644 internal/command/vnc/vnc.go diff --git a/go.mod b/go.mod index 32aa645..565e6c4 100644 --- a/go.mod +++ b/go.mod @@ -17,6 +17,7 @@ require ( github.com/hashicorp/go-version v1.6.0 github.com/manifoldco/promptui v0.9.0 github.com/mitchellh/go-grpc-net-conn v0.0.0-20200427190222-eb030e4876f0 + github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 github.com/spf13/cobra v1.6.0 github.com/stretchr/testify v1.8.1 go.uber.org/zap v1.24.0 @@ -64,6 +65,7 @@ require ( go.uber.org/atomic v1.10.0 // indirect go.uber.org/multierr v1.9.0 // indirect golang.org/x/sys v0.4.0 // indirect + golang.org/x/term v0.4.0 // indirect golang.org/x/text v0.6.0 // indirect google.golang.org/genproto v0.0.0-20230110181048-76db0878b65f // indirect gopkg.in/yaml.v2 v2.4.0 // indirect diff --git a/go.sum b/go.sum index c61d9a9..5414723 100644 --- a/go.sum +++ b/go.sum @@ -155,6 +155,8 @@ github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUA github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE= github.com/russross/blackfriday v1.5.2/go.mod h1:JO/DiYxRf+HjHt06OyowR9PTA263kcR/rfWxYHBV53g= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966 h1:JIAuq3EEf9cgbU6AtGPK4CTG3Zf6CKMNqf0MHTggAUA= +github.com/skratchdot/open-golang v0.0.0-20200116055534-eef842397966/go.mod h1:sUM3LWHvSMaG192sy56D9F7CNvL7jUJVXoqM1QKLnog= github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= @@ -247,6 +249,7 @@ golang.org/x/sys v0.4.0 h1:Zr2JFtRQNX3BCZ8YtxRE9hNJYC8J6I1MVbMg6owUp18= golang.org/x/sys v0.4.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.4.0 h1:O7UWfv5+A2qiuulQk30kVinPoMtoIPeVaKLEgLpVkvg= +golang.org/x/term v0.4.0/go.mod h1:9P2UbLfCdcvo3p/nzKvsmas4TnlujnuoV9hGgYzW1lQ= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= diff --git a/internal/command/root.go b/internal/command/root.go index af35613..b8d035d 100644 --- a/internal/command/root.go +++ b/internal/command/root.go @@ -9,7 +9,9 @@ import ( "github.com/cirruslabs/orchard/internal/command/get" "github.com/cirruslabs/orchard/internal/command/list" "github.com/cirruslabs/orchard/internal/command/portforward" + "github.com/cirruslabs/orchard/internal/command/ssh" "github.com/cirruslabs/orchard/internal/command/stop" + "github.com/cirruslabs/orchard/internal/command/vnc" "github.com/cirruslabs/orchard/internal/command/worker" "github.com/cirruslabs/orchard/internal/version" "github.com/spf13/cobra" @@ -28,6 +30,8 @@ func NewRootCmd() *cobra.Command { get.NewCommand(), list.NewCommand(), portforward.NewCommand(), + ssh.NewCommand(), + vnc.NewCommand(), deletepkg.NewCommand(), stop.NewCommand(), ) diff --git a/internal/command/ssh/ssh.go b/internal/command/ssh/ssh.go new file mode 100644 index 0000000..3ecdcca --- /dev/null +++ b/internal/command/ssh/ssh.go @@ -0,0 +1,179 @@ +package ssh + +import ( + "context" + "errors" + "fmt" + "github.com/cirruslabs/orchard/pkg/client" + "github.com/spf13/cobra" + "golang.org/x/crypto/ssh" + "golang.org/x/term" + "io" + "net" + "os" + "time" +) + +var ErrFailed = errors.New("ssh command failed") + +var username string +var password string + +func NewCommand() *cobra.Command { + command := &cobra.Command{ + Use: "ssh ", + Short: "SSH into the VM", + Args: cobra.ExactArgs(1), + RunE: run, + } + + command.PersistentFlags().StringVarP(&username, "username", "u", "", + "SSH username") + command.PersistentFlags().StringVarP(&password, "password", "p", "", + "SSH password") + + return command +} + +func run(cmd *cobra.Command, args []string) error { + name := args[0] + + client, err := client.New() + if err != nil { + return err + } + + wsConn, err := client.VMs().PortForward(cmd.Context(), name, 22) + if err != nil { + fmt.Printf("failed to forward an SSH port to VM %s: %v\n", name, err) + + return err + } + defer wsConn.Close() + + username, password = ChooseUsernameAndPassword(cmd.Context(), client, name, username, password) + + sshConfig := &ssh.ClientConfig{ + HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error { + return nil + }, + User: username, + Auth: []ssh.AuthMethod{ + ssh.Password(password), + }, + } + + sshConn, chans, reqs, err := ssh.NewClientConn(wsConn, "", sshConfig) + if err != nil { + return fmt.Errorf("%w: failed to establish an SSH connection: %v", ErrFailed, err) + } + + sshClient := ssh.NewClient(sshConn, chans, reqs) + + sshSess, err := sshClient.NewSession() + if err != nil { + return fmt.Errorf("%w: failed to open an SSH session: %v", ErrFailed, err) + } + + // Switch controlling terminal into raw mode, + // otherwise ANSI escape sequences that allow + // for cursor control and more wouldn't work + terminalFd := int(os.Stdin.Fd()) + state, err := term.MakeRaw(terminalFd) + if err != nil { + return fmt.Errorf("%w: failed to switch controlling terminal into raw mode: %v", ErrFailed, err) + } + defer func() { + _ = term.Restore(terminalFd, state) + }() + + width, height, err := term.GetSize(terminalFd) + if err != nil { + return err + } + + if err := sshSess.RequestPty("xterm-256color", height, width, ssh.TerminalModes{}); err != nil { + return fmt.Errorf("%w: failed to request the PTY from the SSH server: %v", ErrFailed, err) + } + + sshSess.Stdout = os.Stdout + sshSess.Stderr = os.Stderr + sshSessStdinPipe, err := sshSess.StdinPipe() + if err != nil { + return err + } + + go func() { + _, _ = io.Copy(sshSessStdinPipe, os.Stdin) + _ = sshSessStdinPipe.Close() + _ = sshSess.Close() + }() + + // Periodically adjust remote terminal size + go func() { + for { + newWidth, newHeight, err := term.GetSize(terminalFd) + if err != nil { + return + } + + if height == newHeight && width == newWidth { + continue + } + + if err := sshSess.WindowChange(newHeight, newWidth); err != nil { + continue + } + + height = newHeight + width = newWidth + + time.Sleep(time.Second) + } + }() + + if err := sshSess.Shell(); err != nil { + return err + } + + sshErr := make(chan error) + + go func() { + sshErr <- sshSess.Wait() + }() + + select { + case err := <-sshErr: + return err + case <-cmd.Context().Done(): + return cmd.Context().Err() + } +} + +func ChooseUsernameAndPassword( + ctx context.Context, + client *client.Client, + vmName string, + usernameFromUser string, + passwordFromUser string, +) (string, string) { + // User settings override everything + if usernameFromUser != "" || passwordFromUser != "" { + return usernameFromUser, passwordFromUser + } + + // Try to get the credentials from the VM's object stored on controller + vm, err := client.VMs().Get(ctx, vmName) + if err != nil { + fmt.Printf("failed to retrieve VM %s's credentials from the API server: %v\n", vmName, err) + } + + if vm.Username != "" && vm.Password != "" { + return vm.Username, vm.Password + } + + // Fall back + fmt.Println("no credentials specified or found, trying default admin:admin credentials...") + + return "admin", "admin" +} diff --git a/internal/command/vnc/vnc.go b/internal/command/vnc/vnc.go new file mode 100644 index 0000000..59cc0a0 --- /dev/null +++ b/internal/command/vnc/vnc.go @@ -0,0 +1,97 @@ +package vnc + +import ( + "fmt" + "github.com/cirruslabs/orchard/internal/command/ssh" + "github.com/cirruslabs/orchard/internal/proxy" + "github.com/cirruslabs/orchard/pkg/client" + "github.com/skratchdot/open-golang/open" + "github.com/spf13/cobra" + "net" +) + +const vncPort = 5900 + +var username string +var password string + +func NewCommand() *cobra.Command { + command := &cobra.Command{ + Use: "vnc ", + Short: "Open VNC session with the VM", + Args: cobra.ExactArgs(1), + RunE: run, + } + + command.PersistentFlags().StringVarP(&username, "username", "u", "", + "VNC username") + command.PersistentFlags().StringVarP(&password, "password", "p", "", + "VNC password") + + return command +} + +func run(cmd *cobra.Command, args []string) (err error) { + name := args[0] + + client, err := client.New() + if err != nil { + return err + } + + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return err + } + defer func() { + if listenerErr := listener.Close(); listenerErr != nil && err == nil { + err = listenerErr + } + }() + + fmt.Printf("forwarding %s -> %s:%d...\n", listener.Addr(), name, vncPort) + + errCh := make(chan error, 1) + + go func() { + for { + conn, err := listener.Accept() + if err != nil { + errCh <- err + + return + } + + go func() { + wsConn, err := client.VMs().PortForward(cmd.Context(), name, vncPort) + if err != nil { + fmt.Printf("failed to forward port: %v\n", err) + + return + } + + if err := proxy.Connections(wsConn, conn); err != nil { + fmt.Printf("failed to forward port: %v\n", err) + } + }() + } + }() + + username, password = ssh.ChooseUsernameAndPassword(cmd.Context(), client, name, username, password) + + openURL := fmt.Sprintf("vnc://%s:%s@%s", username, password, listener.Addr().String()) + openURLSanitized := fmt.Sprintf("vnc://%s@%s", username, listener.Addr().String()) + + fmt.Printf("opening %s...\n", openURLSanitized) + + if err := open.Start(openURL); err != nil { + fmt.Printf("failed to open: %v\n", err) + } + + select { + case <-cmd.Context().Done(): + return cmd.Context().Err() + case err := <-errCh: + return err + } +} diff --git a/internal/controller/controller.go b/internal/controller/controller.go index c77364a..c9d9b70 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -102,8 +102,8 @@ func New(opts ...Option) (*Controller, error) { }) controller.httpServer = &http.Server{ - Handler: h2c.NewHandler(handler, &http2.Server{}), - ReadTimeout: 60 * time.Second, + Handler: h2c.NewHandler(handler, &http2.Server{}), + ReadHeaderTimeout: 60 * time.Second, } return controller, nil diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 491f6a1..3dc46b1 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -3,32 +3,45 @@ package proxy import ( "io" "net" + "strings" ) -func Connections(left net.Conn, right net.Conn) error { - errCh := make(chan error, 2) +func Connections(left net.Conn, right net.Conn) (finalErr error) { + leftErrCh := make(chan error, 1) + rightErrCh := make(chan error, 1) + + recordErr := func(newErr error) { + if newErr != nil && finalErr == nil { + finalErr = newErr + } + } go func() { _, err := io.Copy(left, right) - errCh <- err + rightErrCh <- err }() go func() { _, err := io.Copy(right, left) - errCh <- err + leftErrCh <- err }() - firstErr := <-errCh - - // Force connection closure to unlock the other goroutine - left.Close() - right.Close() - - secondErr := <-errCh - - if firstErr != nil { - return firstErr + // Wait for some goroutine and then unlock the other goroutine + // by closing its source io.Reader + select { + case err := <-rightErrCh: + recordErr(err) + recordErr(left.Close()) + recordErr(<-leftErrCh) + case err := <-leftErrCh: + recordErr(err) + recordErr(right.Close()) + recordErr(<-rightErrCh) } - return secondErr + if strings.Contains(finalErr.Error(), "use of closed network connection") { + finalErr = nil + } + + return finalErr } diff --git a/internal/tests/integration_test.go b/internal/tests/integration_test.go index e94676d..e0dd0b7 100644 --- a/internal/tests/integration_test.go +++ b/internal/tests/integration_test.go @@ -7,6 +7,9 @@ import ( "github.com/cirruslabs/orchard/pkg/client" v1 "github.com/cirruslabs/orchard/pkg/resource/v1" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/crypto/ssh" + "net" "net/http" "testing" "time" @@ -178,3 +181,64 @@ func StartIntegrationTestEnvironment(t *testing.T) *client.Client { } return devClient } + +func TestPortForwarding(t *testing.T) { + ctx := context.Background() + + devClient := StartIntegrationTestEnvironment(t) + + // Create a generic macOS VM + err := devClient.VMs().Create(ctx, &v1.VM{ + Meta: v1.Meta{ + Name: "test-vm", + }, + Image: "ghcr.io/cirruslabs/macos-ventura-base:latest", + CPU: 4, + Memory: 8 * 1024, + Softnet: false, + Headless: true, + }) + require.NoError(t, err) + + // Wait for the VM to start + var vm *v1.VM + + require.True(t, Wait(2*time.Minute, func() bool { + vm, err = devClient.VMs().Get(ctx, "test-vm") + require.NoError(t, err) + + t.Logf("Waiting for the VM to start, current status: %s", vm.Status) + + return vm.Status == v1.VMStatusRunning || vm.Status == v1.VMStatusFailed + }), "failed to start a VM") + + require.Equal(t, v1.VMStatusRunning, vm.Status) + require.Empty(t, vm.StatusMessage) + + // Establish port forwarding to VMs SSH port + wsConn, err := devClient.VMs().PortForward(ctx, "test-vm", 22) + require.NoError(t, err) + + // Make sure we can connect to the VM over SSH via the forwarded port + sshConfig := &ssh.ClientConfig{ + HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error { + return nil + }, + User: "admin", + Auth: []ssh.AuthMethod{ + ssh.Password("admin"), + }, + } + + sshConn, chans, reqs, err := ssh.NewClientConn(wsConn, "", sshConfig) + require.NoError(t, err) + + sshClient := ssh.NewClient(sshConn, chans, reqs) + + sshSession, err := sshClient.NewSession() + require.NoError(t, err) + + unameOutput, err := sshSession.Output("uname -mo") + require.NoError(t, err) + require.Contains(t, string(unameOutput), "Darwin arm64") +} diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 1436486..daeb7cb 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -63,63 +63,51 @@ func New(opts ...Option) (*Worker, error) { } func (worker *Worker) Run(ctx context.Context) error { - tickCh := time.NewTicker(pollInterval) + for { + if err := worker.runNewSession(ctx); err != nil { + return err + } + } +} + +func (worker *Worker) runNewSession(ctx context.Context) error { + subCtx, cancel := context.WithCancel(ctx) + defer cancel() + + if err := worker.registerWorker(subCtx); err != nil { + worker.logger.Warnf("failed to register worker: %v", err) + + return nil + } go func() { _ = retry.Do(func() error { - return worker.watchRPC(ctx) + return worker.watchRPC(subCtx) }, retry.OnRetry(func(n uint, err error) { worker.logger.Warnf("failed to watch RPC: %v", err) - }), retry.Context(ctx), retry.Attempts(0)) + }), retry.Context(subCtx), retry.Attempts(0)) }() + tickCh := time.NewTicker(pollInterval) + for { - if err := worker.registerWorker(ctx); err != nil { - worker.logger.Warnf("failed to register worker: %v", err) + if err := worker.updateWorker(ctx); err != nil { + worker.logger.Errorf("failed to update worker resource: %v", err) - select { - case <-tickCh.C: - // continue - case <-ctx.Done(): - return ctx.Err() - } - - continue + return nil } - for { - if err := worker.updateWorker(ctx); err != nil { - worker.logger.Errorf("failed to update worker resource: %v", err) + if err := worker.syncVMs(subCtx); err != nil { + worker.logger.Warnf("failed to sync VMs: %v", err) - select { - case <-tickCh.C: - // continue - case <-ctx.Done(): - return ctx.Err() - } + return nil + } - break - } - - if err := worker.syncVMs(ctx); err != nil { - worker.logger.Warnf("failed to sync VMs: %v", err) - - select { - case <-tickCh.C: - // continue - case <-ctx.Done(): - return ctx.Err() - } - - break - } - - select { - case <-tickCh.C: - // continue - case <-ctx.Done(): - return ctx.Err() - } + select { + case <-tickCh.C: + // continue + case <-subCtx.Done(): + return subCtx.Err() } } }