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.

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)