feat: Add wireguard to port-forward

kylecarbs · kylecarbs · commit dc6c8d07f9ff · 2022-09-02T16:29:06.000-05:00
This allows replacement of the WebRTC networking!
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -12,10 +12,12 @@ import (
 	"sync"
 	"syscall"
 
+	"cdr.dev/slog"
 	"github.com/pion/udp"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/agent"
 	coderagent "github.com/coder/coder/agent"
 	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/codersdk"
@@ -26,6 +28,7 @@ func portForward() *cobra.Command {
 		tcpForwards  []string // <port>:<port>
 		udpForwards  []string // <port>:<port>
 		unixForwards []string // <path>:<path> OR <port>:<path>
+		wireguard    bool
 	)
 	cmd := &cobra.Command{
 		Use:     "port-forward <workspace>",
@@ -75,7 +78,7 @@ func portForward() *cobra.Command {
 				return err
 			}
 
-			workspace, agent, err := getWorkspaceAndAgent(ctx, cmd, client, codersdk.Me, args[0], false)
+			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, cmd, client, codersdk.Me, args[0], false)
 			if err != nil {
 				return err
 			}
@@ -92,16 +95,21 @@ func portForward() *cobra.Command {
 			err = cliui.Agent(ctx, cmd.ErrOrStderr(), cliui.AgentOptions{
 				WorkspaceName: workspace.Name,
 				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
-					return client.WorkspaceAgent(ctx, agent.ID)
+					return client.WorkspaceAgent(ctx, workspaceAgent.ID)
 				},
 			})
 			if err != nil {
 				return xerrors.Errorf("await agent: %w", err)
 			}
 
-			conn, err := client.DialWorkspaceAgent(ctx, agent.ID, nil)
+			var conn agent.Conn
+			if !wireguard {
+				conn, err = client.DialWorkspaceAgent(ctx, workspaceAgent.ID, nil)
+			} else {
+				conn, err = client.DialWorkspaceAgentTailnet(ctx, slog.Logger{}, workspaceAgent.ID)
+			}
 			if err != nil {
-				return xerrors.Errorf("dial workspace agent: %w", err)
+				return err
 			}
 			defer conn.Close()
 
@@ -159,7 +167,8 @@ func portForward() *cobra.Command {
 	cmd.Flags().StringArrayVarP(&tcpForwards, "tcp", "p", []string{}, "Forward a TCP port from the workspace to the local machine")
 	cmd.Flags().StringArrayVar(&udpForwards, "udp", []string{}, "Forward a UDP port from the workspace to the local machine. The UDP connection has TCP-like semantics to support stateful UDP protocols")
 	cmd.Flags().StringArrayVar(&unixForwards, "unix", []string{}, "Forward a Unix socket in the workspace to a local Unix socket or TCP port")
-
+	cmd.Flags().BoolVarP(&wireguard, "wireguard", "", false, "Specifies whether to use wireguard networking or not.")
+	_ = cmd.Flags().MarkHidden("wireguard")
 	return cmd
 }
 
