mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add session token injection to provisioner (#7461)
This commit is contained in:
@@ -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{},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user