feat: peer wireguard (#2445)

2026-09-24 15:04:27 +08:00 · 2022-06-24 10:25:01 -05:00
parent d21ab2115d
commit 05b67ab1cf
34 changed files with 1935 additions and 236 deletions
@@ -14,17 +14,15 @@ import (
 	"cloud.google.com/go/compute/metadata"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
+	"gopkg.in/natefinch/lumberjack.v2"

 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
-
 	"github.com/coder/coder/agent"
 	"github.com/coder/coder/agent/reaper"
 	"github.com/coder/coder/cli/cliflag"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/retry"
-
-	"gopkg.in/natefinch/lumberjack.v2"
 )

 func workspaceAgent() *cobra.Command {
@@ -33,6 +31,7 @@ func workspaceAgent() *cobra.Command {
 		pprofEnabled bool
 		pprofAddress string
 		noReap       bool
+		wireguard    bool
 	)
 	cmd := &cobra.Command{
 		Use: "agent",
@@ -178,6 +177,9 @@ func workspaceAgent() *cobra.Command {
 					// shells so "gitssh" works!
 					"CODER_AGENT_TOKEN": client.SessionToken,
 				},
+				EnableWireguard:      wireguard,
+				UploadWireguardKeys:  client.UploadWorkspaceAgentKeys,
+				ListenWireguardPeers: client.WireguardPeerListener,
 			})
 			<-cmd.Context().Done()
 			return closer.Close()
@@ -188,5 +190,6 @@ func workspaceAgent() *cobra.Command {
 	cliflag.BoolVarP(cmd.Flags(), &pprofEnabled, "pprof-enable", "", "CODER_AGENT_PPROF_ENABLE", false, "Enable serving pprof metrics on the address defined by --pprof-address.")
 	cliflag.BoolVarP(cmd.Flags(), &noReap, "no-reap", "", "", false, "Do not start a process reaper.")
 	cliflag.StringVarP(cmd.Flags(), &pprofAddress, "pprof-address", "", "CODER_AGENT_PPROF_ADDRESS", "127.0.0.1:6060", "The address to serve pprof.")
+	cliflag.BoolVarP(cmd.Flags(), &wireguard, "wireguard", "", "CODER_AGENT_WIREGUARD", true, "Whether to start the Wireguard interface.")
 	return cmd
 }
@@ -46,7 +46,7 @@ func TestWorkspaceAgent(t *testing.T) {
 		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
 		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)

-		cmd, _ := clitest.New(t, "agent", "--auth", "azure-instance-identity", "--agent-url", client.URL.String())
+		cmd, _ := clitest.New(t, "agent", "--auth", "azure-instance-identity", "--agent-url", client.URL.String(), "--wireguard=false")
 		ctx, cancelFunc := context.WithCancel(context.Background())
 		defer cancelFunc()
 		errC := make(chan error)
@@ -101,7 +101,7 @@ func TestWorkspaceAgent(t *testing.T) {
 		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
 		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)

-		cmd, _ := clitest.New(t, "agent", "--auth", "aws-instance-identity", "--agent-url", client.URL.String())
+		cmd, _ := clitest.New(t, "agent", "--auth", "aws-instance-identity", "--agent-url", client.URL.String(), "--wireguard=false")
 		ctx, cancelFunc := context.WithCancel(context.Background())
 		defer cancelFunc()
 		errC := make(chan error)
@@ -156,7 +156,7 @@ func TestWorkspaceAgent(t *testing.T) {
 		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
 		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)

-		cmd, _ := clitest.New(t, "agent", "--auth", "google-instance-identity", "--agent-url", client.URL.String())
+		cmd, _ := clitest.New(t, "agent", "--auth", "google-instance-identity", "--agent-url", client.URL.String(), "--wireguard=false")
 		ctx, cancelFunc := context.WithCancel(context.Background())
 		defer cancelFunc()
 		errC := make(chan error)
@@ -73,22 +73,23 @@ func Root() *cobra.Command {
 		list(),
 		login(),
 		logout(),
+		parameters(),
+		portForward(),
 		publickey(),
 		resetPassword(),
 		schedules(),
 		server(),
 		show(),
+		ssh(),
 		start(),
 		state(),
 		stop(),
-		ssh(),
 		templates(),
 		update(),
 		users(),
-		portForward(),
-		workspaceAgent(),
 		versionCmd(),
-		parameters(),
+		wireguardPortForward(),
+		workspaceAgent(),
 	)

 	cmd.SetUsageTemplate(usageTemplate())
@@ -0,0 +1,258 @@
+package cli
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"os"
+	"os/signal"
+	"strconv"
+	"sync"
+	"syscall"
+
+	"github.com/google/uuid"
+	"github.com/pion/udp"
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+	"inet.af/netaddr"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/sloghuman"
+	coderagent "github.com/coder/coder/agent"
+	"github.com/coder/coder/cli/cliui"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/peer/peerwg"
+)
+
+func wireguardPortForward() *cobra.Command {
+	var (
+		tcpForwards []string // <port>:<port>
+		udpForwards []string // <port>:<port>
+		// TODO: unix support
+		// unixForwards []string // <path>:<path> OR <port>:<path>
+	)
+	cmd := &cobra.Command{
+		Use:     "wireguard-port-forward <workspace>",
+		Aliases: []string{"wireguard-tunnel"},
+		Args:    cobra.ExactArgs(1),
+		// Hide all wireguard commands for now while we test!
+		Hidden: true,
+		Example: `
+  - Port forward a single TCP port from 1234 in the workspace to port 5678 on
+    your local machine
+
+    ` + cliui.Styles.Code.Render("$ coder port-forward <workspace> --tcp 5678:1234") + `
+
+  - Port forward a single UDP port from port 9000 to port 9000 on your local
+    machine
+
+    ` + cliui.Styles.Code.Render("$ coder port-forward <workspace> --udp 9000") + `
+
+  - Port forward multiple TCP ports and a UDP port
+
+    ` + cliui.Styles.Code.Render("$ coder port-forward <workspace> --tcp 8080:8080 --tcp 9000:3000 --udp 5353:53") + `
+`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			specs, err := parsePortForwards(tcpForwards, nil, nil)
+			if err != nil {
+				return xerrors.Errorf("parse port-forward specs: %w", err)
+			}
+			if len(specs) == 0 {
+				err = cmd.Help()
+				if err != nil {
+					return xerrors.Errorf("generate help output: %w", err)
+				}
+				return xerrors.New("no port-forwards requested")
+			}
+
+			client, err := createClient(cmd)
+			if err != nil {
+				return err
+			}
+
+			workspace, workspaceAgent, err := getWorkspaceAndAgent(cmd, client, codersdk.Me, args[0], false)
+			if err != nil {
+				return err
+			}
+			if workspace.LatestBuild.Transition != codersdk.WorkspaceTransitionStart {
+				return xerrors.New("workspace must be in start transition to port-forward")
+			}
+			if workspace.LatestBuild.Job.CompletedAt == nil {
+				err = cliui.WorkspaceBuild(cmd.Context(), cmd.ErrOrStderr(), client, workspace.LatestBuild.ID, workspace.CreatedAt)
+				if err != nil {
+					return err
+				}
+			}
+
+			err = cliui.Agent(cmd.Context(), cmd.ErrOrStderr(), cliui.AgentOptions{
+				WorkspaceName: workspace.Name,
+				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
+					return client.WorkspaceAgent(ctx, workspaceAgent.ID)
+				},
+			})
+			if err != nil {
+				return xerrors.Errorf("await agent: %w", err)
+			}
+
+			ipv6 := peerwg.UUIDToNetaddr(uuid.New())
+			wgn, err := peerwg.New(
+				slog.Make(sloghuman.Sink(os.Stderr)),
+				[]netaddr.IPPrefix{netaddr.IPPrefixFrom(ipv6, 128)},
+			)
+			if err != nil {
+				return xerrors.Errorf("create wireguard network: %w", err)
+			}
+
+			err = client.PostWireguardPeer(cmd.Context(), workspace.ID, peerwg.Handshake{
+				Recipient:      workspaceAgent.ID,
+				NodePublicKey:  wgn.NodePrivateKey.Public(),
+				DiscoPublicKey: wgn.DiscoPublicKey,
+				IPv6:           ipv6,
+			})
+			if err != nil {
+				return xerrors.Errorf("post wireguard peer: %w", err)
+			}
+
+			err = wgn.AddPeer(peerwg.Handshake{
+				Recipient:      workspaceAgent.ID,
+				DiscoPublicKey: workspaceAgent.DiscoPublicKey,
+				NodePublicKey:  workspaceAgent.WireguardPublicKey,
+				IPv6:           workspaceAgent.IPv6.IP(),
+			})
+			if err != nil {
+				return xerrors.Errorf("add workspace agent as peer: %w", err)
+			}
+
+			// Start all listeners.
+			var (
+				ctx, cancel       = context.WithCancel(cmd.Context())
+				wg                = new(sync.WaitGroup)
+				listeners         = make([]net.Listener, len(specs))
+				closeAllListeners = func() {
+					for _, l := range listeners {
+						if l == nil {
+							continue
+						}
+						_ = l.Close()
+					}
+				}
+			)
+			defer cancel()
+			for i, spec := range specs {
+				l, err := listenAndPortForwardWireguard(ctx, cmd, wgn, wg, spec, workspaceAgent.IPv6.IP())
+				if err != nil {
+					closeAllListeners()
+					return err
+				}
+				listeners[i] = l
+			}
+
+			// Wait for the context to be canceled or for a signal and close
+			// all listeners.
+			var closeErr error
+			go func() {
+				sigs := make(chan os.Signal, 1)
+				signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+
+				select {
+				case <-ctx.Done():
+					closeErr = ctx.Err()
+				case <-sigs:
+					_, _ = fmt.Fprintln(cmd.OutOrStderr(), "Received signal, closing all listeners and active connections")
+					closeErr = xerrors.New("signal received")
+				}
+
+				cancel()
+				closeAllListeners()
+			}()
+
+			_, _ = fmt.Fprintln(cmd.OutOrStderr(), "Ready!")
+			wg.Wait()
+			return closeErr
+		},
+	}
+
+	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")
+
+	return cmd
+}
+
+func listenAndPortForwardWireguard(ctx context.Context, cmd *cobra.Command,
+	wgn *peerwg.Network,
+	wg *sync.WaitGroup,
+	spec portForwardSpec,
+	agentIP netaddr.IP,
+) (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 (
+		l   net.Listener
+		err error
+	)
+	switch spec.listenNetwork {
+	case "tcp":
+		l, err = net.Listen(spec.listenNetwork, spec.listenAddress)
+	case "udp":
+		var host, port string
+		host, port, err = net.SplitHostPort(spec.listenAddress)
+		if err != nil {
+			return nil, xerrors.Errorf("split %q: %w", spec.listenAddress, err)
+		}
+
+		var portInt int
+		portInt, err = strconv.Atoi(port)
+		if err != nil {
+			return nil, xerrors.Errorf("parse port %v from %q as int: %w", port, spec.listenAddress, err)
+		}
+
+		l, err = udp.Listen(spec.listenNetwork, &net.UDPAddr{
+			IP:   net.ParseIP(host),
+			Port: portInt,
+		})
+	// case "unix":
+	// 	l, err = net.Listen(spec.listenNetwork, spec.listenAddress)
+	default:
+		return nil, xerrors.Errorf("unknown listen network %q", spec.listenNetwork)
+	}
+	if err != nil {
+		return nil, xerrors.Errorf("listen '%v://%v': %w", spec.listenNetwork, spec.listenAddress, err)
+	}
+
+	wg.Add(1)
+	go func(spec portForwardSpec) {
+		defer wg.Done()
+		for {
+			netConn, err := l.Accept()
+			if err != nil {
+				_, _ = fmt.Fprintf(cmd.OutOrStderr(), "Error accepting connection from '%v://%v': %+v\n", spec.listenNetwork, spec.listenAddress, err)
+				_, _ = fmt.Fprintln(cmd.OutOrStderr(), "Killing listener")
+				return
+			}
+
+			go func(netConn net.Conn) {
+				defer netConn.Close()
+
+				ipPort := netaddr.MustParseIPPort(spec.dialAddress).WithIP(agentIP)
+
+				var remoteConn net.Conn
+				switch spec.dialNetwork {
+				case "tcp":
+					remoteConn, err = wgn.Netstack.DialContextTCP(ctx, ipPort)
+				case "udp":
+					remoteConn, err = wgn.Netstack.DialContextUDP(ctx, ipPort)
+				}
+				if err != nil {
+					_, _ = fmt.Fprintf(cmd.OutOrStderr(), "Failed to dial '%v://%v' in workspace: %s\n", spec.dialNetwork, spec.dialAddress, err)
+					return
+				}
+				defer remoteConn.Close()
+
+				coderagent.Bicopy(ctx, netConn, remoteConn)
+			}(netConn)
+		}
+	}(spec)
+
+	return l, nil
+}