From 97a014b156c73678596eb35b3515b783f78ff105 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Fri, 2 Sep 2022 16:29:06 -0500
Subject: [PATCH] feat: Add wireguard to port-forward

This allows replacement of the WebRTC networking!
---
 cli/portforward.go | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/cli/portforward.go b/cli/portforward.go
index 66753ac92f7d1..3b78fa8d11a65 100644
--- a/cli/portforward.go
+++ b/cli/portforward.go
@@ -12,11 +12,12 @@ import (
 	"sync"
 	"syscall"
 
+	"cdr.dev/slog"
 	"github.com/pion/udp"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 
-	coderagent "github.com/coder/coder/agent"
+	"github.com/coder/coder/agent"
 	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/codersdk"
 )
@@ -26,6 +27,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 +77,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 +94,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,11 +166,12 @@ 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
 }
 
-func listenAndPortForward(ctx context.Context, cmd *cobra.Command, conn coderagent.Conn, wg *sync.WaitGroup, spec portForwardSpec) (net.Listener, error) {
+func listenAndPortForward(ctx context.Context, cmd *cobra.Command, conn agent.Conn, wg *sync.WaitGroup, spec portForwardSpec) (net.Listener, error) {
 	_, _ = fmt.Fprintf(cmd.OutOrStderr(), "Forwarding '%v://%v' locally to '%v://%v' in the workspace\n", spec.listenNetwork, spec.listenAddress, spec.dialNetwork, spec.dialAddress)
 
 	var (
@@ -219,7 +227,7 @@ func listenAndPortForward(ctx context.Context, cmd *cobra.Command, conn coderage
 				}
 				defer remoteConn.Close()
 
-				coderagent.Bicopy(ctx, netConn, remoteConn)
+				agent.Bicopy(ctx, netConn, remoteConn)
 			}(netConn)
 		}
 	}(spec)
@@ -315,7 +323,7 @@ func parsePort(in string) (uint16, error) {
 }
 
 func parseUnixPath(in string) (string, error) {
-	path, err := coderagent.ExpandRelativeHomePath(strings.TrimSpace(in))
+	path, err := agent.ExpandRelativeHomePath(strings.TrimSpace(in))
 	if err != nil {
 		return "", xerrors.Errorf("tidy path %q: %w", in, err)
 	}