coder · coadler · Jun 24, 2022 · Jun 10, 2022 · Jun 22, 2022 · Jun 22, 2022
diff --git a/Makefile b/Makefile
@@ -57,7 +57,7 @@ coderd/database/dump.sql: $(wildcard coderd/database/migrations/*.sql)
 	go run coderd/database/dump/main.go
 
 # Generates Go code for querying the database.
-coderd/database/querier.go: coderd/database/dump.sql $(wildcard coderd/database/queries/*.sql)
+coderd/database/querier.go: coderd/database/sqlc.yaml coderd/database/dump.sql $(wildcard coderd/database/queries/*.sql)
 	coderd/database/generate.sh
 
 # This target is deprecated, as GNU make has issues passing signals to subprocesses.

diff --git a/agent/agent.go b/agent/agent.go
@@ -27,10 +27,13 @@ import (
 	"go.uber.org/atomic"
 	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/xerrors"
+	"inet.af/netaddr"
+	"tailscale.com/types/key"
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/agent/usershell"
 	"github.com/coder/coder/peer"
+	"github.com/coder/coder/peer/peerwg"
 	"github.com/coder/coder/peerbroker"
 	"github.com/coder/coder/pty"
 	"github.com/coder/retry"
@@ -43,20 +46,31 @@ const (
 )
 
 type Options struct {
+	EnableWireguard        bool
+	PostPublicKeys         PostKeys
+	ListenWireguardPeers   ListenWireguardPeers
 	ReconnectingPTYTimeout time.Duration
 	EnvironmentVariables   map[string]string
 	Logger                 slog.Logger
 }
 
 type Metadata struct {
-	OwnerEmail           string            `json:"owner_email"`
-	OwnerUsername        string            `json:"owner_username"`
-	EnvironmentVariables map[string]string `json:"environment_variables"`
-	StartupScript        string            `json:"startup_script"`
-	Directory            string            `json:"directory"`
+	Addresses            []netaddr.IPPrefix `json:"addresses"`
+	OwnerEmail           string             `json:"owner_email"`
+	OwnerUsername        string             `json:"owner_username"`
+	EnvironmentVariables map[string]string  `json:"environment_variables"`
+	StartupScript        string             `json:"startup_script"`
+	Directory            string             `json:"directory"`
+}
+
+type PublicKeys struct {
+	Public key.NodePublic  `json:"public"`
+	Disco  key.DiscoPublic `json:"disco"`
 }
 
 type Dialer func(ctx context.Context, logger slog.Logger) (Metadata, *peerbroker.Listener, error)
+type PostKeys func(ctx context.Context, keys PublicKeys) error
+type ListenWireguardPeers func(ctx context.Context, logger slog.Logger) (<-chan *peerwg.WireguardPeerMessage, func(), error)
 
 func New(dialer Dialer, options *Options) io.Closer {
 	if options == nil {
@@ -73,6 +87,9 @@ func New(dialer Dialer, options *Options) io.Closer {
 		closeCancel:            cancelFunc,
 		closed:                 make(chan struct{}),
 		envVars:                options.EnvironmentVariables,
+		enableWireguard:        options.EnableWireguard,
+		postKeys:               options.PostPublicKeys,
+		listenWireguardPeers:   options.ListenWireguardPeers,
 	}
 	server.init(ctx)
 	return server
@@ -95,6 +112,11 @@ type agent struct {
 	metadata      atomic.Value
 	startupScript atomic.Bool
 	sshServer     *ssh.Server
+
+	enableWireguard      bool
+	wg                   *peerwg.WireguardNetwork
+	postKeys             PostKeys
+	listenWireguardPeers ListenWireguardPeers
 }
 
 func (a *agent) run(ctx context.Context) {
@@ -138,6 +160,11 @@ func (a *agent) run(ctx context.Context) {
 		}()
 	}
 
+	err = a.startWireguard(ctx, metadata.Addresses)
+	if err != nil {
+		a.logger.Error(ctx, "start wireguard", slog.Error(err))
+	}
+
 	for {
 		conn, err := peerListener.Accept()
 		if err != nil {
@@ -366,17 +393,17 @@ func (a *agent) createCommand(ctx context.Context, rawCommand string, env []stri
 
 	// Load environment variables passed via the agent.
 	// These should override all variables we manually specify.
-	for key, value := range metadata.EnvironmentVariables {
+	for envKey, value := range metadata.EnvironmentVariables {
 		// Expanding environment variables allows for customization
 		// of the $PATH, among other variables. Customers can prepand
 		// or append to the $PATH, so allowing expand is required!
-		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", key, os.ExpandEnv(value)))
+		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", envKey, os.ExpandEnv(value)))
 	}
 
 	// Agent-level environment variables should take over all!
 	// This is used for setting agent-specific variables like "CODER_AGENT_TOKEN".
-	for key, value := range a.envVars {
-		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", key, value))
+	for envKey, value := range a.envVars {
+		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", envKey, value))
 	}
 
 	return cmd, nil

diff --git a/agent/wireguard.go b/agent/wireguard.go
@@ -0,0 +1,64 @@
+package agent
+
+import (
+	"context"
+
+	"golang.org/x/xerrors"
+	"inet.af/netaddr"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/peer/peerwg"
+)
+
+func (a *agent) startWireguard(ctx context.Context, addrs []netaddr.IPPrefix) error {
+	if a.wg != nil {
+		_ = a.wg.Close()
+	}
+
+	if !a.enableWireguard {
+		return nil
+	}
+
+	// We can't create a wireguard network without these.
+	if len(addrs) == 0 || a.listenWireguardPeers == nil || a.postKeys == nil {
+		return xerrors.New("wireguard is enabled, but no addresses were provided or necessary functions were not provided")
+	}
+
+	wg, err := peerwg.NewWireguardNetwork(ctx, a.logger.Named("wireguard"), addrs)
+	if err != nil {
+		return xerrors.Errorf("create wireguard network: %w", err)
+	}
+
+	err = a.postKeys(ctx, PublicKeys{
+		Public: wg.Private.Public(),
+		Disco:  wg.Disco,
+	})
+	if err != nil {
+		a.logger.Warn(ctx, "post keys", slog.Error(err))
+	}
+
+	go func() {
+		for {
+			ch, listenClose, err := a.listenWireguardPeers(ctx, a.logger)
+			if err != nil {
+				a.logger.Warn(ctx, "listen wireguard peers", slog.Error(err))
+				return
+			}
+
+			for {
+				peer := <-ch
+				if peer == nil {
+					break
+				}
+
+				err := wg.AddPeer(*peer)
+				a.logger.Info(ctx, "added wireguard peer", slog.F("peer", peer.Public.ShortString()), slog.Error(err))
+			}
+
+			listenClose()
+		}
+	}()
+
+	a.wg = wg
+	return nil
+}
diff --git a/cli/agent.go b/cli/agent.go
@@ -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,
+				PostPublicKeys:       client.PostWorkspaceAgentKeys,
+				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
 }
diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -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)

diff --git a/cli/root.go b/cli/root.go
@@ -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())