Introduce "orchard {port-forward, vnc} worker WORKER_NAME" (#140)

* Fix potential NPE in Client.wsRequest()

* Introduce "orchard {port-forward, vnc} worker WORKER_NAME"

* portspec.go: simplify logic and respect [LOCAL_PORT]:REMOTE_PORT format

internal/command/portforward/portforward.go

@@ -8,7 +8,7 @@ func NewCommand() *cobra.Command {
 		Short: "Forward TCP port to the resources",
 	}
 
-	command.AddCommand(newPortForwardVMCommand())
+	command.AddCommand(newPortForwardVMCommand(), newPortForwardWorkerCommand())
 
 	return command
 }

internal/command/portforward/portspec.go

@@ -17,31 +17,47 @@ type PortSpec struct {
 func NewPortSpec(portSpecRaw string) (*PortSpec, error) {
 	splits := strings.Split(portSpecRaw, ":")
 
-	if len(splits) > 2 {
-		return nil, fmt.Errorf("%w: expected no more than 2 components delimited by \":\", found %d",
-			ErrInvalidPortSpec, len(splits))
-	}
-
-	localPort, err := strconv.ParseUint(splits[0], 10, 16)
-	if err != nil {
-		return nil, err
-	}
-
-	remotePort := localPort
-
-	if len(splits) > 1 {
-		remotePort, err = strconv.ParseUint(splits[1], 10, 16)
+	switch len(splits) {
+	case 1:
+		remotePort, err := parsePort(splits[0])
 		if err != nil {
 			return nil, err
 		}
-	}
 
-	if localPort < 1 || remotePort < 1 || localPort > 65535 || remotePort > 65535 {
-		return nil, fmt.Errorf("%w: only ports in range [1, 65535] are allowed", ErrInvalidPortSpec)
-	}
+		return &PortSpec{
+			LocalPort:  0,
+			RemotePort: remotePort,
+		}, nil
+	case 2:
+		localPort, err := parsePort(splits[0])
+		if err != nil {
+			return nil, err
+		}
 
-	return &PortSpec{
-		LocalPort:  uint16(localPort),
-		RemotePort: uint16(remotePort),
-	}, nil
+		remotePort, err := parsePort(splits[1])
+		if err != nil {
+			return nil, err
+		}
+
+		return &PortSpec{
+			LocalPort:  localPort,
+			RemotePort: remotePort,
+		}, nil
+	default:
+		return nil, fmt.Errorf("%w: expected 1 or 2 components delimited by \":\", found %d",
+			ErrInvalidPortSpec, len(splits))
+	}
+}
+
+func parsePort(s string) (uint16, error) {
+	port, err := strconv.ParseUint(s, 10, 16)
+	if err != nil {
+		return 0, err
+	}
+
+	if port < 1 || port > 65535 {
+		return 0, fmt.Errorf("%w: only ports in range [1, 65535] are allowed", ErrInvalidPortSpec)
+	}
+
+	return uint16(port), nil
 }

internal/command/portforward/portspec_test.go

@@ -9,7 +9,7 @@ import (
 func TestPortSpecNormal(t *testing.T) {
 	portSpec, err := portforward.NewPortSpec("5555")
 	require.NoError(t, err)
-	require.Equal(t, &portforward.PortSpec{LocalPort: 5555, RemotePort: 5555}, portSpec)
+	require.Equal(t, &portforward.PortSpec{LocalPort: 0, RemotePort: 5555}, portSpec)
 
 	portSpec, err = portforward.NewPortSpec("8000:80")
 	require.NoError(t, err)

internal/command/portforward/worker.go

@@ -0,0 +1,83 @@
+package portforward
+
+import (
+	"fmt"
+	"github.com/cirruslabs/orchard/internal/proxy"
+	"github.com/cirruslabs/orchard/pkg/client"
+	"github.com/spf13/cobra"
+	"net"
+)
+
+func newPortForwardWorkerCommand() *cobra.Command {
+	command := &cobra.Command{
+		Use:   "worker WORKER_NAME [LOCAL_PORT]:REMOTE_PORT",
+		Short: "Forward TCP port to the worker",
+		Args:  cobra.ExactArgs(2),
+		RunE:  runPortForwardWorkerCommand,
+	}
+
+	return command
+}
+
+func runPortForwardWorkerCommand(cmd *cobra.Command, args []string) (err error) {
+	name := args[0]
+	portSpecRaw := args[1]
+
+	portSpec, err := NewPortSpec(portSpecRaw)
+	if err != nil {
+		return err
+	}
+
+	client, err := client.New()
+	if err != nil {
+		return err
+	}
+
+	listener, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", portSpec.LocalPort))
+	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, portSpec.RemotePort)
+
+	errCh := make(chan error, 1)
+
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				errCh <- err
+
+				return
+			}
+
+			go func() {
+				defer conn.Close()
+
+				wsConn, err := client.Workers().PortForward(cmd.Context(), name, portSpec.RemotePort)
+				if err != nil {
+					fmt.Printf("failed to forward port: %v\n", err)
+
+					return
+				}
+				defer wsConn.Close()
+
+				if err := proxy.Connections(wsConn, conn); err != nil {
+					fmt.Printf("failed to forward port: %v\n", err)
+				}
+			}()
+		}
+	}()
+
+	select {
+	case <-cmd.Context().Done():
+		return cmd.Context().Err()
+	case err := <-errCh:
+		return err
+	}
+}

internal/command/vnc/vnc.go

@@ -8,7 +8,7 @@ func NewCommand() *cobra.Command {
 		Short: "Open VNC session with the resource",
 	}
 
-	command.AddCommand(newVNCVMCommand())
+	command.AddCommand(newVNCVMCommand(), newVNCWorkerCommand())
 
 	return command
 }

internal/command/vnc/worker.go

@@ -0,0 +1,104 @@
+package vnc
+
+import (
+	"fmt"
+	"github.com/cirruslabs/orchard/internal/proxy"
+	"github.com/cirruslabs/orchard/pkg/client"
+	"github.com/skratchdot/open-golang/open"
+	"github.com/spf13/cobra"
+	"net"
+)
+
+func newVNCWorkerCommand() *cobra.Command {
+	command := &cobra.Command{
+		Use:   "worker WORKER_NAME",
+		Short: "Open VNC session with the worker",
+		Args:  cobra.ExactArgs(1),
+		RunE:  runVNCWorker,
+	}
+
+	command.PersistentFlags().StringVarP(&username, "username", "u", "",
+		"VNC username")
+	command.PersistentFlags().StringVarP(&password, "password", "p", "",
+		"VNC password")
+	return command
+}
+
+func runVNCWorker(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.Workers().PortForward(cmd.Context(), name, vncPort)
+				if err != nil {
+					fmt.Printf("failed to forward port: %v\n", err)
+
+					return
+				}
+				defer wsConn.Close()
+
+				if err := proxy.Connections(wsConn, conn); err != nil {
+					fmt.Printf("failed to forward port: %v\n", err)
+				}
+			}()
+		}
+	}()
+
+	var credentialsComponent string
+	var credentialsComponentSanitized string
+
+	if username != "" {
+		credentialsComponent = username
+		credentialsComponentSanitized = username
+
+		if password != "" {
+			credentialsComponent += ":" + password
+		}
+
+		credentialsComponent += "@"
+		credentialsComponentSanitized += "@"
+	}
+
+	openURL := fmt.Sprintf("vnc://%s%s", credentialsComponent, listener.Addr().String())
+	openURLSanitized := fmt.Sprintf("vnc://%s%s", credentialsComponentSanitized, 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
+	}
+}