feat: add session token injection to provisioner (#7461)

coderd/provisionerdserver/provisionerdserver_test.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"encoding/json"
 	"net/url"
+	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -15,6 +16,7 @@ import (
 	"golang.org/x/oauth2"
 
 	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/cli/clibase"
 	"github.com/coder/coder/coderd/audit"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/dbfake"
@@ -61,6 +63,7 @@ func TestAcquireJob(t *testing.T) {
 			Auditor:               mockAuditor(),
 			TemplateScheduleStore: testTemplateScheduleStore(),
 			Tracer:                trace.NewNoopTracerProvider().Tracer("noop"),
+			DeploymentValues:      &codersdk.DeploymentValues{},
 		}
 		job, err := srv.AcquireJob(context.Background(), nil)
 		require.NoError(t, err)
@@ -102,6 +105,10 @@ func TestAcquireJob(t *testing.T) {
 		t.Parallel()
 		srv := setup(t, false)
 		gitAuthProvider := "github"
+		// Set the max session token lifetime so we can assert we
+		// create an API key with an expiration within the bounds of the
+		// deployment config.
+		srv.DeploymentValues.MaxTokenLifetime = clibase.Duration(time.Hour)
 		srv.GitAuthConfigs = []*gitauth.Config{{
 			ID:           gitAuthProvider,
 			OAuth2Config: &testutil.OAuth2Config{},
@@ -192,12 +199,16 @@ func TestAcquireJob(t *testing.T) {
 			})),
 		})
 
-		published := make(chan struct{})
-		closeSubscribe, err := srv.Pubsub.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, _ []byte) {
-			close(published)
+		startPublished := make(chan struct{})
+		var closed bool
+		closeStartSubscribe, err := srv.Pubsub.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, _ []byte) {
+			if !closed {
+				close(startPublished)
+				closed = true
+			}
 		})
 		require.NoError(t, err)
-		defer closeSubscribe()
+		defer closeStartSubscribe()
 
 		var job *proto.AcquiredJob
 
@@ -211,11 +222,21 @@ func TestAcquireJob(t *testing.T) {
 			}
 		}
 
-		<-published
+		<-startPublished
 
 		got, err := json.Marshal(job.Type)
 		require.NoError(t, err)
 
+		// Validate that a session token is generated during the job.
+		sessionToken := job.Type.(*proto.AcquiredJob_WorkspaceBuild_).WorkspaceBuild.Metadata.WorkspaceOwnerSessionToken
+		require.NotEmpty(t, sessionToken)
+		toks := strings.Split(sessionToken, "-")
+		require.Len(t, toks, 2, "invalid api key")
+		key, err := srv.Database.GetAPIKeyByID(ctx, toks[0])
+		require.NoError(t, err)
+		require.Equal(t, int64(srv.DeploymentValues.MaxTokenLifetime.Value().Seconds()), key.LifetimeSeconds)
+		require.WithinDuration(t, time.Now().Add(srv.DeploymentValues.MaxTokenLifetime.Value()), key.ExpiresAt, time.Minute)
+
 		want, err := json.Marshal(&proto.AcquiredJob_WorkspaceBuild_{
 			WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
 				WorkspaceBuildId: build.ID.String(),
@@ -247,13 +268,59 @@ func TestAcquireJob(t *testing.T) {
 					WorkspaceOwnerId:              user.ID.String(),
 					TemplateName:                  template.Name,
 					TemplateVersion:               version.Name,
+					WorkspaceOwnerSessionToken:    sessionToken,
 				},
 			},
 		})
 		require.NoError(t, err)
 
 		require.JSONEq(t, string(want), string(got))
+
+		// Assert that we delete the session token whenever
+		// a stop is issued.
+		stopbuild := dbgen.WorkspaceBuild(t, srv.Database, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			BuildNumber:       2,
+			JobID:             uuid.New(),
+			TemplateVersionID: version.ID,
+			Transition:        database.WorkspaceTransitionStop,
+			Reason:            database.BuildReasonInitiator,
+		})
+		_ = dbgen.ProvisionerJob(t, srv.Database, database.ProvisionerJob{
+			ID:            stopbuild.ID,
+			InitiatorID:   user.ID,
+			Provisioner:   database.ProvisionerTypeEcho,
+			StorageMethod: database.ProvisionerStorageMethodFile,
+			FileID:        file.ID,
+			Type:          database.ProvisionerJobTypeWorkspaceBuild,
+			Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
+				WorkspaceBuildID: stopbuild.ID,
+			})),
+		})
+
+		stopPublished := make(chan struct{})
+		closeStopSubscribe, err := srv.Pubsub.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, _ []byte) {
+			close(stopPublished)
+		})
+		require.NoError(t, err)
+		defer closeStopSubscribe()
+
+		// Grab jobs until we find the workspace build job. There is also
+		// an import version job that we need to ignore.
+		job, err = srv.AcquireJob(ctx, nil)
+		require.NoError(t, err)
+		_, ok := job.Type.(*proto.AcquiredJob_WorkspaceBuild_)
+		require.True(t, ok, "acquired job not a workspace build?")
+
+		<-stopPublished
+
+		// Validate that a session token is deleted during a stop job.
+		sessionToken = job.Type.(*proto.AcquiredJob_WorkspaceBuild_).WorkspaceBuild.Metadata.WorkspaceOwnerSessionToken
+		require.Empty(t, sessionToken)
+		_, err = srv.Database.GetAPIKeyByID(ctx, key.ID)
+		require.ErrorIs(t, err, sql.ErrNoRows)
 	})
+
 	t.Run("TemplateVersionDryRun", func(t *testing.T) {
 		t.Parallel()
 		srv := setup(t, false)
@@ -1205,6 +1272,7 @@ func setup(t *testing.T, ignoreLogErrors bool) *provisionerdserver.Server {
 		Auditor:               mockAuditor(),
 		TemplateScheduleStore: testTemplateScheduleStore(),
 		Tracer:                trace.NewNoopTracerProvider().Tracer("noop"),
+		DeploymentValues:      &codersdk.DeploymentValues{},
 	}
 }