feat: Separate workspace agent for tests (#567)

This adds tests for Google Cloud authentication, and lays
the ground-work for future agent auth types in the future.

agent/agent.go

@@ -26,36 +26,6 @@ import (
 	"golang.org/x/xerrors"
 )
 
-func DialSSH(conn *peer.Conn) (net.Conn, error) {
-	channel, err := conn.Dial(context.Background(), "ssh", &peer.ChannelOptions{
-		Protocol: "ssh",
-	})
-	if err != nil {
-		return nil, err
-	}
-	return channel.NetConn(), nil
-}
-
-func DialSSHClient(conn *peer.Conn) (*gossh.Client, error) {
-	netConn, err := DialSSH(conn)
-	if err != nil {
-		return nil, err
-	}
-	sshConn, channels, requests, err := gossh.NewClientConn(netConn, "localhost:22", &gossh.ClientConfig{
-		Config: gossh.Config{
-			Ciphers: []string{"arcfour"},
-		},
-		// SSH host validation isn't helpful, because obtaining a peer
-		// connection already signifies user-intent to dial a workspace.
-		// #nosec
-		HostKeyCallback: gossh.InsecureIgnoreHostKey(),
-	})
-	if err != nil {
-		return nil, err
-	}
-	return gossh.NewClient(sshConn, channels, requests), nil
-}
-
 type Options struct {
 	Logger slog.Logger
 }
@@ -64,7 +34,7 @@ type Dialer func(ctx context.Context, options *peer.ConnOptions) (*peerbroker.Li
 
 func New(dialer Dialer, options *peer.ConnOptions) io.Closer {
 	ctx, cancelFunc := context.WithCancel(context.Background())
-	server := &server{
+	server := &agent{
 		clientDialer: dialer,
 		options:      options,
 		closeCancel:  cancelFunc,
@@ -74,7 +44,7 @@ func New(dialer Dialer, options *peer.ConnOptions) io.Closer {
 	return server
 }
 
-type server struct {
+type agent struct {
 	clientDialer Dialer
 	options      *peer.ConnOptions
 
@@ -86,7 +56,7 @@ type server struct {
 	sshServer *ssh.Server
 }
 
-func (s *server) run(ctx context.Context) {
+func (s *agent) run(ctx context.Context) {
 	var peerListener *peerbroker.Listener
 	var err error
 	// An exponential back-off occurs when the connection is failing to dial.
@@ -103,7 +73,7 @@ func (s *server) run(ctx context.Context) {
 			s.options.Logger.Warn(context.Background(), "failed to dial", slog.Error(err))
 			continue
 		}
-		s.options.Logger.Debug(context.Background(), "connected")
+		s.options.Logger.Info(context.Background(), "connected")
 		break
 	}
 	select {
@@ -129,7 +99,7 @@ func (s *server) run(ctx context.Context) {
 	}
 }
 
-func (s *server) handlePeerConn(ctx context.Context, conn *peer.Conn) {
+func (s *agent) handlePeerConn(ctx context.Context, conn *peer.Conn) {
 	go func() {
 		<-conn.Closed()
 		s.connCloseWait.Done()
@@ -156,7 +126,7 @@ func (s *server) handlePeerConn(ctx context.Context, conn *peer.Conn) {
 	}
 }
 
-func (s *server) init(ctx context.Context) {
+func (s *agent) init(ctx context.Context) {
 	// Clients' should ignore the host key when connecting.
 	// The agent needs to authenticate with coderd to SSH,
 	// so SSH authentication doesn't improve security.
@@ -221,7 +191,7 @@ func (s *server) init(ctx context.Context) {
 	go s.run(ctx)
 }
 
-func (*server) handleSSHSession(session ssh.Session) error {
+func (*agent) handleSSHSession(session ssh.Session) error {
 	var (
 		command string
 		args    = []string{}
@@ -316,7 +286,7 @@ func (*server) handleSSHSession(session ssh.Session) error {
 }
 
 // isClosed returns whether the API is closed or not.
-func (s *server) isClosed() bool {
+func (s *agent) isClosed() bool {
 	select {
 	case <-s.closed:
 		return true
@@ -325,7 +295,7 @@ func (s *server) isClosed() bool {
 	}
 }
 
-func (s *server) Close() error {
+func (s *agent) Close() error {
 	s.closeMutex.Lock()
 	defer s.closeMutex.Unlock()
 	if s.isClosed() {

agent/agent_test.go

@@ -39,7 +39,8 @@ func TestAgent(t *testing.T) {
 		t.Cleanup(func() {
 			_ = conn.Close()
 		})
-		sshClient, err := agent.DialSSHClient(conn)
+		client := agent.Conn{conn}
+		sshClient, err := client.SSHClient()
 		require.NoError(t, err)
 		session, err := sshClient.NewSession()
 		require.NoError(t, err)
@@ -64,7 +65,8 @@ func TestAgent(t *testing.T) {
 		t.Cleanup(func() {
 			_ = conn.Close()
 		})
-		sshClient, err := agent.DialSSHClient(conn)
+		client := &agent.Conn{conn}
+		sshClient, err := client.SSHClient()
 		require.NoError(t, err)
 		session, err := sshClient.NewSession()
 		require.NoError(t, err)

agent/conn.go

@@ -0,0 +1,50 @@
+package agent
+
+import (
+	"context"
+	"net"
+
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/peer"
+)
+
+// Conn wraps a peer connection with helper functions to
+// communicate with the agent.
+type Conn struct {
+	*peer.Conn
+}
+
+// SSH dials the built-in SSH server.
+func (c *Conn) SSH() (net.Conn, error) {
+	channel, err := c.Dial(context.Background(), "ssh", &peer.ChannelOptions{
+		Protocol: "ssh",
+	})
+	if err != nil {
+		return nil, xerrors.Errorf("dial: %w", err)
+	}
+	return channel.NetConn(), nil
+}
+
+// SSHClient calls SSH to create a client that uses a weak cipher
+// for high throughput.
+func (c *Conn) SSHClient() (*ssh.Client, error) {
+	netConn, err := c.SSH()
+	if err != nil {
+		return nil, xerrors.Errorf("ssh: %w", err)
+	}
+	sshConn, channels, requests, err := ssh.NewClientConn(netConn, "localhost:22", &ssh.ClientConfig{
+		Config: ssh.Config{
+			Ciphers: []string{"arcfour"},
+		},
+		// SSH host validation isn't helpful, because obtaining a peer
+		// connection already signifies user-intent to dial a workspace.
+		// #nosec
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	})
+	if err != nil {
+		return nil, xerrors.Errorf("ssh conn: %w", err)
+	}
+	return ssh.NewClient(sshConn, channels, requests), nil
+}

cli/ssh.go

@@ -1,17 +1,15 @@
 package cli
 
 import (
-	"fmt"
 	"os"
 
-	"github.com/pion/webrtc/v3"
 	"github.com/spf13/cobra"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/term"
+	"golang.org/x/xerrors"
 
-	"github.com/coder/coder/agent"
-	"github.com/coder/coder/peer"
-	"github.com/coder/coder/peerbroker"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/database"
 )
 
 func workspaceSSH() *cobra.Command {
@@ -26,58 +24,76 @@ func workspaceSSH() *cobra.Command {
 			if err != nil {
 				return err
 			}
+			if workspace.LatestBuild.Transition == database.WorkspaceTransitionDelete {
+				return xerrors.New("workspace is deleting...")
+			}
 			resources, err := client.WorkspaceResourcesByBuild(cmd.Context(), workspace.LatestBuild.ID)
 			if err != nil {
 				return err
 			}
-
+			resourceByAddress := make(map[string]codersdk.WorkspaceResource)
 			for _, resource := range resources {
-				_, _ = fmt.Printf("Got resource: %+v\n", resource)
 				if resource.Agent == nil {
 					continue
 				}
+				resourceByAddress[resource.Address] = resource
+			}
+			var resourceAddress string
+			if len(args) >= 2 {
+				resourceAddress = args[1]
+			} else {
+				// No resource name was provided!
+				if len(resourceByAddress) > 1 {
+					// List available resources to connect into?
+					return xerrors.Errorf("multiple agents")
+				}
+				for _, resource := range resourceByAddress {
+					resourceAddress = resource.Address
+					break
+				}
+			}
+			resource, exists := resourceByAddress[resourceAddress]
+			if !exists {
+				resourceKeys := make([]string, 0)
+				for resourceKey := range resourceByAddress {
+					resourceKeys = append(resourceKeys, resourceKey)
+				}
+				return xerrors.Errorf("no sshable agent with address %q: %+v", resourceAddress, resourceKeys)
+			}
+			if resource.Agent.LastConnectedAt == nil {
+				return xerrors.Errorf("agent hasn't connected yet")
+			}
 
-				dialed, err := client.DialWorkspaceAgent(cmd.Context(), resource.ID)
-				if err != nil {
-					return err
-				}
-				stream, err := dialed.NegotiateConnection(cmd.Context())
-				if err != nil {
-					return err
-				}
-				conn, err := peerbroker.Dial(stream, []webrtc.ICEServer{{
-					URLs: []string{"stun:stun.l.google.com:19302"},
-				}}, &peer.ConnOptions{})
-				if err != nil {
-					return err
-				}
-				client, err := agent.DialSSHClient(conn)
-				if err != nil {
-					return err
-				}
+			conn, err := client.DialWorkspaceAgent(cmd.Context(), resource.ID, nil, nil)
+			if err != nil {
+				return err
+			}
+			sshClient, err := conn.SSHClient()
+			if err != nil {
+				return err
+			}
 
-				session, err := client.NewSession()
-				if err != nil {
-					return err
-				}
-				_, _ = term.MakeRaw(int(os.Stdin.Fd()))
-				err = session.RequestPty("xterm-256color", 128, 128, ssh.TerminalModes{
-					ssh.OCRNL: 1,
-				})
-				if err != nil {
-					return err
-				}
-				session.Stdin = os.Stdin
-				session.Stdout = os.Stdout
-				session.Stderr = os.Stderr
-				err = session.Shell()
-				if err != nil {
-					return err
-				}
-				err = session.Wait()
-				if err != nil {
-					return err
-				}
+			sshSession, err := sshClient.NewSession()
+			if err != nil {
+				return err
+			}
+			_, _ = term.MakeRaw(int(os.Stdin.Fd()))
+			err = sshSession.RequestPty("xterm-256color", 128, 128, ssh.TerminalModes{
+				ssh.OCRNL: 1,
+			})
+			if err != nil {
+				return err
+			}
+			sshSession.Stdin = os.Stdin
+			sshSession.Stdout = os.Stdout
+			sshSession.Stderr = os.Stderr
+			err = sshSession.Shell()
+			if err != nil {
+				return err
+			}
+			err = sshSession.Wait()
+			if err != nil {
+				return err
 			}
 
 			return nil

cli/workspaceagent.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"time"
 
+	"cloud.google.com/go/compute/metadata"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 
@@ -19,25 +20,24 @@ import (
 )
 
 func workspaceAgent() *cobra.Command {
-	return &cobra.Command{
+	var (
+		rawURL string
+		auth   string
+	)
+	cmd := &cobra.Command{
 		Use: "agent",
 		// This command isn't useful to manually execute.
 		Hidden: true,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			coderURLRaw, exists := os.LookupEnv("CODER_URL")
-			if !exists {
+			if rawURL == "" {
 				return xerrors.New("CODER_URL must be set")
 			}
-			coderURL, err := url.Parse(coderURLRaw)
+			coderURL, err := url.Parse(rawURL)
 			if err != nil {
-				return xerrors.Errorf("parse %q: %w", coderURLRaw, err)
+				return xerrors.Errorf("parse %q: %w", rawURL, err)
 			}
-			logger := slog.Make(sloghuman.Sink(cmd.OutOrStdout()))
+			logger := slog.Make(sloghuman.Sink(cmd.OutOrStdout())).Leveled(slog.LevelDebug)
 			client := codersdk.New(coderURL)
-			auth, exists := os.LookupEnv("CODER_AUTH")
-			if !exists {
-				auth = "token"
-			}
 			switch auth {
 			case "token":
 				sessionToken, exists := os.LookupEnv("CODER_TOKEN")
@@ -46,16 +46,26 @@ func workspaceAgent() *cobra.Command {
 				}
 				client.SessionToken = sessionToken
 			case "google-instance-identity":
+				// This is *only* done for testing to mock client authentication.
+				// This will never be set in a production scenario.
+				var gcpClient *metadata.Client
+				gcpClientRaw := cmd.Context().Value("gcp-client")
+				if gcpClientRaw != nil {
+					gcpClient, _ = gcpClientRaw.(*metadata.Client)
+				}
+
 				ctx, cancelFunc := context.WithTimeout(cmd.Context(), 30*time.Second)
 				defer cancelFunc()
 				for retry.New(100*time.Millisecond, 5*time.Second).Wait(ctx) {
 					var response codersdk.WorkspaceAgentAuthenticateResponse
-					response, err = client.AuthWorkspaceGoogleInstanceIdentity(cmd.Context(), "", nil)
+
+					response, err = client.AuthWorkspaceGoogleInstanceIdentity(ctx, "", gcpClient)
 					if err != nil {
 						logger.Warn(ctx, "authenticate workspace with Google Instance Identity", slog.Error(err))
 						continue
 					}
 					client.SessionToken = response.SessionToken
+					logger.Info(ctx, "authenticated with Google Instance Identity")
 					break
 				}
 				if err != nil {
@@ -73,4 +83,12 @@ func workspaceAgent() *cobra.Command {
 			return closer.Close()
 		},
 	}
+	defaultAuth := os.Getenv("CODER_AUTH")
+	if defaultAuth == "" {
+		defaultAuth = "token"
+	}
+	cmd.Flags().StringVarP(&auth, "auth", "", defaultAuth, "Specify the authentication type to use for the agent.")
+	cmd.Flags().StringVarP(&rawURL, "url", "", os.Getenv("CODER_URL"), "Specify the URL to access Coder.")
+
+	return cmd
 }

cli/workspaceagent_test.go

@@ -0,0 +1,70 @@
+package cli_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/provisioner/echo"
+	"github.com/coder/coder/provisionersdk/proto"
+)
+
+func TestWorkspaceAgent(t *testing.T) {
+	t.Parallel()
+	t.Run("GoogleCloud", func(t *testing.T) {
+		t.Parallel()
+		instanceID := "instanceidentifier"
+		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
+		client := coderdtest.New(t, &coderdtest.Options{
+			GoogleTokenValidator: validator,
+		})
+		user := coderdtest.CreateFirstUser(t, client)
+		coderdtest.NewProvisionerDaemon(t, client)
+		version := coderdtest.CreateProjectVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse: echo.ParseComplete,
+			Provision: []*proto.Provision_Response{{
+				Type: &proto.Provision_Response_Complete{
+					Complete: &proto.Provision_Complete{
+						Resources: []*proto.Resource{{
+							Name: "somename",
+							Type: "someinstance",
+							Agent: &proto.Agent{
+								Auth: &proto.Agent_InstanceId{
+									InstanceId: instanceID,
+								},
+							},
+						}},
+					},
+				},
+			}},
+		})
+		project := coderdtest.CreateProject(t, client, user.OrganizationID, version.ID)
+		coderdtest.AwaitProjectVersionJob(t, client, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, "me", project.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+
+		cmd, _ := clitest.New(t, "workspaces", "agent", "--auth", "google-instance-identity", "--url", client.URL.String())
+		ctx, cancelFunc := context.WithCancel(context.Background())
+		defer cancelFunc()
+		go func() {
+			// A linting error occurs for weakly typing the context value here,
+			// but it seems reasonable for a one-off test.
+			// nolint
+			ctx = context.WithValue(ctx, "gcp-client", metadata)
+			err := cmd.ExecuteContext(ctx)
+			require.NoError(t, err)
+		}()
+		coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
+		resources, err := client.WorkspaceResourcesByBuild(ctx, workspace.LatestBuild.ID)
+		require.NoError(t, err)
+		dialer, err := client.DialWorkspaceAgent(ctx, resources[0].ID, nil, nil)
+		require.NoError(t, err)
+		defer dialer.Close()
+		_, err = dialer.Ping()
+		require.NoError(t, err)
+		cancelFunc()
+	})
+}

coderd/coderdtest/coderdtest.go

@@ -1,10 +1,18 @@
 package coderdtest
 
 import (
+	"bytes"
 	"context"
+	"crypto/rand"
+	"crypto/rsa"
 	"database/sql"
+	"encoding/base64"
+	"encoding/json"
 	"io"
+	"io/ioutil"
+	"math/big"
 	"net"
+	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"os"
@@ -12,6 +20,8 @@ import (
 	"testing"
 	"time"
 
+	"cloud.google.com/go/compute/metadata"
+	"github.com/golang-jwt/jwt"
 	"github.com/google/uuid"
 	"github.com/moby/moby/pkg/namesgenerator"
 	"github.com/stretchr/testify/require"
@@ -22,6 +32,7 @@ import (
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/coderd"
 	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/cryptorand"
 	"github.com/coder/coder/database"
 	"github.com/coder/coder/database/databasefake"
 	"github.com/coder/coder/database/postgres"
@@ -260,6 +271,76 @@ func CreateWorkspace(t *testing.T, client *codersdk.Client, user string, project
 	return workspace
 }
 
+// NewGoogleInstanceIdentity returns a metadata client and ID token validator for faking
+// instance authentication for Google Cloud.
+// nolint:revive
+func NewGoogleInstanceIdentity(t *testing.T, instanceID string, expired bool) (*idtoken.Validator, *metadata.Client) {
+	keyID, err := cryptorand.String(12)
+	require.NoError(t, err)
+	claims := jwt.MapClaims{
+		"google": map[string]interface{}{
+			"compute_engine": map[string]string{
+				"instance_id": instanceID,
+			},
+		},
+	}
+	if !expired {
+		claims["exp"] = time.Now().AddDate(1, 0, 0).Unix()
+	}
+	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
+	token.Header["kid"] = keyID
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+	signedKey, err := token.SignedString(privateKey)
+	require.NoError(t, err)
+
+	// Taken from: https://github.com/googleapis/google-api-go-client/blob/4bb729045d611fa77bdbeb971f6a1204ba23161d/idtoken/validate.go#L57-L75
+	type jwk struct {
+		Kid string `json:"kid"`
+		N   string `json:"n"`
+		E   string `json:"e"`
+	}
+	type certResponse struct {
+		Keys []jwk `json:"keys"`
+	}
+
+	validator, err := idtoken.NewValidator(context.Background(), option.WithHTTPClient(&http.Client{
+		Transport: roundTripper(func(r *http.Request) (*http.Response, error) {
+			data, err := json.Marshal(certResponse{
+				Keys: []jwk{{
+					Kid: keyID,
+					N:   base64.RawURLEncoding.EncodeToString(privateKey.N.Bytes()),
+					E:   base64.RawURLEncoding.EncodeToString(new(big.Int).SetInt64(int64(privateKey.E)).Bytes()),
+				}},
+			})
+			require.NoError(t, err)
+			return &http.Response{
+				StatusCode: http.StatusOK,
+				Body:       ioutil.NopCloser(bytes.NewReader(data)),
+				Header:     make(http.Header),
+			}, nil
+		}),
+	}))
+	require.NoError(t, err)
+
+	return validator, metadata.NewClient(&http.Client{
+		Transport: roundTripper(func(r *http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode: http.StatusOK,
+				Body:       ioutil.NopCloser(bytes.NewReader([]byte(signedKey))),
+				Header:     make(http.Header),
+			}, nil
+		}),
+	})
+}
+
 func randomUsername() string {
 	return strings.ReplaceAll(namesgenerator.GetRandomName(0), "_", "-")
 }
+
+// Used to easily create an HTTP transport!
+type roundTripper func(req *http.Request) (*http.Response, error)
+
+func (r roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	return r(req)
+}

coderd/coderdtest/coderdtest_test.go

@@ -23,5 +23,6 @@ func TestNew(t *testing.T) {
 	workspace := coderdtest.CreateWorkspace(t, client, "me", project.ID)
 	coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 	coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
+	_, _ = coderdtest.NewGoogleInstanceIdentity(t, "example", false)
 	closer.Close()
 }

coderd/workspaceresourceauth_test.go

@@ -1,27 +1,14 @@
 package coderd_test
 
 import (
-	"bytes"
 	"context"
-	"crypto/rand"
-	"crypto/rsa"
-	"encoding/base64"
-	"encoding/json"
-	"io/ioutil"
-	"math/big"
 	"net/http"
 	"testing"
-	"time"
 
-	"cloud.google.com/go/compute/metadata"
-	"github.com/golang-jwt/jwt"
 	"github.com/stretchr/testify/require"
-	"google.golang.org/api/idtoken"
-	"google.golang.org/api/option"
 
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/codersdk"
-	"github.com/coder/coder/cryptorand"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
 )
@@ -31,12 +18,11 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 	t.Run("Expired", func(t *testing.T) {
 		t.Parallel()
 		instanceID := "instanceidentifier"
-		signedKey, keyID, privateKey := createSignedToken(t, instanceID, &jwt.MapClaims{})
-		validator := createValidator(t, keyID, privateKey)
+		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, true)
 		client := coderdtest.New(t, &coderdtest.Options{
 			GoogleTokenValidator: validator,
 		})
-		_, err := client.AuthWorkspaceGoogleInstanceIdentity(context.Background(), "", createMetadataClient(signedKey))
+		_, err := client.AuthWorkspaceGoogleInstanceIdentity(context.Background(), "", metadata)
 		var apiErr *codersdk.Error
 		require.ErrorAs(t, err, &apiErr)
 		require.Equal(t, http.StatusUnauthorized, apiErr.StatusCode())
@@ -45,12 +31,11 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 	t.Run("InstanceNotFound", func(t *testing.T) {
 		t.Parallel()
 		instanceID := "instanceidentifier"
-		signedKey, keyID, privateKey := createSignedToken(t, instanceID, nil)
-		validator := createValidator(t, keyID, privateKey)
+		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
 		client := coderdtest.New(t, &coderdtest.Options{
 			GoogleTokenValidator: validator,
 		})
-		_, err := client.AuthWorkspaceGoogleInstanceIdentity(context.Background(), "", createMetadataClient(signedKey))
+		_, err := client.AuthWorkspaceGoogleInstanceIdentity(context.Background(), "", metadata)
 		var apiErr *codersdk.Error
 		require.ErrorAs(t, err, &apiErr)
 		require.Equal(t, http.StatusNotFound, apiErr.StatusCode())
@@ -59,8 +44,7 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 	t.Run("Success", func(t *testing.T) {
 		t.Parallel()
 		instanceID := "instanceidentifier"
-		signedKey, keyID, privateKey := createSignedToken(t, instanceID, nil)
-		validator := createValidator(t, keyID, privateKey)
+		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
 		client := coderdtest.New(t, &coderdtest.Options{
 			GoogleTokenValidator: validator,
 		})
@@ -104,85 +88,7 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 		workspace := coderdtest.CreateWorkspace(t, client, "me", project.ID)
 		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 
-		_, err := client.AuthWorkspaceGoogleInstanceIdentity(context.Background(), "", createMetadataClient(signedKey))
+		_, err := client.AuthWorkspaceGoogleInstanceIdentity(context.Background(), "", metadata)
 		require.NoError(t, err)
 	})
 }
-
-// Used to easily create an HTTP transport!
-type roundTripper func(req *http.Request) (*http.Response, error)
-
-func (r roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
-	return r(req)
-}
-
-// Create's a new Google metadata client to authenticate.
-func createMetadataClient(signedKey string) *metadata.Client {
-	return metadata.NewClient(&http.Client{
-		Transport: roundTripper(func(r *http.Request) (*http.Response, error) {
-			return &http.Response{
-				StatusCode: http.StatusOK,
-				Body:       ioutil.NopCloser(bytes.NewReader([]byte(signedKey))),
-				Header:     make(http.Header),
-			}, nil
-		}),
-	})
-}
-
-// Create's a signed JWT with a randomly generated private key.
-func createSignedToken(t *testing.T, instanceID string, claims *jwt.MapClaims) (signedKey string, keyID string, privateKey *rsa.PrivateKey) {
-	keyID, err := cryptorand.String(12)
-	require.NoError(t, err)
-	if claims == nil {
-		claims = &jwt.MapClaims{
-			"exp": time.Now().AddDate(1, 0, 0).Unix(),
-			"google": map[string]interface{}{
-				"compute_engine": map[string]string{
-					"instance_id": instanceID,
-				},
-			},
-		}
-	}
-	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
-	token.Header["kid"] = keyID
-	privateKey, err = rsa.GenerateKey(rand.Reader, 2048)
-	require.NoError(t, err)
-	signedKey, err = token.SignedString(privateKey)
-	require.NoError(t, err)
-	return signedKey, keyID, privateKey
-}
-
-// Create's a validator that verifies against the provided private key.
-// In a production scenario, the validator calls against the Google OAuth API
-// to obtain certificates.
-func createValidator(t *testing.T, keyID string, privateKey *rsa.PrivateKey) *idtoken.Validator {
-	// Taken from: https://github.com/googleapis/google-api-go-client/blob/4bb729045d611fa77bdbeb971f6a1204ba23161d/idtoken/validate.go#L57-L75
-	type jwk struct {
-		Kid string `json:"kid"`
-		N   string `json:"n"`
-		E   string `json:"e"`
-	}
-	type certResponse struct {
-		Keys []jwk `json:"keys"`
-	}
-
-	validator, err := idtoken.NewValidator(context.Background(), option.WithHTTPClient(&http.Client{
-		Transport: roundTripper(func(r *http.Request) (*http.Response, error) {
-			data, err := json.Marshal(certResponse{
-				Keys: []jwk{{
-					Kid: keyID,
-					N:   base64.RawURLEncoding.EncodeToString(privateKey.N.Bytes()),
-					E:   base64.RawURLEncoding.EncodeToString(new(big.Int).SetInt64(int64(privateKey.E)).Bytes()),
-				}},
-			})
-			require.NoError(t, err)
-			return &http.Response{
-				StatusCode: http.StatusOK,
-				Body:       ioutil.NopCloser(bytes.NewReader(data)),
-				Header:     make(http.Header),
-			}, nil
-		}),
-	}))
-	require.NoError(t, err)
-	return validator
-}

coderd/workspaceresources_test.go

@@ -14,7 +14,6 @@ import (
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/peer"
-	"github.com/coder/coder/peerbroker"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
 )
@@ -95,14 +94,7 @@ func TestWorkspaceAgentListen(t *testing.T) {
 		_ = agentCloser.Close()
 	})
 	resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
-	workspaceClient, err := client.DialWorkspaceAgent(context.Background(), resources[0].ID)
-	require.NoError(t, err)
-	t.Cleanup(func() {
-		_ = workspaceClient.DRPCConn().Close()
-	})
-	stream, err := workspaceClient.NegotiateConnection(context.Background())
-	require.NoError(t, err)
-	conn, err := peerbroker.Dial(stream, nil, &peer.ConnOptions{
+	conn, err := client.DialWorkspaceAgent(context.Background(), resources[0].ID, nil, &peer.ConnOptions{
 		Logger: slogtest.Make(t, nil).Named("client").Leveled(slog.LevelDebug),
 	})
 	require.NoError(t, err)

codersdk/workspaceresources.go

@@ -15,6 +15,7 @@ import (
 	"golang.org/x/xerrors"
 	"nhooyr.io/websocket"
 
+	"github.com/coder/coder/agent"
 	"github.com/coder/coder/database"
 	"github.com/coder/coder/httpmw"
 	"github.com/coder/coder/peer"
@@ -90,7 +91,7 @@ func (c *Client) WorkspaceResource(ctx context.Context, id uuid.UUID) (Workspace
 }
 
 // DialWorkspaceAgent creates a connection to the specified resource.
-func (c *Client) DialWorkspaceAgent(ctx context.Context, resource uuid.UUID) (proto.DRPCPeerBrokerClient, error) {
+func (c *Client) DialWorkspaceAgent(ctx context.Context, resource uuid.UUID, iceServers []webrtc.ICEServer, opts *peer.ConnOptions) (*agent.Conn, error) {
 	serverURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceresources/%s/dial", resource.String()))
 	if err != nil {
 		return nil, xerrors.Errorf("parse url: %w", err)
@@ -123,7 +124,25 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, resource uuid.UUID) (pr
 	if err != nil {
 		return nil, xerrors.Errorf("multiplex client: %w", err)
 	}
-	return proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(session)), nil
+	client := proto.NewDRPCPeerBrokerClient(provisionersdk.Conn(session))
+	stream, err := client.NegotiateConnection(ctx)
+	if err != nil {
+		return nil, xerrors.Errorf("negotiate connection: %w", err)
+	}
+	peerConn, err := peerbroker.Dial(stream, iceServers, opts)
+	if err != nil {
+		return nil, xerrors.Errorf("dial peer: %w", err)
+	}
+	go func() {
+		// The stream is kept alive to renegotiate the RTC connection
+		// if need-be. The calling context can be canceled to end
+		// the negotiation stream, but not the peer connection.
+		<-peerConn.Closed()
+		_ = conn.Close(websocket.StatusNormalClosure, "")
+	}()
+	return &agent.Conn{
+		Conn: peerConn,
+	}, nil
 }
 
 // ListenWorkspaceAgent connects as a workspace agent.