feat: peer wireguard (#2445)

codersdk/workspaceagents.go

@@ -23,6 +23,7 @@ import (
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/turnconn"
 	"github.com/coder/coder/peer"
+	"github.com/coder/coder/peer/peerwg"
 	"github.com/coder/coder/peerbroker"
 	"github.com/coder/coder/peerbroker/proto"
 	"github.com/coder/coder/provisionersdk"
@@ -252,6 +253,97 @@ func (c *Client) ListenWorkspaceAgent(ctx context.Context, logger slog.Logger) (
 	return agentMetadata, listener, json.NewDecoder(res.Body).Decode(&agentMetadata)
 }
 
+// PostWireguardPeer announces your public keys and IPv6 address to the
+// specified recipient.
+func (c *Client) PostWireguardPeer(ctx context.Context, workspaceID uuid.UUID, peerMsg peerwg.Handshake) error {
+	res, err := c.Request(ctx, http.MethodPost, fmt.Sprintf("/api/v2/workspaceagents/%s/peer?workspace=%s",
+		peerMsg.Recipient,
+		workspaceID.String(),
+	), peerMsg)
+	if err != nil {
+		return err
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusNoContent {
+		return readBodyAsError(res)
+	}
+
+	_, _ = io.Copy(io.Discard, res.Body)
+	return nil
+}
+
+// WireguardPeerListener listens for wireguard peer messages. Peer messages are
+// sent when a new client wants to connect. Once receiving a peer message, the
+// peer should be added to the NetworkMap of the wireguard interface.
+func (c *Client) WireguardPeerListener(ctx context.Context, logger slog.Logger) (<-chan peerwg.Handshake, func(), error) {
+	serverURL, err := c.URL.Parse("/api/v2/workspaceagents/me/wireguardlisten")
+	if err != nil {
+		return nil, nil, xerrors.Errorf("parse url: %w", err)
+	}
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		return nil, nil, xerrors.Errorf("create cookie jar: %w", err)
+	}
+	jar.SetCookies(serverURL, []*http.Cookie{{
+		Name:  httpmw.SessionTokenKey,
+		Value: c.SessionToken,
+	}})
+	httpClient := &http.Client{
+		Jar: jar,
+	}
+
+	conn, res, err := websocket.Dial(ctx, serverURL.String(), &websocket.DialOptions{
+		HTTPClient: httpClient,
+		// Need to disable compression to avoid a data-race.
+		CompressionMode: websocket.CompressionDisabled,
+	})
+	if err != nil {
+		if res == nil {
+			return nil, nil, xerrors.Errorf("websocket dial: %w", err)
+		}
+		return nil, nil, readBodyAsError(res)
+	}
+
+	ch := make(chan peerwg.Handshake, 1)
+	go func() {
+		defer conn.Close(websocket.StatusGoingAway, "")
+		defer close(ch)
+
+		for {
+			_, message, err := conn.Read(ctx)
+			if err != nil {
+				break
+			}
+
+			var msg peerwg.Handshake
+			err = msg.UnmarshalText(message)
+			if err != nil {
+				logger.Error(ctx, "unmarshal wireguard peer message", slog.Error(err))
+				continue
+			}
+
+			ch <- msg
+		}
+	}()
+
+	return ch, func() { _ = conn.Close(websocket.StatusGoingAway, "") }, nil
+}
+
+// UploadWorkspaceAgentKeys uploads the public keys of the workspace agent that
+// were generated on startup. These keys are used by clients to communicate with
+// the workspace agent over the wireguard interface.
+func (c *Client) UploadWorkspaceAgentKeys(ctx context.Context, keys agent.WireguardPublicKeys) error {
+	res, err := c.Request(ctx, http.MethodPost, "/api/v2/workspaceagents/me/keys", keys)
+	if err != nil {
+		return xerrors.Errorf("do request: %w", err)
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusNoContent {
+		return readBodyAsError(res)
+	}
+	return nil
+}
+
 // DialWorkspaceAgent creates a connection to the specified resource.
 func (c *Client) DialWorkspaceAgent(ctx context.Context, agentID uuid.UUID, options *peer.ConnOptions) (*agent.Conn, error) {
 	serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/dial", agentID.String()))

codersdk/workspaceresources.go

@@ -8,6 +8,8 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"inet.af/netaddr"
+	"tailscale.com/types/key"
 )
 
 type WorkspaceAgentStatus string
@@ -45,6 +47,9 @@ type WorkspaceAgent struct {
 	StartupScript        string               `json:"startup_script,omitempty"`
 	Directory            string               `json:"directory,omitempty"`
 	Apps                 []WorkspaceApp       `json:"apps"`
+	WireguardPublicKey   key.NodePublic       `json:"wireguard_public_key"`
+	DiscoPublicKey       key.DiscoPublic      `json:"disco_public_key"`
+	IPv6                 netaddr.IPPrefix     `json:"ipv6"`
 }
 
 type WorkspaceAgentResourceMetadata struct {