feat: peer wireguard

coder · coadler · Jun 24, 2022 · Jun 10, 2022 · Jun 22, 2022 · Jun 22, 2022
commit 2361e15d44367d5ddde98ed21dab14babf982bec
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
@@ -178,6 +178,9 @@ func workspaceAgent() *cobra.Command {
 					// shells so "gitssh" works!
 					"CODER_AGENT_TOKEN": client.SessionToken,
 				},
+				EnableWireguard:      true,
+				PostPublicKeys:       client.PostWorkspaceAgentKeys,
+				ListenWireguardPeers: client.WireguardPeerListener,
 			})
 			<-cmd.Context().Done()
 			return closer.Close()

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())