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

This commit is contained in:
Jon Ayers
2023-05-17 23:29:22 -05:00
committed by GitHub
parent 00a2413c03
commit ec5ef51b49
13 changed files with 578 additions and 219 deletions
@@ -27,6 +27,7 @@ import (
"cdr.dev/slog"
"github.com/coder/coder/coderd/apikey"
"github.com/coder/coder/coderd/audit"
"github.com/coder/coder/coderd/database"
"github.com/coder/coder/coderd/database/dbauthz"
@@ -62,6 +63,7 @@ type Server struct {
QuotaCommitter *atomic.Pointer[proto.QuotaCommitter]
Auditor *atomic.Pointer[audit.Auditor]
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
DeploymentValues *codersdk.DeploymentValues
AcquireJobDebounce time.Duration
OIDCConfig httpmw.OAuth2Config
@@ -193,6 +195,20 @@ func (server *Server) AcquireJob(ctx context.Context, _ *proto.Empty) (*proto.Ac
}
}
var sessionToken string
switch workspaceBuild.Transition {
case database.WorkspaceTransitionStart:
sessionToken, err = server.regenerateSessionToken(ctx, owner, workspace)
if err != nil {
return nil, failJob(fmt.Sprintf("regenerate session token: %s", err))
}
case database.WorkspaceTransitionStop, database.WorkspaceTransitionDelete:
err = deleteSessionToken(ctx, server.Database, workspace)
if err != nil {
return nil, failJob(fmt.Sprintf("delete session token: %s", err))
}
}
// Compute parameters for the workspace to consume.
parameters, err := parameter.Compute(ctx, server.Database, parameter.ComputeScope{
TemplateImportJobID: templateVersion.JobID,
@@ -286,6 +302,7 @@ func (server *Server) AcquireJob(ctx context.Context, _ *proto.Empty) (*proto.Ac
WorkspaceOwnerId: owner.ID.String(),
TemplateName: template.Name,
TemplateVersion: templateVersion.Name,
WorkspaceOwnerSessionToken: sessionToken,
},
LogLevel: input.LogLevel,
},
@@ -1410,6 +1427,64 @@ func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.
return nil
}
func workspaceSessionTokenName(workspace database.Workspace) string {
return fmt.Sprintf("%s_%s_session_token", workspace.OwnerID, workspace.ID)
}
func (server *Server) regenerateSessionToken(ctx context.Context, user database.User, workspace database.Workspace) (string, error) {
newkey, sessionToken, err := apikey.Generate(apikey.CreateParams{
UserID: user.ID,
LoginType: user.LoginType,
DeploymentValues: server.DeploymentValues,
TokenName: workspaceSessionTokenName(workspace),
LifetimeSeconds: int64(server.DeploymentValues.MaxTokenLifetime.Value().Seconds()),
})
if err != nil {
return "", xerrors.Errorf("generate API key: %w", err)
}
err = server.Database.InTx(func(tx database.Store) error {
err := deleteSessionToken(ctx, tx, workspace)
if err != nil {
return xerrors.Errorf("delete session token: %w", err)
}
_, err = tx.InsertAPIKey(ctx, newkey)
if err != nil {
return xerrors.Errorf("insert API key: %w", err)
}
return nil
}, nil)
if err != nil {
return "", xerrors.Errorf("create API key: %w", err)
}
return sessionToken, nil
}
func deleteSessionToken(ctx context.Context, db database.Store, workspace database.Workspace) error {
err := db.InTx(func(tx database.Store) error {
key, err := tx.GetAPIKeyByName(ctx, database.GetAPIKeyByNameParams{
UserID: workspace.OwnerID,
TokenName: workspaceSessionTokenName(workspace),
})
if err == nil {
err = tx.DeleteAPIKeyByID(ctx, key.ID)
}
if err != nil && !xerrors.Is(err, sql.ErrNoRows) {
return xerrors.Errorf("get api key by name: %w", err)
}
return nil
}, nil)
if err != nil {
return xerrors.Errorf("in tx: %w", err)
}
return nil
}
// obtainOIDCAccessToken returns a valid OpenID Connect access token
// for the user if it's able to obtain one, otherwise it returns an empty string.
func obtainOIDCAccessToken(ctx context.Context, db database.Store, oidcConfig httpmw.OAuth2Config, userID uuid.UUID) (string, error) {
@@ -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{},
}
}