diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 77ab7dc5ae54c..632efc4c8e103 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1764,6 +1764,13 @@ func scopedOrgRoleIdentifiers(names []string, orgID uuid.UUID) []rbac.RoleIdenti return out } +func (q *querier) AcquireExternalAuthLinkRefreshLease(ctx context.Context, arg database.AcquireExternalAuthLinkRefreshLeaseParams) (database.ExternalAuthLink, error) { + fetch := func(ctx context.Context, arg database.AcquireExternalAuthLinkRefreshLeaseParams) (database.ExternalAuthLink, error) { + return q.db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{UserID: arg.UserID, ProviderID: arg.ProviderID}) + } + return fetchAndQuery(q.log, q.auth, policy.ActionUpdatePersonal, fetch, q.db.AcquireExternalAuthLinkRefreshLease)(ctx, arg) +} + func (q *querier) AcquireLock(ctx context.Context, id int64) error { return q.db.AcquireLock(ctx, id) } @@ -7102,6 +7109,13 @@ func (q *querier) RegisterWorkspaceProxy(ctx context.Context, arg database.Regis return updateWithReturn(q.log, q.auth, fetch, q.db.RegisterWorkspaceProxy)(ctx, arg) } +func (q *querier) ReleaseExternalAuthLinkRefreshLease(ctx context.Context, arg database.ReleaseExternalAuthLinkRefreshLeaseParams) error { + fetch := func(ctx context.Context, arg database.ReleaseExternalAuthLinkRefreshLeaseParams) (database.ExternalAuthLink, error) { + return q.db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{UserID: arg.UserID, ProviderID: arg.ProviderID}) + } + return fetchAndExec(q.log, q.auth, policy.ActionUpdatePersonal, fetch, q.db.ReleaseExternalAuthLinkRefreshLease)(ctx, arg) +} + func (q *querier) RemoveUserFromGroups(ctx context.Context, arg database.RemoveUserFromGroupsParams) ([]uuid.UUID, error) { // This is a system function to clear user groups in group sync. if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil { @@ -7659,13 +7673,6 @@ func (q *querier) UpdateExternalAuthLink(ctx context.Context, arg database.Updat return fetchAndQuery(q.log, q.auth, policy.ActionUpdatePersonal, fetch, q.db.UpdateExternalAuthLink)(ctx, arg) } -func (q *querier) UpdateExternalAuthLinkRefreshToken(ctx context.Context, arg database.UpdateExternalAuthLinkRefreshTokenParams) error { - fetch := func(ctx context.Context, arg database.UpdateExternalAuthLinkRefreshTokenParams) (database.ExternalAuthLink, error) { - return q.db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{UserID: arg.UserID, ProviderID: arg.ProviderID}) - } - return fetchAndExec(q.log, q.auth, policy.ActionUpdatePersonal, fetch, q.db.UpdateExternalAuthLinkRefreshToken)(ctx, arg) -} - func (q *querier) UpdateGitSSHKey(ctx context.Context, arg database.UpdateGitSSHKeyParams) (database.GitSSHKey, error) { fetch := func(ctx context.Context, arg database.UpdateGitSSHKeyParams) (database.GitSSHKey, error) { return q.db.GetGitSSHKey(ctx, arg.UserID) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index fa784ded25b63..738b4a2de9997 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3261,13 +3261,6 @@ func (s *MethodTestSuite) TestUser() { dbm.EXPECT().InsertExternalAuthLink(gomock.Any(), arg).Return(database.ExternalAuthLink{}, nil).AnyTimes() check.Args(arg).Asserts(u, policy.ActionUpdatePersonal) })) - s.Run("UpdateExternalAuthLinkRefreshToken", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - link := testutil.Fake(s.T(), faker, database.ExternalAuthLink{}) - arg := database.UpdateExternalAuthLinkRefreshTokenParams{OAuthRefreshToken: "", OAuthRefreshTokenKeyID: "", ProviderID: link.ProviderID, UserID: link.UserID, UpdatedAt: link.UpdatedAt, OldOauthRefreshToken: link.OAuthRefreshToken} - dbm.EXPECT().GetExternalAuthLink(gomock.Any(), database.GetExternalAuthLinkParams{ProviderID: link.ProviderID, UserID: link.UserID}).Return(link, nil).AnyTimes() - dbm.EXPECT().UpdateExternalAuthLinkRefreshToken(gomock.Any(), arg).Return(nil).AnyTimes() - check.Args(arg).Asserts(link, policy.ActionUpdatePersonal) - })) s.Run("UpdateExternalAuthLink", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { link := testutil.Fake(s.T(), faker, database.ExternalAuthLink{}) arg := database.UpdateExternalAuthLinkParams{ProviderID: link.ProviderID, UserID: link.UserID, OAuthAccessToken: link.OAuthAccessToken, OAuthRefreshToken: link.OAuthRefreshToken, OAuthExpiry: link.OAuthExpiry, UpdatedAt: link.UpdatedAt} @@ -3275,6 +3268,23 @@ func (s *MethodTestSuite) TestUser() { dbm.EXPECT().UpdateExternalAuthLink(gomock.Any(), arg).Return(link, nil).AnyTimes() check.Args(arg).Asserts(link, policy.ActionUpdatePersonal).Returns(link) })) + s.Run("AcquireExternalAuthLinkRefreshLease", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + link := testutil.Fake(s.T(), faker, database.ExternalAuthLink{}) + dbm.EXPECT().GetExternalAuthLink(gomock.Any(), database.GetExternalAuthLinkParams{ProviderID: link.ProviderID, UserID: link.UserID}).Return(link, nil).AnyTimes() + timeout := 10 * time.Second + arg := database.AcquireExternalAuthLinkRefreshLeaseParams{ProviderID: link.ProviderID, UserID: link.UserID, TimeoutMs: timeout.Milliseconds()} + dbm.EXPECT().AcquireExternalAuthLinkRefreshLease(gomock.Any(), arg).Return(link, nil).AnyTimes() + check.Args(arg).Asserts(link, policy.ActionUpdatePersonal).Returns(link) + })) + s.Run("ReleaseExternalAuthLinkRefreshLease", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + link := testutil.Fake(s.T(), faker, database.ExternalAuthLink{ + RefreshLeaseExpiresAt: sql.NullTime{Time: dbtime.Now().Add(time.Minute), Valid: true}, + }) + dbm.EXPECT().GetExternalAuthLink(gomock.Any(), database.GetExternalAuthLinkParams{ProviderID: link.ProviderID, UserID: link.UserID}).Return(link, nil).AnyTimes() + arg := database.ReleaseExternalAuthLinkRefreshLeaseParams{ProviderID: link.ProviderID, UserID: link.UserID, RefreshLeaseExpiresAt: link.RefreshLeaseExpiresAt} + dbm.EXPECT().ReleaseExternalAuthLinkRefreshLease(gomock.Any(), arg).Return(nil).AnyTimes() + check.Args(arg).Asserts(link, policy.ActionUpdatePersonal).Returns() + })) s.Run("UpdateUserLink", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { link := testutil.Fake(s.T(), faker, database.UserLink{}) arg := database.UpdateUserLinkParams{OAuthAccessToken: link.OAuthAccessToken, OAuthRefreshToken: link.OAuthRefreshToken, OAuthExpiry: link.OAuthExpiry, UserID: link.UserID, LoginType: link.LoginType, Claims: database.UserLinkClaims{}} diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index ae63a2867688d..492686145adbc 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -105,6 +105,14 @@ func (m queryMetricsStore) DeleteOrganization(ctx context.Context, id uuid.UUID) return r0 } +func (m queryMetricsStore) AcquireExternalAuthLinkRefreshLease(ctx context.Context, arg database.AcquireExternalAuthLinkRefreshLeaseParams) (database.ExternalAuthLink, error) { + start := time.Now() + r0, r1 := m.s.AcquireExternalAuthLinkRefreshLease(ctx, arg) + m.queryLatencies.WithLabelValues("AcquireExternalAuthLinkRefreshLease").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "AcquireExternalAuthLinkRefreshLease").Inc() + return r0, r1 +} + func (m queryMetricsStore) AcquireLock(ctx context.Context, pgAdvisoryXactLock int64) error { start := time.Now() r0 := m.s.AcquireLock(ctx, pgAdvisoryXactLock) @@ -5009,6 +5017,14 @@ func (m queryMetricsStore) RegisterWorkspaceProxy(ctx context.Context, arg datab return r0, r1 } +func (m queryMetricsStore) ReleaseExternalAuthLinkRefreshLease(ctx context.Context, arg database.ReleaseExternalAuthLinkRefreshLeaseParams) error { + start := time.Now() + r0 := m.s.ReleaseExternalAuthLinkRefreshLease(ctx, arg) + m.queryLatencies.WithLabelValues("ReleaseExternalAuthLinkRefreshLease").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ReleaseExternalAuthLinkRefreshLease").Inc() + return r0 +} + func (m queryMetricsStore) RemoveUserFromGroups(ctx context.Context, arg database.RemoveUserFromGroupsParams) ([]uuid.UUID, error) { start := time.Now() r0, r1 := m.s.RemoveUserFromGroups(ctx, arg) @@ -5409,14 +5425,6 @@ func (m queryMetricsStore) UpdateExternalAuthLink(ctx context.Context, arg datab return r0, r1 } -func (m queryMetricsStore) UpdateExternalAuthLinkRefreshToken(ctx context.Context, arg database.UpdateExternalAuthLinkRefreshTokenParams) error { - start := time.Now() - r0 := m.s.UpdateExternalAuthLinkRefreshToken(ctx, arg) - m.queryLatencies.WithLabelValues("UpdateExternalAuthLinkRefreshToken").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateExternalAuthLinkRefreshToken").Inc() - return r0 -} - func (m queryMetricsStore) UpdateGitSSHKey(ctx context.Context, arg database.UpdateGitSSHKeyParams) (database.GitSSHKey, error) { start := time.Now() r0, r1 := m.s.UpdateGitSSHKey(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 6864efcb86ea6..0c08caa669ed8 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -45,6 +45,21 @@ func (m *MockStore) EXPECT() *MockStoreMockRecorder { return m.recorder } +// AcquireExternalAuthLinkRefreshLease mocks base method. +func (m *MockStore) AcquireExternalAuthLinkRefreshLease(ctx context.Context, arg database.AcquireExternalAuthLinkRefreshLeaseParams) (database.ExternalAuthLink, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "AcquireExternalAuthLinkRefreshLease", ctx, arg) + ret0, _ := ret[0].(database.ExternalAuthLink) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// AcquireExternalAuthLinkRefreshLease indicates an expected call of AcquireExternalAuthLinkRefreshLease. +func (mr *MockStoreMockRecorder) AcquireExternalAuthLinkRefreshLease(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcquireExternalAuthLinkRefreshLease", reflect.TypeOf((*MockStore)(nil).AcquireExternalAuthLinkRefreshLease), ctx, arg) +} + // AcquireLock mocks base method. func (m *MockStore) AcquireLock(ctx context.Context, pgAdvisoryXactLock int64) error { m.ctrl.T.Helper() @@ -9504,6 +9519,20 @@ func (mr *MockStoreMockRecorder) RegisterWorkspaceProxy(ctx, arg any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RegisterWorkspaceProxy", reflect.TypeOf((*MockStore)(nil).RegisterWorkspaceProxy), ctx, arg) } +// ReleaseExternalAuthLinkRefreshLease mocks base method. +func (m *MockStore) ReleaseExternalAuthLinkRefreshLease(ctx context.Context, arg database.ReleaseExternalAuthLinkRefreshLeaseParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReleaseExternalAuthLinkRefreshLease", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// ReleaseExternalAuthLinkRefreshLease indicates an expected call of ReleaseExternalAuthLinkRefreshLease. +func (mr *MockStoreMockRecorder) ReleaseExternalAuthLinkRefreshLease(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReleaseExternalAuthLinkRefreshLease", reflect.TypeOf((*MockStore)(nil).ReleaseExternalAuthLinkRefreshLease), ctx, arg) +} + // RemoveUserFromGroups mocks base method. func (m *MockStore) RemoveUserFromGroups(ctx context.Context, arg database.RemoveUserFromGroupsParams) ([]uuid.UUID, error) { m.ctrl.T.Helper() @@ -10237,20 +10266,6 @@ func (mr *MockStoreMockRecorder) UpdateExternalAuthLink(ctx, arg any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateExternalAuthLink", reflect.TypeOf((*MockStore)(nil).UpdateExternalAuthLink), ctx, arg) } -// UpdateExternalAuthLinkRefreshToken mocks base method. -func (m *MockStore) UpdateExternalAuthLinkRefreshToken(ctx context.Context, arg database.UpdateExternalAuthLinkRefreshTokenParams) error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpdateExternalAuthLinkRefreshToken", ctx, arg) - ret0, _ := ret[0].(error) - return ret0 -} - -// UpdateExternalAuthLinkRefreshToken indicates an expected call of UpdateExternalAuthLinkRefreshToken. -func (mr *MockStoreMockRecorder) UpdateExternalAuthLinkRefreshToken(ctx, arg any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateExternalAuthLinkRefreshToken", reflect.TypeOf((*MockStore)(nil).UpdateExternalAuthLinkRefreshToken), ctx, arg) -} - // UpdateGitSSHKey mocks base method. func (m *MockStore) UpdateGitSSHKey(ctx context.Context, arg database.UpdateGitSSHKeyParams) (database.GitSSHKey, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index e32e7513112be..343491f73f015 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -732,6 +732,60 @@ CREATE TYPE workspace_transition AS ENUM ( 'delete' ); +CREATE TABLE external_auth_links ( + provider_id text NOT NULL, + user_id uuid NOT NULL, + created_at timestamp with time zone NOT NULL, + updated_at timestamp with time zone NOT NULL, + oauth_access_token text NOT NULL, + oauth_refresh_token text NOT NULL, + oauth_expiry timestamp with time zone NOT NULL, + oauth_access_token_key_id text, + oauth_refresh_token_key_id text, + oauth_extra jsonb, + oauth_refresh_failure_reason text DEFAULT ''::text NOT NULL, + refresh_lease_expires_at timestamp with time zone +); + +COMMENT ON COLUMN external_auth_links.oauth_access_token_key_id IS 'The ID of the key used to encrypt the OAuth access token. If this is NULL, the access token is not encrypted'; + +COMMENT ON COLUMN external_auth_links.oauth_refresh_token_key_id IS 'The ID of the key used to encrypt the OAuth refresh token. If this is NULL, the refresh token is not encrypted'; + +COMMENT ON COLUMN external_auth_links.oauth_refresh_failure_reason IS 'This error means the refresh token is invalid. Cached so we can avoid calling the external provider again for the same error.'; + +COMMENT ON COLUMN external_auth_links.refresh_lease_expires_at IS 'Indicates a replica is refreshing the token; prevents concurrent refreshes.'; + +CREATE FUNCTION acquire_external_auth_link_refresh_lease(arg_provider_id text, arg_user_id uuid, timeout_ms bigint) RETURNS SETOF external_auth_links + LANGUAGE plpgsql + AS $$ +DECLARE r external_auth_links; +BEGIN + UPDATE external_auth_links + SET + refresh_lease_expires_at = NOW() + (timeout_ms || ' ms')::interval + WHERE + provider_id = arg_provider_id + AND user_id = arg_user_id + AND (refresh_lease_expires_at IS NULL OR refresh_lease_expires_at < NOW()) + RETURNING * INTO r; + -- Got the lease, return the one row. + IF FOUND THEN + RETURN NEXT r; + RETURN; + END IF; + -- Differentiate between unable to get the lease and the row being gone. + IF EXISTS (SELECT 1 FROM external_auth_links WHERE provider_id = arg_provider_id AND user_id = arg_user_id) THEN + RAISE EXCEPTION 'row is currently leased by another replica' + USING ERRCODE = 'check_violation', + CONSTRAINT = 'external_auth_link_active_lease'; + END IF; + -- Row is gone, return nothing. + RETURN; +END; +$$; + +COMMENT ON FUNCTION acquire_external_auth_link_refresh_lease(arg_provider_id text, arg_user_id uuid, timeout_ms bigint) IS 'Acquire a lease on the external auth link and return the row. If there is already an active lease, an exception is raised.'; + CREATE FUNCTION aggregate_usage_event() RETURNS trigger LANGUAGE plpgsql AS $$ @@ -2333,26 +2387,6 @@ COMMENT ON COLUMN dbcrypt_keys.revoked_at IS 'The time at which the key was revo COMMENT ON COLUMN dbcrypt_keys.test IS 'A column used to test the encryption.'; -CREATE TABLE external_auth_links ( - provider_id text NOT NULL, - user_id uuid NOT NULL, - created_at timestamp with time zone NOT NULL, - updated_at timestamp with time zone NOT NULL, - oauth_access_token text NOT NULL, - oauth_refresh_token text NOT NULL, - oauth_expiry timestamp with time zone NOT NULL, - oauth_access_token_key_id text, - oauth_refresh_token_key_id text, - oauth_extra jsonb, - oauth_refresh_failure_reason text DEFAULT ''::text NOT NULL -); - -COMMENT ON COLUMN external_auth_links.oauth_access_token_key_id IS 'The ID of the key used to encrypt the OAuth access token. If this is NULL, the access token is not encrypted'; - -COMMENT ON COLUMN external_auth_links.oauth_refresh_token_key_id IS 'The ID of the key used to encrypt the OAuth refresh token. If this is NULL, the refresh token is not encrypted'; - -COMMENT ON COLUMN external_auth_links.oauth_refresh_failure_reason IS 'This error means the refresh token is invalid. Cached so we can avoid calling the external provider again for the same error.'; - CREATE TABLE files ( hash character varying(64) NOT NULL, created_at timestamp with time zone NOT NULL, diff --git a/coderd/database/migrations/000577_external_auth_lock.down.sql b/coderd/database/migrations/000577_external_auth_lock.down.sql new file mode 100644 index 0000000000000..82f935fabd602 --- /dev/null +++ b/coderd/database/migrations/000577_external_auth_lock.down.sql @@ -0,0 +1,2 @@ +ALTER TABLE external_auth_links DROP COLUMN IF EXISTS refresh_lease_expires_at; +DROP FUNCTION IF EXISTS acquire_external_auth_link_refresh_lease; diff --git a/coderd/database/migrations/000577_external_auth_lock.up.sql b/coderd/database/migrations/000577_external_auth_lock.up.sql new file mode 100644 index 0000000000000..56ee05ae57659 --- /dev/null +++ b/coderd/database/migrations/000577_external_auth_lock.up.sql @@ -0,0 +1,32 @@ +ALTER TABLE external_auth_links ADD COLUMN IF NOT EXISTS refresh_lease_expires_at timestamp WITH time zone DEFAULT NULL; +COMMENT ON COLUMN external_auth_links.refresh_lease_expires_at IS 'Indicates a replica is refreshing the token; prevents concurrent refreshes.'; + +CREATE OR REPLACE FUNCTION acquire_external_auth_link_refresh_lease(arg_provider_id text, arg_user_id uuid, timeout_ms bigint) +RETURNS SETOF external_auth_links AS $$ +DECLARE r external_auth_links; +BEGIN + UPDATE external_auth_links + SET + refresh_lease_expires_at = NOW() + (timeout_ms || ' ms')::interval + WHERE + provider_id = arg_provider_id + AND user_id = arg_user_id + AND (refresh_lease_expires_at IS NULL OR refresh_lease_expires_at < NOW()) + RETURNING * INTO r; + -- Got the lease, return the one row. + IF FOUND THEN + RETURN NEXT r; + RETURN; + END IF; + -- Differentiate between unable to get the lease and the row being gone. + IF EXISTS (SELECT 1 FROM external_auth_links WHERE provider_id = arg_provider_id AND user_id = arg_user_id) THEN + RAISE EXCEPTION 'row is currently leased by another replica' + USING ERRCODE = 'check_violation', + CONSTRAINT = 'external_auth_link_active_lease'; + END IF; + -- Row is gone, return nothing. + RETURN; +END; +$$ LANGUAGE plpgsql; + +COMMENT ON FUNCTION acquire_external_auth_link_refresh_lease IS 'Acquire a lease on the external auth link and return the row. If there is already an active lease, an exception is raised.'; diff --git a/coderd/database/models.go b/coderd/database/models.go index 0e914ec66f594..68130280eb144 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -5343,6 +5343,8 @@ type ExternalAuthLink struct { OAuthExtra pqtype.NullRawMessage `db:"oauth_extra" json:"oauth_extra"` // This error means the refresh token is invalid. Cached so we can avoid calling the external provider again for the same error. OauthRefreshFailureReason string `db:"oauth_refresh_failure_reason" json:"oauth_refresh_failure_reason"` + // Indicates a replica is refreshing the token; prevents concurrent refreshes. + RefreshLeaseExpiresAt sql.NullTime `db:"refresh_lease_expires_at" json:"refresh_lease_expires_at"` } type File struct { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 58c197125b575..ac1945f20f0c9 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -13,6 +13,9 @@ import ( ) type sqlcQuerier interface { + // Set the lease to expire according to the provided timeout. If there is + // already a lease, an exception is raised. + AcquireExternalAuthLinkRefreshLease(ctx context.Context, arg AcquireExternalAuthLinkRefreshLeaseParams) (ExternalAuthLink, error) // Blocks until the lock is acquired. // // This must be called from within a transaction. The lock will be automatically @@ -1329,6 +1332,8 @@ type sqlcQuerier interface { PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) ReduceWorkspaceAgentShareLevelToAuthenticatedByTemplate(ctx context.Context, templateID uuid.UUID) error RegisterWorkspaceProxy(ctx context.Context, arg RegisterWorkspaceProxyParams) (WorkspaceProxy, error) + // The lease is only removed if it is the current lease. + ReleaseExternalAuthLinkRefreshLease(ctx context.Context, arg ReleaseExternalAuthLinkRefreshLeaseParams) error RemoveUserFromGroups(ctx context.Context, arg RemoveUserFromGroupsParams) ([]uuid.UUID, error) // Mutates only created_at on the target row; ids are unchanged so // consumers can keep tracking queued messages by id. @@ -1495,12 +1500,8 @@ type sqlcQuerier interface { // rows in place. UpdateEncryptedAIProviderSettings(ctx context.Context, arg UpdateEncryptedAIProviderSettingsParams) (AIProvider, error) UpdateEncryptedUserAIProviderKey(ctx context.Context, arg UpdateEncryptedUserAIProviderKeyParams) (UserAIProviderKey, error) + // If a refresh lease is provided, the row is only updated if the lease matches. UpdateExternalAuthLink(ctx context.Context, arg UpdateExternalAuthLinkParams) (ExternalAuthLink, error) - // Optimistic lock: only update the row if the refresh token in the database - // still matches the one we read before attempting the refresh. This prevents - // a concurrent caller that lost a token-refresh race from overwriting a valid - // token stored by the winner. - UpdateExternalAuthLinkRefreshToken(ctx context.Context, arg UpdateExternalAuthLinkRefreshTokenParams) error UpdateGitSSHKey(ctx context.Context, arg UpdateGitSSHKeyParams) (GitSSHKey, error) UpdateGroupByID(ctx context.Context, arg UpdateGroupByIDParams) (Group, error) UpdateInactiveUsersToDormant(ctx context.Context, arg UpdateInactiveUsersToDormantParams) ([]UpdateInactiveUsersToDormantRow, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 5849295977668..e364a9400eadf 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -14000,6 +14000,38 @@ func (q *sqlQuerier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest strin return err } +const acquireExternalAuthLinkRefreshLease = `-- name: AcquireExternalAuthLinkRefreshLease :one +SELECT provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason, refresh_lease_expires_at from acquire_external_auth_link_refresh_lease($1, $2, $3) +` + +type AcquireExternalAuthLinkRefreshLeaseParams struct { + ProviderID string `db:"provider_id" json:"provider_id"` + UserID uuid.UUID `db:"user_id" json:"user_id"` + TimeoutMs int64 `db:"timeout_ms" json:"timeout_ms"` +} + +// Set the lease to expire according to the provided timeout. If there is +// already a lease, an exception is raised. +func (q *sqlQuerier) AcquireExternalAuthLinkRefreshLease(ctx context.Context, arg AcquireExternalAuthLinkRefreshLeaseParams) (ExternalAuthLink, error) { + row := q.db.QueryRowContext(ctx, acquireExternalAuthLinkRefreshLease, arg.ProviderID, arg.UserID, arg.TimeoutMs) + var i ExternalAuthLink + err := row.Scan( + &i.ProviderID, + &i.UserID, + &i.CreatedAt, + &i.UpdatedAt, + &i.OAuthAccessToken, + &i.OAuthRefreshToken, + &i.OAuthExpiry, + &i.OAuthAccessTokenKeyID, + &i.OAuthRefreshTokenKeyID, + &i.OAuthExtra, + &i.OauthRefreshFailureReason, + &i.RefreshLeaseExpiresAt, + ) + return i, err +} + const deleteExternalAuthLink = `-- name: DeleteExternalAuthLink :exec DELETE FROM external_auth_links WHERE provider_id = $1 AND user_id = $2 ` @@ -14015,7 +14047,7 @@ func (q *sqlQuerier) DeleteExternalAuthLink(ctx context.Context, arg DeleteExter } const getExternalAuthLink = `-- name: GetExternalAuthLink :one -SELECT provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason FROM external_auth_links WHERE provider_id = $1 AND user_id = $2 +SELECT provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason, refresh_lease_expires_at FROM external_auth_links WHERE provider_id = $1 AND user_id = $2 ` type GetExternalAuthLinkParams struct { @@ -14038,12 +14070,13 @@ func (q *sqlQuerier) GetExternalAuthLink(ctx context.Context, arg GetExternalAut &i.OAuthRefreshTokenKeyID, &i.OAuthExtra, &i.OauthRefreshFailureReason, + &i.RefreshLeaseExpiresAt, ) return i, err } const getExternalAuthLinksByUserID = `-- name: GetExternalAuthLinksByUserID :many -SELECT provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason FROM external_auth_links WHERE user_id = $1 +SELECT provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason, refresh_lease_expires_at FROM external_auth_links WHERE user_id = $1 ` func (q *sqlQuerier) GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]ExternalAuthLink, error) { @@ -14067,6 +14100,7 @@ func (q *sqlQuerier) GetExternalAuthLinksByUserID(ctx context.Context, userID uu &i.OAuthRefreshTokenKeyID, &i.OAuthExtra, &i.OauthRefreshFailureReason, + &i.RefreshLeaseExpiresAt, ); err != nil { return nil, err } @@ -14104,7 +14138,7 @@ INSERT INTO external_auth_links ( $8, $9, $10 -) RETURNING provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason +) RETURNING provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason, refresh_lease_expires_at ` type InsertExternalAuthLinkParams struct { @@ -14146,42 +14180,71 @@ func (q *sqlQuerier) InsertExternalAuthLink(ctx context.Context, arg InsertExter &i.OAuthRefreshTokenKeyID, &i.OAuthExtra, &i.OauthRefreshFailureReason, + &i.RefreshLeaseExpiresAt, ) return i, err } +const releaseExternalAuthLinkRefreshLease = `-- name: ReleaseExternalAuthLinkRefreshLease :exec +UPDATE + external_auth_links +SET + refresh_lease_expires_at = NULL +WHERE + provider_id = $1 + AND user_id = $2 + AND refresh_lease_expires_at = $3 +` + +type ReleaseExternalAuthLinkRefreshLeaseParams struct { + ProviderID string `db:"provider_id" json:"provider_id"` + UserID uuid.UUID `db:"user_id" json:"user_id"` + RefreshLeaseExpiresAt sql.NullTime `db:"refresh_lease_expires_at" json:"refresh_lease_expires_at"` +} + +// The lease is only removed if it is the current lease. +func (q *sqlQuerier) ReleaseExternalAuthLinkRefreshLease(ctx context.Context, arg ReleaseExternalAuthLinkRefreshLeaseParams) error { + _, err := q.db.ExecContext(ctx, releaseExternalAuthLinkRefreshLease, arg.ProviderID, arg.UserID, arg.RefreshLeaseExpiresAt) + return err +} + const updateExternalAuthLink = `-- name: UpdateExternalAuthLink :one UPDATE external_auth_links SET - updated_at = $3, - oauth_access_token = $4, - oauth_access_token_key_id = $5, - oauth_refresh_token = $6, - oauth_refresh_token_key_id = $7, - oauth_expiry = $8, - oauth_extra = $9, - -- Only 'UpdateExternalAuthLinkRefreshToken' supports updating the oauth_refresh_failure_reason. - -- Any updates to the external auth link, will be assumed to change the state and clear - -- any cached errors. - oauth_refresh_failure_reason = '' -WHERE provider_id = $1 AND user_id = $2 RETURNING provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason + updated_at = $4, + oauth_access_token = $5, + oauth_access_token_key_id = $6, + oauth_refresh_token = $7, + oauth_refresh_token_key_id = $8, + oauth_expiry = $9, + oauth_extra = $10, + oauth_refresh_failure_reason = $11 +WHERE + provider_id = $1 + AND user_id = $2 + AND (refresh_lease_expires_at = $3 OR $3 IS NULL) +RETURNING provider_id, user_id, created_at, updated_at, oauth_access_token, oauth_refresh_token, oauth_expiry, oauth_access_token_key_id, oauth_refresh_token_key_id, oauth_extra, oauth_refresh_failure_reason, refresh_lease_expires_at ` type UpdateExternalAuthLinkParams struct { - ProviderID string `db:"provider_id" json:"provider_id"` - UserID uuid.UUID `db:"user_id" json:"user_id"` - UpdatedAt time.Time `db:"updated_at" json:"updated_at"` - OAuthAccessToken string `db:"oauth_access_token" json:"oauth_access_token"` - OAuthAccessTokenKeyID sql.NullString `db:"oauth_access_token_key_id" json:"oauth_access_token_key_id"` - OAuthRefreshToken string `db:"oauth_refresh_token" json:"oauth_refresh_token"` - OAuthRefreshTokenKeyID sql.NullString `db:"oauth_refresh_token_key_id" json:"oauth_refresh_token_key_id"` - OAuthExpiry time.Time `db:"oauth_expiry" json:"oauth_expiry"` - OAuthExtra pqtype.NullRawMessage `db:"oauth_extra" json:"oauth_extra"` -} - + ProviderID string `db:"provider_id" json:"provider_id"` + UserID uuid.UUID `db:"user_id" json:"user_id"` + RefreshLeaseExpiresAt sql.NullTime `db:"refresh_lease_expires_at" json:"refresh_lease_expires_at"` + UpdatedAt time.Time `db:"updated_at" json:"updated_at"` + OAuthAccessToken string `db:"oauth_access_token" json:"oauth_access_token"` + OAuthAccessTokenKeyID sql.NullString `db:"oauth_access_token_key_id" json:"oauth_access_token_key_id"` + OAuthRefreshToken string `db:"oauth_refresh_token" json:"oauth_refresh_token"` + OAuthRefreshTokenKeyID sql.NullString `db:"oauth_refresh_token_key_id" json:"oauth_refresh_token_key_id"` + OAuthExpiry time.Time `db:"oauth_expiry" json:"oauth_expiry"` + OAuthExtra pqtype.NullRawMessage `db:"oauth_extra" json:"oauth_extra"` + OauthRefreshFailureReason string `db:"oauth_refresh_failure_reason" json:"oauth_refresh_failure_reason"` +} + +// If a refresh lease is provided, the row is only updated if the lease matches. func (q *sqlQuerier) UpdateExternalAuthLink(ctx context.Context, arg UpdateExternalAuthLinkParams) (ExternalAuthLink, error) { row := q.db.QueryRowContext(ctx, updateExternalAuthLink, arg.ProviderID, arg.UserID, + arg.RefreshLeaseExpiresAt, arg.UpdatedAt, arg.OAuthAccessToken, arg.OAuthAccessTokenKeyID, @@ -14189,6 +14252,7 @@ func (q *sqlQuerier) UpdateExternalAuthLink(ctx context.Context, arg UpdateExter arg.OAuthRefreshTokenKeyID, arg.OAuthExpiry, arg.OAuthExtra, + arg.OauthRefreshFailureReason, ) var i ExternalAuthLink err := row.Scan( @@ -14203,57 +14267,11 @@ func (q *sqlQuerier) UpdateExternalAuthLink(ctx context.Context, arg UpdateExter &i.OAuthRefreshTokenKeyID, &i.OAuthExtra, &i.OauthRefreshFailureReason, + &i.RefreshLeaseExpiresAt, ) return i, err } -const updateExternalAuthLinkRefreshToken = `-- name: UpdateExternalAuthLinkRefreshToken :exec -UPDATE - external_auth_links -SET - -- oauth_refresh_failure_reason can be set to cache the failure reason - -- for subsequent refresh attempts. - oauth_refresh_failure_reason = $1, - oauth_refresh_token = $2, - updated_at = $3 -WHERE - provider_id = $4 -AND - user_id = $5 -AND - oauth_refresh_token = $6 -AND - -- Required for sqlc to generate a parameter for the oauth_refresh_token_key_id - $7 :: text = $7 :: text -` - -type UpdateExternalAuthLinkRefreshTokenParams struct { - OauthRefreshFailureReason string `db:"oauth_refresh_failure_reason" json:"oauth_refresh_failure_reason"` - OAuthRefreshToken string `db:"oauth_refresh_token" json:"oauth_refresh_token"` - UpdatedAt time.Time `db:"updated_at" json:"updated_at"` - ProviderID string `db:"provider_id" json:"provider_id"` - UserID uuid.UUID `db:"user_id" json:"user_id"` - OldOauthRefreshToken string `db:"old_oauth_refresh_token" json:"old_oauth_refresh_token"` - OAuthRefreshTokenKeyID string `db:"oauth_refresh_token_key_id" json:"oauth_refresh_token_key_id"` -} - -// Optimistic lock: only update the row if the refresh token in the database -// still matches the one we read before attempting the refresh. This prevents -// a concurrent caller that lost a token-refresh race from overwriting a valid -// token stored by the winner. -func (q *sqlQuerier) UpdateExternalAuthLinkRefreshToken(ctx context.Context, arg UpdateExternalAuthLinkRefreshTokenParams) error { - _, err := q.db.ExecContext(ctx, updateExternalAuthLinkRefreshToken, - arg.OauthRefreshFailureReason, - arg.OAuthRefreshToken, - arg.UpdatedAt, - arg.ProviderID, - arg.UserID, - arg.OldOauthRefreshToken, - arg.OAuthRefreshTokenKeyID, - ) - return err -} - const getFileByHashAndCreator = `-- name: GetFileByHashAndCreator :one SELECT hash, created_at, created_by, mimetype, data, id diff --git a/coderd/database/queries/externalauth.sql b/coderd/database/queries/externalauth.sql index e5d0ec548bf47..aeaf714420f41 100644 --- a/coderd/database/queries/externalauth.sql +++ b/coderd/database/queries/externalauth.sql @@ -33,39 +33,34 @@ INSERT INTO external_auth_links ( ) RETURNING *; -- name: UpdateExternalAuthLink :one +-- If a refresh lease is provided, the row is only updated if the lease matches. UPDATE external_auth_links SET - updated_at = $3, - oauth_access_token = $4, - oauth_access_token_key_id = $5, - oauth_refresh_token = $6, - oauth_refresh_token_key_id = $7, - oauth_expiry = $8, - oauth_extra = $9, - -- Only 'UpdateExternalAuthLinkRefreshToken' supports updating the oauth_refresh_failure_reason. - -- Any updates to the external auth link, will be assumed to change the state and clear - -- any cached errors. - oauth_refresh_failure_reason = '' -WHERE provider_id = $1 AND user_id = $2 RETURNING *; + updated_at = $4, + oauth_access_token = $5, + oauth_access_token_key_id = $6, + oauth_refresh_token = $7, + oauth_refresh_token_key_id = $8, + oauth_expiry = $9, + oauth_extra = $10, + oauth_refresh_failure_reason = $11 +WHERE + provider_id = $1 + AND user_id = $2 + AND (refresh_lease_expires_at = $3 OR $3 IS NULL) +RETURNING *; + +-- name: AcquireExternalAuthLinkRefreshLease :one +-- Set the lease to expire according to the provided timeout. If there is +-- already a lease, an exception is raised. +SELECT * from acquire_external_auth_link_refresh_lease(@provider_id, @user_id, @timeout_ms); --- name: UpdateExternalAuthLinkRefreshToken :exec --- Optimistic lock: only update the row if the refresh token in the database --- still matches the one we read before attempting the refresh. This prevents --- a concurrent caller that lost a token-refresh race from overwriting a valid --- token stored by the winner. +-- name: ReleaseExternalAuthLinkRefreshLease :exec +-- The lease is only removed if it is the current lease. UPDATE external_auth_links SET - -- oauth_refresh_failure_reason can be set to cache the failure reason - -- for subsequent refresh attempts. - oauth_refresh_failure_reason = @oauth_refresh_failure_reason, - oauth_refresh_token = @oauth_refresh_token, - updated_at = @updated_at + refresh_lease_expires_at = NULL WHERE - provider_id = @provider_id -AND - user_id = @user_id -AND - oauth_refresh_token = @old_oauth_refresh_token -AND - -- Required for sqlc to generate a parameter for the oauth_refresh_token_key_id - @oauth_refresh_token_key_id :: text = @oauth_refresh_token_key_id :: text; + provider_id = @provider_id + AND user_id = @user_id + AND refresh_lease_expires_at = @refresh_lease_expires_at; diff --git a/coderd/externalauth.go b/coderd/externalauth.go index 51b7727c00d69..04ab62d3af66d 100644 --- a/coderd/externalauth.go +++ b/coderd/externalauth.go @@ -203,15 +203,17 @@ func (api *API) postExternalAuthDeviceByID(rw http.ResponseWriter, r *http.Reque } } else { _, err = api.Database.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ - ProviderID: config.ID, - UserID: apiKey.UserID, - UpdatedAt: dbtime.Now(), - OAuthAccessToken: token.AccessToken, - OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required - OAuthRefreshToken: token.RefreshToken, - OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required - OAuthExpiry: token.Expiry, - OAuthExtra: pqtype.NullRawMessage{}, + ProviderID: config.ID, + UserID: apiKey.UserID, + UpdatedAt: dbtime.Now(), + OAuthAccessToken: token.AccessToken, + OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthRefreshToken: token.RefreshToken, + OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthExpiry: token.Expiry, + OAuthExtra: pqtype.NullRawMessage{}, + OauthRefreshFailureReason: "", + RefreshLeaseExpiresAt: sql.NullTime{}, }) if err != nil { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ @@ -305,15 +307,17 @@ func (api *API) externalAuthCallback(externalAuthConfig *externalauth.Config) ht } } else { _, err = api.Database.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ - ProviderID: externalAuthConfig.ID, - UserID: apiKey.UserID, - UpdatedAt: dbtime.Now(), - OAuthAccessToken: state.Token.AccessToken, - OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required - OAuthRefreshToken: state.Token.RefreshToken, - OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required - OAuthExpiry: state.Token.Expiry, - OAuthExtra: extra, + ProviderID: externalAuthConfig.ID, + UserID: apiKey.UserID, + UpdatedAt: dbtime.Now(), + OAuthAccessToken: state.Token.AccessToken, + OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthRefreshToken: state.Token.RefreshToken, + OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthExpiry: state.Token.Expiry, + OAuthExtra: extra, + OauthRefreshFailureReason: "", + RefreshLeaseExpiresAt: sql.NullTime{}, }) if err != nil { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go index c6d87e34b4f31..e4e9809fe2f85 100644 --- a/coderd/externalauth/externalauth.go +++ b/coderd/externalauth/externalauth.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "encoding/json" + "errors" "fmt" "io" "mime" @@ -55,6 +56,18 @@ const ( // defaultRefreshRetryTimeout bounds the total time spent retrying a // transient refresh failure across all attempts. defaultRefreshRetryTimeout = 10 * time.Second + + // defaultRefreshLeaseInitialBackoff is the starting wait between polls to + // check whether another replica has finished refreshing. + defaultRefreshLeaseInitialBackoff = 50 * time.Millisecond + + // defaultRefreshLeaseMaxBackoff is the maximum wait between polls to check + // whether another replica has finished refreshing. + defaultRefreshLeaseMaxBackoff = 500 * time.Millisecond + + // externalAuthLinkActiveLeaseConstraint indicates the lease could not be + // acquired because something else has an active lease. + externalAuthLinkActiveLeaseConstraint database.CheckConstraint = "external_auth_link_active_lease" ) // SingleflightGroup exposes a subset of singleflight.Group for easier testing. @@ -164,6 +177,14 @@ type Config struct { // RefreshGroup deduplicates concurrent requests. RefreshGroup SingleflightGroup + // RefreshLeaseInitialBackoff is the starting wait between polls to check + // whether another replica has finished refreshing. + RefreshLeaseInitialBackoff time.Duration + + // RefreshLeaseMaxBackoff is the maximum wait between polls to check whether + // another replica has finished refreshing. + RefreshLeaseMaxBackoff time.Duration + gitProviderMu sync.Mutex // gitProvider memoizes the provider so the GitHub ETag response // cache survives across Git calls. @@ -232,7 +253,15 @@ func IsInvalidTokenError(err error) bool { } // RefreshToken automatically refreshes the token if expired and permitted. +// Tokens are then validated, whether or not they were refreshed. func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAuthLink database.ExternalAuthLink) (database.ExternalAuthLink, error) { + // If the token is expired and refresh is disabled, prompt the user to + // authenticate manually again. + token := externalAuthLink.OAuthToken() + if c.NoRefresh && !token.Valid() { + return externalAuthLink, InvalidTokenError("token expired and refreshing is disabled") + } + // Prevent parallel refreshes by waiting for the result of any already // in-flight refresh. Otherwise, the parallel calls will fail with a bad // refresh token error as they can only be used once. @@ -246,9 +275,24 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu if c.RefreshRetryTimeout > 0 { timeout += c.RefreshRetryTimeout } - rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) + ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) defer cancel() - return c.innerRefreshToken(rctx, db, externalAuthLink) + + // Although TokenSource().Token() will also check if the token is expired, + // do so ahead of time here to avoid the overhead of acquiring the lease + // when no refresh is required. Also validate the token under this lock, + // since we already have it anyway. + if !token.Valid() { + return c.refreshAndValidateWithLease(ctx, db, externalAuthLink, timeout) + } + + // Validate the token even if we did not refresh. This is done within the + // group but outside the lease, meaning multiple instances may validate at + // the same time, but the lock overhead seems greater than the overhead of + // occasional concurrent validation, and is not strictly necessary like it + // is with refreshing, so avoid the lock when not refreshing. + _, err := c.validateWithRetry(ctx, token) + return externalAuthLink, err }) select { case results := <-ch: @@ -263,37 +307,100 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu } } -func (c *Config) innerRefreshToken(ctx context.Context, db database.Store, externalAuthLink database.ExternalAuthLink) (database.ExternalAuthLink, error) { - // If the token is expired and refresh is disabled, we prompt - // the user to authenticate again. - if c.NoRefresh && - // If the time is set to 0, then it should never expire. - // This is true for github, which has no expiry. - !externalAuthLink.OAuthExpiry.IsZero() && - externalAuthLink.OAuthExpiry.Before(dbtime.Now()) { - return externalAuthLink, InvalidTokenError("token expired, refreshing is either disabled or refreshing failed and will not be retried") +// refreshAndValidateWithLease wraps the refresh and subsequent validation with +// concurrency protection between multiple instances by using a lease column on +// the link's row. +func (c *Config) refreshAndValidateWithLease(ctx context.Context, db database.Store, externalAuthLink database.ExternalAuthLink, timeout time.Duration) (newLink database.ExternalAuthLink, refreshErr error) { + // There may be other replicas also wanting to refresh; try to get a lease + // on the row. This also ensures we have the latest link. + var leasedLink database.ExternalAuthLink + initial := defaultRefreshLeaseInitialBackoff + if c.RefreshLeaseInitialBackoff > 0 { + initial = c.RefreshLeaseInitialBackoff + } + maximum := defaultRefreshLeaseMaxBackoff + if c.RefreshLeaseMaxBackoff > 0 { + maximum = c.RefreshLeaseMaxBackoff + } + r := retry.New(initial, maximum) + // Make sure to release the lease if we manage to get one before returning. + defer func() { + if leasedLink.RefreshLeaseExpiresAt.Valid { + refreshErr = errors.Join(refreshErr, db.ReleaseExternalAuthLinkRefreshLease(ctx, database.ReleaseExternalAuthLinkRefreshLeaseParams{ + ProviderID: externalAuthLink.ProviderID, + UserID: externalAuthLink.UserID, + // The row will only update if we hold the current lease. This is + // somewhat redundant since if we lost the lease our context would be + // expired anyway, so it is not actually possible to get the sql.ErrNoRows + // that would result from this. + RefreshLeaseExpiresAt: leasedLink.RefreshLeaseExpiresAt, + })) + } + }() + for !leasedLink.RefreshLeaseExpiresAt.Valid { + // Acquiring the lease also returns the link, so we can get the expiry date + // that the database sets and check to see if it was refreshed in the + // meantime. + var err error + leasedLink, err = db.AcquireExternalAuthLinkRefreshLease(ctx, database.AcquireExternalAuthLinkRefreshLeaseParams{ + ProviderID: externalAuthLink.ProviderID, + UserID: externalAuthLink.UserID, + TimeoutMs: timeout.Milliseconds(), + }) + switch { + // Something still holds the lock; keep waiting. + case database.IsCheckViolation(err, externalAuthLinkActiveLeaseConstraint): + if !r.Wait(ctx) { + return externalAuthLink, ctx.Err() + } + // Some kind of DB error or the row does not exist. + case err != nil: + return externalAuthLink, err + // Something else refreshed either while we were waiting or before we got a + // hold of the lease but after we initially fetched the link. + case leasedLink.OAuthRefreshToken != externalAuthLink.OAuthRefreshToken: + if leasedLink.OauthRefreshFailureReason != "" { + return externalAuthLink, refreshError(leasedLink, leasedLink.OauthRefreshFailureReason) + } + return leasedLink, nil + } } - refreshToken := externalAuthLink.OAuthRefreshToken + // Otherwise the token has still not been updated; refresh it now. + newLink, refreshErr = c.refreshAndValidateToken(ctx, db, leasedLink) + return newLink, refreshErr +} - // This is additional defensive programming. Because TokenSource is an interface, - // we cannot be sure that the implementation will treat an 'IsZero' time - // as "not-expired". The default implementation does, but a custom implementation - // might not. Removing the refreshToken will guarantee a refresh will fail. - if c.NoRefresh { - refreshToken = "" - } +// refreshError converts a failure reason to an error. +func refreshError(link database.ExternalAuthLink, reason string) error { + return InvalidTokenError(fmt.Sprintf("token expired and refreshing failed %s with: %s", + // Do not return the exact time, because then we have to know what timezone + // the user is in. This approximate time is good enough. + humanize.Time(link.UpdatedAt), + reason, + )) +} - existingToken := &oauth2.Token{ - AccessToken: externalAuthLink.OAuthAccessToken, - RefreshToken: refreshToken, - Expiry: externalAuthLink.OAuthExpiry, +// refreshAndValidateToken does the actual token refresh, persists the result to +// the database, then validates the token. The provided link must be up to date +// with the currently held lease. +func (c *Config) refreshAndValidateToken(ctx context.Context, db database.Store, externalAuthLink database.ExternalAuthLink) (database.ExternalAuthLink, error) { + existingToken := externalAuthLink.OAuthToken() + + // This is additional defensive programming. Because TokenSource is an + // interface, we cannot be sure that the implementation will treat an 'IsZero' + // time as "not-expired". The default implementation does, but a custom + // implementation might not. Removing the refresh token will guarantee a + // refresh will fail. + if c.NoRefresh { + existingToken.RefreshToken = "" } // NOTE: TokenSource(...).Token() will short-circuit if the token: // - is not expired (returns original token) // - is expired and has no refresh token (returns error) - // This means we will avoid making useless HTTP requests. + // This means we will avoid making useless HTTP requests, and only get errors + // when an actual refresh attempt is made. // // External providers (GitHub in particular) intermittently fail token // refreshes with transient errors such as 5xx responses, network timeouts, @@ -304,50 +411,41 @@ func (c *Config) innerRefreshToken(ctx context.Context, db database.Store, exter // will never succeed and retrying wastes the refresh quota. token, err := c.refreshTokenWithRetry(ctx, existingToken) if err != nil { - // A refresh attempt can fail for numerous reasons. If it fails because - // of a bad refresh token, then the refresh token is invalid, and we - // should get rid of it. Keeping it around will cause additional refresh - // attempts that will fail and cost us api rate limits. - // - // The error message is saved for debugging purposes. + // A refresh attempt can fail for numerous reasons. If it fails because of a + // bad refresh token, then the refresh token is invalid, and we should get + // rid of it. Keeping it around will cause additional refresh attempts that + // will fail and cost us api rate limits. Also save the error message for + // debugging purposes. if isFailedRefresh(existingToken, err) { - // Before caching the failure, re-read the external auth link from the - // database. A nearly-concurrent request may have already refreshed the - // token successfully, consuming the single-use refresh token (e.g., - // GitHub App tokens). In that case our "bad_refresh_token" error is a - // false positive from losing the race, and we should use the winner's - // updated token instead of poisoning the database with a cached failure. - currentLink, readErr := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ - ProviderID: externalAuthLink.ProviderID, - UserID: externalAuthLink.UserID, - }) - if readErr == nil && currentLink.OAuthRefreshToken != externalAuthLink.OAuthRefreshToken { - return currentLink, nil - } - reason := err.Error() if len(reason) > failureReasonLimit { // Limit the length of the error message to prevent // spamming the database with long error messages. reason = reason[:failureReasonLimit] } - dbExecErr := db.UpdateExternalAuthLinkRefreshToken(ctx, database.UpdateExternalAuthLinkRefreshTokenParams{ - // Adding a reason will prevent further attempts to try and refresh the token. - OauthRefreshFailureReason: reason, - // Remove the invalid refresh token so it is never used again. The cached - // `reason` can be used to know why this field was zeroed out. + _, updateErr := db.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ + ProviderID: externalAuthLink.ProviderID, + UserID: externalAuthLink.UserID, + UpdatedAt: dbtime.Now(), + // Remove the invalid refresh token so it is never used again. OAuthRefreshToken: "", - OAuthRefreshTokenKeyID: externalAuthLink.OAuthRefreshTokenKeyID.String, - UpdatedAt: dbtime.Now(), - ProviderID: externalAuthLink.ProviderID, - UserID: externalAuthLink.UserID, - // Optimistic lock: only clear the token if it hasn't been - // updated by a concurrent caller that won the refresh race. - OldOauthRefreshToken: externalAuthLink.OAuthRefreshToken, + OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required + // The cached reason can be used to know why the token was zeroed out. + OauthRefreshFailureReason: reason, + // Preserve the access token, expiry, and extra info as they are. + OAuthAccessToken: externalAuthLink.OAuthAccessToken, + OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthExpiry: externalAuthLink.OAuthExpiry, + OAuthExtra: externalAuthLink.OAuthExtra, + // The row will only update if we hold the current lease. This is + // somewhat redundant since if we lost the lease our context would be + // expired anyway, so it is not actually possible to get the + // sql.ErrNoRows that would result from this. + RefreshLeaseExpiresAt: externalAuthLink.RefreshLeaseExpiresAt, }) - if dbExecErr != nil { + if updateErr != nil { // This error should be rare. - return externalAuthLink, InvalidTokenError(fmt.Sprintf("refresh token failed: %q, then removing refresh token failed: %q", err.Error(), dbExecErr.Error())) + return externalAuthLink, InvalidTokenError(fmt.Sprintf("refresh token failed: %q, then removing refresh token failed: %q", err.Error(), updateErr.Error())) } // The refresh token was cleared externalAuthLink.OAuthRefreshToken = "" @@ -365,19 +463,12 @@ func (c *Config) innerRefreshToken(ctx context.Context, db database.Store, exter if externalAuthLink.OauthRefreshFailureReason != "" { // A cached refresh failure error exists. So the refresh token was set, but was invalid, and zeroed out. // Return this cached error for the original refresh attempt. - return externalAuthLink, InvalidTokenError(fmt.Sprintf("token expired and refreshing failed %s with: %s", - // Do not return the exact time, because then we have to know what timezone the - // user is in. This approximate time is good enough. - humanize.Time(externalAuthLink.UpdatedAt), - externalAuthLink.OauthRefreshFailureReason, - )) + return externalAuthLink, refreshError(externalAuthLink, externalAuthLink.OauthRefreshFailureReason) } - return externalAuthLink, InvalidTokenError("token expired, refreshing is either disabled or refreshing failed and will not be retried") + return externalAuthLink, InvalidTokenError("token expired, refreshing failed and will not be retried") } - // Non-expired tokens are short-circuited as noted above; reaching here - // means refresh failed. return externalAuthLink, InvalidTokenError(fmt.Sprintf("refresh token: %s", err.Error())) } @@ -386,71 +477,53 @@ func (c *Config) innerRefreshToken(ctx context.Context, db database.Store, exter return externalAuthLink, xerrors.Errorf("generate token extra: %w", err) } - // Persist the refreshed token to the DB before validation. GitHub - // rotates refresh tokens on every use, so the old refresh token is - // already invalid on the IDP side. If we validated first and the - // validation endpoint was unavailable (e.g. rate-limited 403), the - // new token would be silently lost and the user would be forced to - // re-authenticate manually. - originalAccessToken := externalAuthLink.OAuthAccessToken - if token.AccessToken != originalAccessToken { - updatedAuthLink, err := db.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ - ProviderID: c.ID, - UserID: externalAuthLink.UserID, - UpdatedAt: dbtime.Now(), - OAuthAccessToken: token.AccessToken, - OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required - OAuthRefreshToken: token.RefreshToken, - OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required - OAuthExpiry: token.Expiry, - OAuthExtra: extra, - }) - if err != nil { - return updatedAuthLink, xerrors.Errorf("persist refreshed token: %w", err) - } - externalAuthLink = updatedAuthLink + // Persist the refreshed token to the DB before validation. GitHub rotates + // refresh tokens on every use, so the old refresh token is already invalid on + // the IDP side. If we validated first and the validation endpoint was + // unavailable (e.g. rate-limited 403), the new token would be silently lost + // and the user would be forced to re-authenticate manually. + updatedAuthLink, err := db.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ + ProviderID: externalAuthLink.ProviderID, + UserID: externalAuthLink.UserID, + UpdatedAt: dbtime.Now(), + OAuthAccessToken: token.AccessToken, + OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthRefreshToken: token.RefreshToken, + OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthExpiry: token.Expiry, + OAuthExtra: extra, + // If there was any failure before, we can clear it now. + OauthRefreshFailureReason: "", + // The row will only update if we hold the current lease. This is somewhat + // redundant since if we lost the lease our context would be expired anyway, + // so it is not actually possible to get the sql.ErrNoRows that would result + // from this. + RefreshLeaseExpiresAt: externalAuthLink.RefreshLeaseExpiresAt, + }) + if err != nil { + return updatedAuthLink, xerrors.Errorf("persist refreshed token: %w", err) } - r := retry.New(50*time.Millisecond, 200*time.Millisecond) - // See the comment below why the retry and cancel is required. - retryCtx, retryCtxCancel := context.WithTimeout(ctx, time.Second) - defer retryCtxCancel() -validate: - valid, user, err := c.ValidateToken(ctx, token) + user, err := c.validateWithRetry(ctx, token) if err != nil { - return externalAuthLink, xerrors.Errorf("validate external auth token: %w", err) - } - if !valid { - // A customer using GitHub in Australia reported that validating immediately - // after refreshing the token would intermittently fail with a 401. Waiting - // a few milliseconds with the exact same token on the exact same request - // would resolve the issue. It seems likely that the write is not propagating - // to the read replica in time. - // - // We do an exponential backoff here to give the write time to propagate. - if c.Type == string(codersdk.EnhancedExternalAuthProviderGitHub) && r.Wait(retryCtx) { - goto validate - } - // The token is no longer valid! - return externalAuthLink, InvalidTokenError("token failed to validate") + return updatedAuthLink, err } - // Update the associated user's github.com user ID if the token - // is for github.com and validation returned user info. - if token.AccessToken != originalAccessToken && IsGithubDotComURL(c.AuthCodeURL("")) && user != nil { + // Update the associated user's github.com user ID if the token is for + // github.com and validation returned user info. + if IsGithubDotComURL(c.AuthCodeURL("")) && user != nil { err = db.UpdateUserGithubComUserID(ctx, database.UpdateUserGithubComUserIDParams{ - ID: externalAuthLink.UserID, + ID: updatedAuthLink.UserID, GithubComUserID: sql.NullInt64{ Int64: user.ID, Valid: true, }, }) if err != nil { - return externalAuthLink, xerrors.Errorf("update user github com user id: %w", err) + return updatedAuthLink, xerrors.Errorf("update user github com user id: %w", err) } } - - return externalAuthLink, nil + return updatedAuthLink, nil } // refreshTokenWithRetry exchanges the refresh token for a new access token, @@ -489,12 +562,8 @@ func (c *Config) refreshTokenWithRetry(ctx context.Context, existingToken *oauth defer retryCancel() backoff := retry.New(initial, maximum) - var ( - token *oauth2.Token - err error - ) for { - token, err = c.TokenSource(ctx, existingToken).Token() + token, err := c.TokenSource(ctx, existingToken).Token() if err == nil || isFailedRefresh(existingToken, err) { return token, err } @@ -512,17 +581,48 @@ func (c *Config) refreshTokenWithRetry(ctx context.Context, existingToken *oauth } } +// validateWithRetry validates the provided link, retrying on failure. On +// success return the user info. +func (c *Config) validateWithRetry(ctx context.Context, token *oauth2.Token) (*codersdk.ExternalAuthUser, error) { + r := retry.New(50*time.Millisecond, 200*time.Millisecond) + // See the comment below why the retry and cancel is required. + retryCtx, retryCtxCancel := context.WithTimeout(ctx, time.Second) + defer retryCtxCancel() + +validate: + valid, user, err := c.ValidateToken(ctx, token) + if err != nil { + return nil, xerrors.Errorf("validate external auth token: %w", err) + } + if !valid { + // A customer using GitHub in Australia reported that validating immediately + // after refreshing the token would intermittently fail with a 401. Waiting + // a few milliseconds with the exact same token on the exact same request + // would resolve the issue. It seems likely that the write is not propagating + // to the read replica in time. + // + // We do an exponential backoff here to give the write time to propagate. + if c.Type == string(codersdk.EnhancedExternalAuthProviderGitHub) && r.Wait(retryCtx) { + goto validate + } + // The token is no longer valid! + return nil, InvalidTokenError("token failed to validate") + } + + return user, nil +} + // ValidateToken checks if the Git token provided is valid. // The user is optionally returned if the provider supports it. // Returns valid=true when: the provider confirmed the token, // no ValidateURL is configured, or the validation endpoint // returned a rate-limited response (403 with rate-limit headers // or 429). -func (c *Config) ValidateToken(ctx context.Context, link *oauth2.Token) (bool, *codersdk.ExternalAuthUser, error) { - if link == nil { +func (c *Config) ValidateToken(ctx context.Context, token *oauth2.Token) (bool, *codersdk.ExternalAuthUser, error) { + if token == nil { return false, nil, xerrors.New("validate external auth token: token is nil") } - if !link.Expiry.IsZero() && link.Expiry.Before(dbtime.Now()) { + if !token.Valid() { return false, nil, nil } @@ -535,7 +635,7 @@ func (c *Config) ValidateToken(ctx context.Context, link *oauth2.Token) (bool, * return false, nil, err } - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", link.AccessToken)) + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken)) res, err := c.InstrumentedOAuth2Config.Do(ctx, promoauth.SourceValidateToken, req) if err != nil { return false, nil, err diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go index 9e954219d82fa..1001668dce356 100644 --- a/coderd/externalauth/externalauth_test.go +++ b/coderd/externalauth/externalauth_test.go @@ -3,6 +3,7 @@ package externalauth_test import ( "bytes" "context" + "database/sql" "encoding/json" "fmt" "io" @@ -19,6 +20,7 @@ import ( "github.com/coreos/go-oidc/v3/oidc" "github.com/golang-jwt/jwt/v4" "github.com/google/uuid" + "github.com/lib/pq" "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -33,6 +35,7 @@ import ( "github.com/coder/coder/v2/coderd" "github.com/coder/coder/v2/coderd/coderdtest/oidctest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" @@ -117,6 +120,7 @@ func TestRefreshToken(t *testing.T) { t.Run("NoRefreshExpired", func(t *testing.T) { t.Parallel() + fake, config, link := setupOauth2Test(t, testConfig{ FakeIDPOpts: []oidctest.FakeIDPOpt{ oidctest.WithRefresh(func(_ string) error { @@ -132,17 +136,22 @@ func TestRefreshToken(t *testing.T) { }, ExternalAuthOpt: func(cfg *externalauth.Config) { cfg.NoRefresh = true + // Should abort before entering the group. + cfg.RefreshGroup = nil + }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - // Expire the link - link.OAuthExpiry = expired + mDB := mockDB(t) - _, err := config.RefreshToken(ctx, nil, link) + // There should be no database calls since we return early. + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + _, err := config.RefreshToken(ctx, mDB, link) require.Error(t, err) require.True(t, externalauth.IsInvalidTokenError(err)) - require.Contains(t, err.Error(), "refreshing is either disabled or refreshing failed") + require.Contains(t, err.Error(), "token expired and refreshing is disabled") }) // NoRefreshNoExpiry tests that an oauth token without an expiry is always valid. @@ -167,18 +176,23 @@ func TestRefreshToken(t *testing.T) { }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) + mDB := mockDB(t) + + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) // Zero time used link.OAuthExpiry = time.Time{} - _, err := config.RefreshToken(ctx, nil, link) + // Since the token is not expired, no refresh lease will be acquired and + // it will only be validated. + _, err := config.RefreshToken(ctx, mDB, link) require.NoError(t, err) require.True(t, validated, "token should have been validated") }) t.Run("FalseIfTokenSourceFails", func(t *testing.T) { t.Parallel() + config := &externalauth.Config{ InstrumentedOAuth2Config: &testutil.OAuth2Config{ TokenSourceFunc: func() (*oauth2.Token, error) { @@ -188,9 +202,14 @@ func TestRefreshToken(t *testing.T) { RefreshGroup: new(singleflight.Group), } - _, err := config.RefreshToken(context.Background(), nil, database.ExternalAuthLink{ + link := database.ExternalAuthLink{ OAuthExpiry: expired, - }) + } + + mDB := mockDB(t, withLease(link)) + + ctx := testutil.Context(t, testutil.WaitLong) + _, err := config.RefreshToken(ctx, mDB, link) require.Error(t, err) require.True(t, externalauth.IsInvalidTokenError(err)) require.Contains(t, err.Error(), "failure") @@ -199,11 +218,6 @@ func TestRefreshToken(t *testing.T) { t.Run("ValidateServerError", func(t *testing.T) { t.Parallel() - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) - mDB.EXPECT().UpdateExternalAuthLink(gomock.Any(), gomock.Any()). - Return(database.ExternalAuthLink{}, nil).AnyTimes() - const staticError = "static error" validated := false fake, config, link := setupOauth2Test(t, testConfig{ @@ -217,7 +231,11 @@ func TestRefreshToken(t *testing.T) { }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) + mDB := mockDB(t, + withLease(link), + withUpdatePassthrough()) + + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) link.OAuthExpiry = expired _, err := config.RefreshToken(ctx, mDB, link) @@ -229,10 +247,10 @@ func TestRefreshToken(t *testing.T) { require.True(t, validated, "token should have been attempted to be validated") }) - // RefreshRetries tests that refresh token retry behavior works as expected. - // If a refresh token fails because the token itself is invalid, no more - // refresh attempts should ever happen. An invalid refresh token does - // not magically become valid at some point in the future. + // RefreshRetries tests that refresh token external retry behavior works as + // expected. If a refresh token fails because the token itself is invalid, no + // more refresh attempts should ever happen. An invalid refresh token does not + // magically become valid at some point in the future. // // Internal retries are disabled in this subtest via a negative // RefreshRetryTimeout so each RefreshToken call results in exactly one @@ -242,10 +260,6 @@ func TestRefreshToken(t *testing.T) { t.Parallel() var refreshErr *oauth2.RetrieveError - - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) - refreshCount := 0 fake, config, link := setupOauth2Test(t, testConfig{ FakeIDPOpts: []oidctest.FakeIDPOpt{ @@ -266,11 +280,21 @@ func TestRefreshToken(t *testing.T) { // (Windows). cfg.RefreshRetryTimeout = -1 }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - // Expire the link - link.OAuthExpiry = expired + // Allow acquiring and releasing the lease for all the temporary error + // attempts, the bad refresh token attempt, and then finally the last + // attempt with no refresh token set. + mDB := mockDB(t, + withLease(link), + withLease(link), + withLease(link), + withLease(link)) + + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) // Make the failure a server internal error. Not related to the token // This should be retried since this error is temporary. @@ -290,29 +314,56 @@ func TestRefreshToken(t *testing.T) { require.Equal(t, refreshCount, totalRefreshes) } - // Try again with a bad refresh token error. This will invalidate the - // refresh token, and not retry again. Expect DB calls to check for - // concurrent refresh (GetExternalAuthLink) and then remove the refresh token. - mDB.EXPECT().GetExternalAuthLink(gomock.Any(), gomock.Any()).Return(link, nil).Times(1) - mDB.EXPECT().UpdateExternalAuthLinkRefreshToken(gomock.Any(), gomock.Any()).Return(nil).Times(1) + // The final attempt will be a permanent error and we should see the + // database update with the error. Need to extract it from the mock call + // like this rather than use the returned link from RefreshToken as it does + // not return the updated link in the error case. + mDB.EXPECT().UpdateExternalAuthLink(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, p database.UpdateExternalAuthLinkParams) (database.ExternalAuthLink, error) { + link = database.ExternalAuthLink{ + ProviderID: p.ProviderID, + UserID: p.UserID, + OAuthAccessToken: p.OAuthAccessToken, + // This should be zeroed out. + OAuthRefreshToken: p.OAuthRefreshToken, + OAuthExpiry: p.OAuthExpiry, + OauthRefreshFailureReason: p.OauthRefreshFailureReason, + } + return link, nil + }).Times(1) + refreshErr = &oauth2.RetrieveError{ // github error Response: &http.Response{ StatusCode: http.StatusOK, }, ErrorCode: "bad_refresh_token", } - _, err := config.RefreshToken(ctx, mDB, link) + + zeroedLink, err := config.RefreshToken(ctx, mDB, link) require.Error(t, err) totalRefreshes++ require.True(t, externalauth.IsInvalidTokenError(err)) require.Equal(t, refreshCount, totalRefreshes) + // Although the fully updated link with the error is not returned, it does + // zero out the token. + require.Empty(t, zeroedLink.OAuthRefreshToken) + + // Databasae link should have a reason and no refresh token. + require.NotEmpty(t, link.OauthRefreshFailureReason) + require.Empty(t, link.OAuthRefreshToken) + + // Once more, this time with the zeroed-out refresh token due to the bad + // refresh error. + withLease(link)(mDB) // When the refresh token is empty, no api calls should be made - link.OAuthRefreshToken = "" // mock'd db, so manually set the token to '' _, err = config.RefreshToken(ctx, mDB, link) require.Error(t, err) require.True(t, externalauth.IsInvalidTokenError(err)) require.Equal(t, refreshCount, totalRefreshes) + // It should return the original cached error, not "no refresh token". + require.ErrorContains(t, err, "a long while ago") + require.ErrorContains(t, err, "bad_refresh_token") }) // RefreshTokenWithBackoff tests that refreshes which fail with transient @@ -348,11 +399,13 @@ func TestRefreshToken(t *testing.T) { cfg.RefreshRetryTimeout = 5 * time.Second }, DB: db, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) oldAccessToken := link.OAuthAccessToken - link.OAuthExpiry = expired updated, err := config.RefreshToken(ctx, db, link) require.NoError(t, err, "transient errors should be retried until success") @@ -370,8 +423,7 @@ func TestRefreshToken(t *testing.T) { t.Run("RefreshTokenBackoffPermanentError", func(t *testing.T) { t.Parallel() - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) + db, _ := dbtestutil.NewDB(t) var refreshCalls atomic.Int64 fake, config, link := setupOauth2Test(t, testConfig{ @@ -393,20 +445,14 @@ func TestRefreshToken(t *testing.T) { cfg.RefreshRetryMaxBackoff = 5 * time.Millisecond cfg.RefreshRetryTimeout = time.Second }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, + DB: db, }) - // The race-detection re-read returns the same refresh token so it - // does not look like a concurrent winner. The cached-failure write - // then proceeds. Each runs exactly once for a single refresh attempt. - mDB.EXPECT().GetExternalAuthLink(gomock.Any(), gomock.Any()). - Return(link, nil).Times(1) - mDB.EXPECT().UpdateExternalAuthLinkRefreshToken(gomock.Any(), gomock.Any()). - Return(nil).Times(1) - - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - link.OAuthExpiry = expired - - _, err := config.RefreshToken(ctx, mDB, link) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + _, err := config.RefreshToken(ctx, db, link) require.Error(t, err) require.True(t, externalauth.IsInvalidTokenError(err)) require.Equal(t, int64(1), refreshCalls.Load(), @@ -419,9 +465,6 @@ func TestRefreshToken(t *testing.T) { t.Run("ConcurrentRefreshGroup", func(t *testing.T) { t.Parallel() - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) - parallelRequests := 5 ch := make(chan string) refreshedToken := &oauth2.Token{ @@ -457,16 +500,12 @@ func TestRefreshToken(t *testing.T) { } link := database.ExternalAuthLink{OAuthExpiry: expired} - refreshedLink := database.ExternalAuthLink{ - OAuthAccessToken: refreshedToken.AccessToken, - OAuthRefreshToken: refreshedToken.RefreshToken, - OAuthExpiry: refreshedToken.Expiry, - } - // The single winning call will update the link. - mDB.EXPECT().UpdateExternalAuthLink(gomock.Any(), gomock.Cond(func(params database.UpdateExternalAuthLinkParams) bool { - return params.ProviderID == link.ProviderID && params.UserID == link.UserID - })).Return(refreshedLink, nil).Times(1) + // Only one call should try to acquire and release a lease and only one call + // should update the link. + mDB := mockDB(t, + withLease(link), + withUpdatePassthrough()) // When we fire off all requests in parallel... ctx := testutil.Context(t, testutil.WaitLong) @@ -486,7 +525,8 @@ func TestRefreshToken(t *testing.T) { // All calls should have picked up the winning token. for i := range parallelRequests { - require.Equal(t, refreshedLink, results[i]) + require.Equal(t, refreshedToken.AccessToken, results[i].OAuthAccessToken) + require.Equal(t, refreshedToken.RefreshToken, results[i].OAuthRefreshToken) } // Only one refresh call should have actually been made. @@ -495,55 +535,42 @@ func TestRefreshToken(t *testing.T) { // ConcurrentRefreshRace tests what happens a request reads the refresh token // from the database, then another request finishes and updates the token and - // releases the refresh group lock before this request can join. + // releases the refresh group lock before this request can join that group. // - // This request will then fail with `bad_refresh_token` for providers that - // have single-use refresh tokens. It should re-read the token from the - // database after making this failed request to check whether the token was - // updated by another request and returns that rather than incorrectly - // recording in the database that the request failed. + // This request would fail with `bad_refresh_token` for providers that have + // single-use refresh tokens. It should instead re-read the token from the + // database to check whether the token was updated by another request and + // returns that rather than incorrectly recording in the database that the + // request failed. t.Run("ConcurrentRefreshRace", func(t *testing.T) { t.Parallel() - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) - fake, config, link := setupOauth2Test(t, testConfig{ FakeIDPOpts: []oidctest.FakeIDPOpt{ oidctest.WithRefresh(func(_ string) error { - return &oauth2.RetrieveError{ - Response: &http.Response{ - StatusCode: http.StatusOK, - }, - ErrorCode: "bad_refresh_token", - } + return xerrors.New("should not reach this") }), }, - ExternalAuthOpt: func(cfg *externalauth.Config) {}, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - link.OAuthExpiry = time.Now().Add(time.Hour * -1) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) - // Simulate a concurrent winner: when the loser re-reads the - // DB, the refresh token has changed (the winner stored a new - // one). The loser should return the updated link instead of - // caching the failure. winnerLink := link winnerLink.OAuthRefreshToken = "winner-refresh-token" winnerLink.OAuthAccessToken = "winner-access-token" - mDB.EXPECT().GetExternalAuthLink(gomock.Any(), database.GetExternalAuthLinkParams{ - ProviderID: link.ProviderID, - UserID: link.UserID, - }).Return(winnerLink, nil).Times(1) - - // UpdateExternalAuthLinkRefreshToken should NOT be called - // because the re-read detected the concurrent refresh. + // Simulate that another caller updated the link. + // UpdateExternalAuthLinkRefreshToken should NOT be called because trying to + // get the lease detected the nearly-concurrent refresh. It should instead + // return the winning token. + mDB := mockDB(t, withLease(winnerLink)) result, err := config.RefreshToken(ctx, mDB, link) require.NoError(t, err, "loser should succeed using the winner's token") - require.Equal(t, "winner-access-token", result.OAuthAccessToken) - require.Equal(t, "winner-refresh-token", result.OAuthRefreshToken) + require.Equal(t, winnerLink.OAuthAccessToken, result.OAuthAccessToken) + require.Equal(t, winnerLink.OAuthRefreshToken, result.OAuthRefreshToken) }) // ConcurrentContextCancel tests that if one request is canceled, it does not @@ -583,8 +610,7 @@ func TestRefreshToken(t *testing.T) { } } } - // Should never reach here. - return xerrors.New("bad_refresh_token") + return xerrors.New("should not reach this") }), oidctest.WithDynamicUserInfo(func(_ string) (jwt.MapClaims, error) { return jwt.MapClaims{}, nil @@ -595,11 +621,13 @@ func TestRefreshToken(t *testing.T) { cfg.RefreshGroup = &group{notify: ch} }, DB: db, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) oldAccessToken := link.OAuthAccessToken oldRefreshToken := link.OAuthRefreshToken - link.OAuthExpiry = expired var wg sync.WaitGroup // Start the first call with the cancelable context. @@ -632,7 +660,7 @@ func TestRefreshToken(t *testing.T) { wg.Wait() // DB link should have been updated. - dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ + dbLink, err := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ ProviderID: link.ProviderID, UserID: link.UserID, }) @@ -646,15 +674,89 @@ func TestRefreshToken(t *testing.T) { require.Equal(t, int64(1), refreshCalls.Load()) }) + t.Run("LeaseAcquisitionError", func(t *testing.T) { + t.Parallel() + + fake, config, link := setupOauth2Test(t, testConfig{ + FakeIDPOpts: []oidctest.FakeIDPOpt{ + oidctest.WithRefresh(func(_ string) error { + return nil + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) {}, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, + }) + + mDB := mockDB(t, withLeaseErrors(link, xerrors.New("acquire error"), nil)) + + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + _, err := config.RefreshToken(ctx, mDB, link) + require.Error(t, err) + require.ErrorContains(t, err, "acquire error") + }) + + t.Run("ReturnsReleaseError", func(t *testing.T) { + t.Parallel() + + fake, config, link := setupOauth2Test(t, testConfig{ + FakeIDPOpts: []oidctest.FakeIDPOpt{ + oidctest.WithRefresh(func(_ string) error { + return nil + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) {}, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, + }) + + mDB := mockDB(t, + withLeaseErrors(link, nil, xerrors.New("release error")), + withUpdatePassthrough()) + + // Although the refresh was successful, an error is still returned due to + // the release having failed. + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + refreshed, err := config.RefreshToken(ctx, mDB, link) + require.Error(t, err) + require.ErrorContains(t, err, "release error") + require.NotEqual(t, link.OAuthAccessToken, refreshed.OAuthAccessToken) + require.NotEqual(t, link.OAuthRefreshToken, refreshed.OAuthRefreshToken) + }) + + t.Run("ReturnsCombinedWithReleaseError", func(t *testing.T) { + t.Parallel() + + fake, config, link := setupOauth2Test(t, testConfig{ + FakeIDPOpts: []oidctest.FakeIDPOpt{ + oidctest.WithRefresh(func(_ string) error { + return nil + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) {}, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, + }) + + mDB := mockDB(t, + withLeaseErrors(link, nil, xerrors.New("release error")), + withUpdateError(xerrors.New("update error"))) + + // Both the release and update errors should be returned. + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + _, err := config.RefreshToken(ctx, mDB, link) + require.Error(t, err) + require.ErrorContains(t, err, "release error") + require.ErrorContains(t, err, "update error") + }) + // ValidateFailure tests if the token is no longer valid with a 401 response. t.Run("ValidateFailure", func(t *testing.T) { t.Parallel() - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) - mDB.EXPECT().UpdateExternalAuthLink(gomock.Any(), gomock.Any()). - Return(database.ExternalAuthLink{}, nil).AnyTimes() - const staticError = "static error" validated := false fake, config, link := setupOauth2Test(t, testConfig{ @@ -668,9 +770,13 @@ func TestRefreshToken(t *testing.T) { }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) link.OAuthExpiry = expired + mDB := mockDB(t, + withLease(link), + withUpdatePassthrough()) + _, err := config.RefreshToken(ctx, mDB, link) require.ErrorContains(t, err, "token failed to validate") require.True(t, externalauth.IsInvalidTokenError(err)) @@ -700,13 +806,18 @@ func TestRefreshToken(t *testing.T) { ExternalAuthOpt: func(cfg *externalauth.Config) { cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + // Unlimited lifetime, this is what GitHub returns tokens as. + link.OAuthExpiry = time.Time{} + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - // Unlimited lifetime, this is what GitHub returns tokens as - link.OAuthExpiry = time.Time{} + // Since the token is not expired, no lock or refresh lease will be acquired + // and it will only be validated. + mDB := mockDB(t) - _, err := config.RefreshToken(ctx, nil, link) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + _, err := config.RefreshToken(ctx, mDB, link) require.NoError(t, err) require.Equal(t, 2, validateCalls, "token should have been attempted to be validated more than once") }) @@ -731,9 +842,12 @@ func TestRefreshToken(t *testing.T) { }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) + // Since the token is not expired, no lock or refresh lease will be acquired + // and it will only be validated. + mDB := mockDB(t) - _, err := config.RefreshToken(ctx, nil, link) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + _, err := config.RefreshToken(ctx, mDB, link) require.NoError(t, err) require.Equal(t, 1, validateCalls, "token is validated") }) @@ -760,24 +874,25 @@ func TestRefreshToken(t *testing.T) { cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() }, DB: db, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - // Force a refresh - link.OAuthExpiry = expired - + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) updated, err := config.RefreshToken(ctx, db, link) require.NoError(t, err) require.Equal(t, 1, validateCalls, "token is validated") require.Equal(t, 1, refreshCalls, "token is refreshed") require.NotEqualf(t, link.OAuthAccessToken, updated.OAuthAccessToken, "token is updated") - dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ + dbLink, err := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ ProviderID: link.ProviderID, UserID: link.UserID, }) require.NoError(t, err) require.Equal(t, updated.OAuthAccessToken, dbLink.OAuthAccessToken, "token is updated in the DB") }) + t.Run("WithExtra", func(t *testing.T) { t.Parallel() @@ -796,11 +911,12 @@ func TestRefreshToken(t *testing.T) { cfg.ValidateURL = "" }, DB: db, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - // Force a refresh - link.OAuthExpiry = expired + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) updated, err := config.RefreshToken(ctx, db, link) require.NoError(t, err) @@ -846,16 +962,16 @@ func TestRefreshToken(t *testing.T) { cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() }, DB: db, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) oldAccessToken := link.OAuthAccessToken oldRefreshToken := link.OAuthRefreshToken - // Expire the token to force a refresh. - link.OAuthExpiry = expired - // First call: refresh succeeds, validation fails (403). _, err := config.RefreshToken(ctx, db, link) require.Error(t, err, "expected error because validation returned 403") @@ -864,7 +980,7 @@ func TestRefreshToken(t *testing.T) { // Critical assertion: the DB must contain the NEW tokens from the // successful refresh, not the old (now-stale) ones. - dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ + dbLink, err := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ ProviderID: link.ProviderID, UserID: link.UserID, }) @@ -896,7 +1012,8 @@ func TestRefreshToken(t *testing.T) { db, _ := dbtestutil.NewDB(t) var refreshCalls atomic.Int64 - cancelOnRefresh, cancel := context.WithCancel(context.Background()) + ctx := testutil.Context(t, testutil.WaitLong) + cancelOnRefresh, cancel := context.WithCancel(ctx) defer cancel() fake, config, link := setupOauth2Test(t, testConfig{ @@ -916,20 +1033,20 @@ func TestRefreshToken(t *testing.T) { cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() }, DB: db, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(cancelOnRefresh, fake.HTTPClient(nil)) - oldAccessToken := link.OAuthAccessToken oldRefreshToken := link.OAuthRefreshToken - link.OAuthExpiry = expired - _, err := config.RefreshToken(ctx, db, link) + octx := oidc.ClientContext(cancelOnRefresh, fake.HTTPClient(nil)) + _, err := config.RefreshToken(octx, db, link) require.ErrorIs(t, err, context.Canceled) - require.Equal(t, int64(1), refreshCalls.Load()) require.Eventually(t, func() bool { - dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ + dbLink, err := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ ProviderID: link.ProviderID, UserID: link.UserID, }) @@ -940,6 +1057,8 @@ func TestRefreshToken(t *testing.T) { dbLink.OAuthAccessToken != oldAccessToken && dbLink.OAuthRefreshToken != oldRefreshToken }, testutil.WaitShort, testutil.IntervalFast, "never saw refresh token db updated") + + require.Equal(t, int64(1), refreshCalls.Load()) }) // SaveBeforeValidate_RateLimited tests the full path: refresh @@ -974,20 +1093,20 @@ func TestRefreshToken(t *testing.T) { cfg.ValidateURL = rateLimitValidate.URL }, DB: db, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) // Use a real HTTP transport for non-IDP requests so the // validate request can reach the httptest server. - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(&http.Client{ + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(&http.Client{ Transport: http.DefaultTransport, })) oldAccessToken := link.OAuthAccessToken oldRefreshToken := link.OAuthRefreshToken - // Expire the token to force a refresh. - link.OAuthExpiry = expired - // RefreshToken should succeed: the IDP refresh works, the // early save persists the token, and ValidateToken returns // (true, nil, nil) because the 403 has rate-limit headers. @@ -998,7 +1117,7 @@ func TestRefreshToken(t *testing.T) { "returned token should be the new one from the refresh") // Verify the DB has the new token. - dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ + dbLink, err := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ ProviderID: link.ProviderID, UserID: link.UserID, }) @@ -1017,9 +1136,6 @@ func TestRefreshToken(t *testing.T) { t.Run("SaveBeforeValidate_DBError", func(t *testing.T) { t.Parallel() - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) - fake, config, link := setupOauth2Test(t, testConfig{ FakeIDPOpts: []oidctest.FakeIDPOpt{ oidctest.WithRefresh(func(_ string) error { @@ -1029,14 +1145,16 @@ func TestRefreshToken(t *testing.T) { ExternalAuthOpt: func(cfg *externalauth.Config) { cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - link.OAuthExpiry = expired + mDB := mockDB(t, + withLease(link), + withUpdateError(xerrors.New("db connection lost"))) - mDB.EXPECT(). - UpdateExternalAuthLink(gomock.Any(), gomock.Any()). - Return(database.ExternalAuthLink{}, xerrors.New("db connection lost")) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) _, err := config.RefreshToken(ctx, mDB, link) require.Error(t, err) @@ -1045,18 +1163,243 @@ func TestRefreshToken(t *testing.T) { "DB errors should not be treated as invalid token") }) + t.Run("WaitsForConcurrentReplicaOK", func(t *testing.T) { + t.Parallel() + + fake, config, link := setupOauth2Test(t, testConfig{ + FakeIDPOpts: []oidctest.FakeIDPOpt{ + oidctest.WithRefresh(func(_ string) error { + return xerrors.New("should not be called") + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) { + cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() + // Faster polling for faster tests. + cfg.RefreshLeaseInitialBackoff = time.Millisecond + cfg.RefreshLeaseMaxBackoff = time.Millisecond + }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, + }) + + winnerLink := link + winnerLink.OAuthAccessToken = "winner-access-token" + winnerLink.OAuthRefreshToken = "winner-refresh-token" + + // Simulate another replica already having a lease. On the second attempt, + // simulate the token having been refreshed by the other replica. That + // refreshed token should be returned instead of trying to initiate a new + // refresh. + mDB := mockDB(t, + withLeaseErrors(link, &pq.Error{ + Code: pq.ErrorCode("23514"), // check_violation + Constraint: "external_auth_link_active_lease", + }, nil), + withLease(winnerLink)) + + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + + updated, err := config.RefreshToken(ctx, mDB, link) + require.NoError(t, err) + require.Equal(t, winnerLink.OAuthAccessToken, updated.OAuthAccessToken) + require.Equal(t, winnerLink.OAuthRefreshToken, updated.OAuthRefreshToken) + }) + + t.Run("WaitsForConcurrentReplicaError", func(t *testing.T) { + t.Parallel() + + fake, config, link := setupOauth2Test(t, testConfig{ + FakeIDPOpts: []oidctest.FakeIDPOpt{ + oidctest.WithRefresh(func(_ string) error { + return xerrors.New("should not be called") + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) { + cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() + // Faster polling for faster tests. + cfg.RefreshLeaseInitialBackoff = time.Millisecond + cfg.RefreshLeaseMaxBackoff = time.Millisecond + }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, + }) + + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) + + winnerLink := link + winnerLink.OAuthRefreshToken = "" + winnerLink.OauthRefreshFailureReason = "failed to refresh" + + // Simulate another replica already having a lease. On the second attempt, + // simulate the token having been failed to be refreshed by the other + // replica. That error should be returned instead of trying to initiate a + // new refresh. + mDB := mockDB(t, + withLeaseErrors(link, &pq.Error{ + Code: pq.ErrorCode("23514"), // check_violation + Constraint: "external_auth_link_active_lease", + }, nil), + withLease(winnerLink)) + + _, err := config.RefreshToken(ctx, mDB, link) + require.Error(t, err) + require.ErrorContains(t, err, winnerLink.OauthRefreshFailureReason) + }) + + t.Run("AcquireLeaseAtomicity", func(t *testing.T) { + t.Parallel() + if testing.Short() { + t.SkipNow() + } + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + one := dbtestutil.StartTx(t, db, nil) + two := dbtestutil.StartTx(t, db, nil) + + user := dbgen.User(t, db, database.User{}) + link := dbgen.ExternalAuthLink(t, db, database.ExternalAuthLink{ + UserID: user.ID, + }) + + // The winner acquires the lease inside an open transaction, holding the + // row lock without committing. + acquired, err := one.AcquireExternalAuthLinkRefreshLease(ctx, database.AcquireExternalAuthLinkRefreshLeaseParams{ + ProviderID: link.ProviderID, + UserID: link.UserID, + TimeoutMs: time.Hour.Milliseconds(), + }) + require.NoError(t, err) + require.True(t, acquired.RefreshLeaseExpiresAt.Valid) + + // The loser's UPDATE targets the same row and must block on the winner's + // uncommitted row lock rather than return anything. + loserErr := make(chan error, 1) + go func() { + _, err := two.AcquireExternalAuthLinkRefreshLease(ctx, database.AcquireExternalAuthLinkRefreshLeaseParams{ + ProviderID: link.ProviderID, + UserID: link.UserID, + TimeoutMs: time.Hour.Milliseconds(), + }) + loserErr <- err + }() + + select { + case err := <-loserErr: + t.Fatalf("loser returned %v before the winner committed; expected it to block on the row lock", err) + case <-time.After(testutil.IntervalMedium): + } + + // Commit the winner. The loser unblocks, re-checks its predicate against + // the committed row, and errors with the active lease check violation. + require.NoError(t, one.Done()) + err = testutil.RequireReceive(ctx, t, loserErr) + require.True(t, database.IsCheckViolation(err, "external_auth_link_active_lease")) + + // The stored lease is the winner's, untouched by the loser. + final, err := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ + ProviderID: link.ProviderID, + UserID: link.UserID, + }) + require.NoError(t, err) + require.Equal(t, final.RefreshLeaseExpiresAt.Valid, acquired.RefreshLeaseExpiresAt.Valid) + }) + + t.Run("LeaseMissingLink", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + + fake, config, link := setupOauth2Test(t, testConfig{ + FakeIDPOpts: []oidctest.FakeIDPOpt{ + oidctest.WithRefresh(func(_ string) error { + return nil + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) { + cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() + // Faster polling for faster tests. + cfg.RefreshLeaseInitialBackoff = time.Millisecond + cfg.RefreshLeaseMaxBackoff = time.Millisecond + }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, + }) + + // Insert another link to ensure the link acquisition function's select + // fallback matches on the right provider/user. + ctx := testutil.Context(t, testutil.WaitLong) + _, err := db.InsertExternalAuthLink(ctx, database.InsertExternalAuthLinkParams{ + ProviderID: "decoy-provider", + UserID: uuid.New(), + CreatedAt: dbtime.Now(), + UpdatedAt: dbtime.Now(), + OAuthAccessToken: "x", + OAuthRefreshToken: "x", + OAuthExpiry: dbtime.Now().Add(time.Hour), + }) + require.NoError(t, err) + + ctx = oidc.ClientContext(ctx, fake.HTTPClient(nil)) + _, err = config.RefreshToken(ctx, db, link) + require.ErrorIs(t, err, sql.ErrNoRows) + }) + + t.Run("OverridesStaleLease", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + + fake, config, link := setupOauth2Test(t, testConfig{ + DB: db, + FakeIDPOpts: []oidctest.FakeIDPOpt{ + oidctest.WithRefresh(func(_ string) error { + return nil + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) { + cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() + // Faster polling for faster tests. + cfg.RefreshLeaseInitialBackoff = time.Millisecond + cfg.RefreshLeaseMaxBackoff = time.Millisecond + }, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, + }) + + // Simulate another replica having a stale lease. + ctx := testutil.Context(t, testutil.WaitLong) + _, err := db.AcquireExternalAuthLinkRefreshLease(ctx, database.AcquireExternalAuthLinkRefreshLeaseParams{ + ProviderID: link.ProviderID, + UserID: link.UserID, + TimeoutMs: -time.Hour.Milliseconds(), + }) + require.NoError(t, err) + + ctx = oidc.ClientContext(ctx, fake.HTTPClient(nil)) + _, err = config.RefreshToken(ctx, db, link) + require.NoError(t, err) + }) + // OptimisticLockPreventsStaleOverwrite verifies that the - // UpdateExternalAuthLinkRefreshToken WHERE clause prevents a - // stale caller from overwriting a valid refresh token saved - // by a concurrent winner. + // UpdateExternalAuthLink WHERE clause prevents a stale caller from + // overwriting a valid refresh token saved by a concurrent winner. t.Run("OptimisticLockPreventsStaleOverwrite", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) + wait := make(chan struct{}) fake, config, link := setupOauth2Test(t, testConfig{ FakeIDPOpts: []oidctest.FakeIDPOpt{ oidctest.WithRefresh(func(_ string) error { + wait <- struct{}{} + <-wait return nil }), oidctest.WithDynamicUserInfo(func(_ string) (jwt.MapClaims, error) { @@ -1067,39 +1410,57 @@ func TestRefreshToken(t *testing.T) { cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() }, DB: db, + ExternalAuthLinkOpts: func(link *database.ExternalAuthLink) { + link.OAuthExpiry = expired + }, }) - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) + ctx := oidc.ClientContext(testutil.Context(t, testutil.WaitLong), fake.HTTPClient(nil)) // Snapshot the original tokens before any refresh. oldRefreshToken := link.OAuthRefreshToken - // Expire the token to force a refresh. - link.OAuthExpiry = expired + var ( + updated database.ExternalAuthLink + err error + ) - // Caller A: refresh and save successfully. - updated, err := config.RefreshToken(ctx, db, link) - require.NoError(t, err) - require.NotEqual(t, oldRefreshToken, updated.OAuthRefreshToken, - "caller A should have a new refresh token") - - // Caller B had a stale read of the original link. It tries to - // destroy the refresh token using the OLD refresh token in the - // optimistic lock. Because caller A already wrote a different - // refresh token, this WHERE clause matches nothing. - err = db.UpdateExternalAuthLinkRefreshToken(ctx, database.UpdateExternalAuthLinkRefreshTokenParams{ + // Caller A begins a refresh and takes the lease. + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + updated, err = config.RefreshToken(ctx, db, link) + assert.NoError(t, err) + assert.NotEqual(t, oldRefreshToken, updated.OAuthRefreshToken, + "caller A should have a new refresh token") + }() + + // Once caller A has the lease, simulate a caller B trying to update the + // link with an error. It should fail because caller A has the lease. + <-wait + _, err = db.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ + ProviderID: link.ProviderID, + UserID: link.UserID, + // Write the error an 8lear the token. OauthRefreshFailureReason: "simulated failure from stale caller B", OAuthRefreshToken: "", - OAuthRefreshTokenKeyID: "", UpdatedAt: dbtime.Now(), - ProviderID: link.ProviderID, - UserID: link.UserID, - OldOauthRefreshToken: oldRefreshToken, + // This should prevent the write because it does not match. + RefreshLeaseExpiresAt: sql.NullTime{Time: dbtime.Now().Add(time.Hour), Valid: true}, + // Preserve then. + OAuthAccessToken: link.OAuthAccessToken, + OAuthExpiry: link.OAuthExpiry, + OAuthExtra: link.OAuthExtra, }) - require.NoError(t, err, "optimistic lock write should not error, it is a no-op") + require.ErrorIs(t, err, sql.ErrNoRows) + + // Let caller A finish. + close(wait) + wg.Wait() // Verify DB still has caller A's valid token. - dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ + dbLink, err := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ ProviderID: link.ProviderID, UserID: link.UserID, }) @@ -1160,38 +1521,22 @@ func TestRefreshTokenWithScopes(t *testing.T) { expired := dbtime.Now().Add(-time.Hour) - // mockDBPassthrough returns a mock store that echoes the - // UpdateExternalAuthLink params back as a populated ExternalAuthLink, - // letting the test read what RefreshToken decided to persist. - mockDBPassthrough := func(t *testing.T) database.Store { - t.Helper() - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) - mDB.EXPECT().UpdateExternalAuthLink(gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, p database.UpdateExternalAuthLinkParams) (database.ExternalAuthLink, error) { - return database.ExternalAuthLink{ - ProviderID: p.ProviderID, - UserID: p.UserID, - OAuthAccessToken: p.OAuthAccessToken, - OAuthRefreshToken: p.OAuthRefreshToken, - OAuthExpiry: p.OAuthExpiry, - }, nil - }).AnyTimes() - return mDB - } - t.Run("EchoesConfiguredScopesOnRefresh", func(t *testing.T) { t.Parallel() client, captured := fakeAS(t, []byte(`{"access_token":"new","refresh_token":"new-r","token_type":"bearer","expires_in":3600}`)) cfg := newConfig(t, []string{"openid", "offline_access", "api://app/session:role-any"}) - ctx := context.WithValue(context.Background(), oauth2.HTTPClient, client) - _, err := cfg.RefreshToken(ctx, mockDBPassthrough(t), database.ExternalAuthLink{ + ctx := context.WithValue(testutil.Context(t, testutil.WaitLong), oauth2.HTTPClient, client) + link := database.ExternalAuthLink{ OAuthAccessToken: "old", OAuthRefreshToken: "old-r", OAuthExpiry: expired, - }) + } + mDB := mockDB(t, + withLease(link), + withUpdatePassthrough()) + _, err := cfg.RefreshToken(ctx, mDB, link) require.NoError(t, err) require.Equal(t, "refresh_token", captured.Get("grant_type")) @@ -1206,12 +1551,16 @@ func TestRefreshTokenWithScopes(t *testing.T) { []byte(`{"access_token":"new","refresh_token":"new-r","token_type":"bearer","expires_in":3600}`)) cfg := newConfig(t, nil) - ctx := context.WithValue(context.Background(), oauth2.HTTPClient, client) - _, err := cfg.RefreshToken(ctx, mockDBPassthrough(t), database.ExternalAuthLink{ + ctx := context.WithValue(testutil.Context(t, testutil.WaitLong), oauth2.HTTPClient, client) + link := database.ExternalAuthLink{ OAuthAccessToken: "old", OAuthRefreshToken: "old-r", OAuthExpiry: expired, - }) + } + mDB := mockDB(t, + withLease(link), + withUpdatePassthrough()) + _, err := cfg.RefreshToken(ctx, mDB, link) require.NoError(t, err) require.Equal(t, "refresh_token", captured.Get("grant_type")) @@ -1227,12 +1576,16 @@ func TestRefreshTokenWithScopes(t *testing.T) { []byte(`{"access_token":"new","token_type":"bearer","expires_in":3600}`)) cfg := newConfig(t, nil) - ctx := context.WithValue(context.Background(), oauth2.HTTPClient, client) - link, err := cfg.RefreshToken(ctx, mockDBPassthrough(t), database.ExternalAuthLink{ + ctx := context.WithValue(testutil.Context(t, testutil.WaitLong), oauth2.HTTPClient, client) + link := database.ExternalAuthLink{ OAuthAccessToken: "old", OAuthRefreshToken: "prior-r", OAuthExpiry: expired, - }) + } + mDB := mockDB(t, + withLease(link), + withUpdatePassthrough()) + link, err := cfg.RefreshToken(ctx, mDB, link) require.NoError(t, err) require.Equal(t, "prior-r", link.OAuthRefreshToken, "prior refresh_token must be preserved when AS omits a new one (RFC 6749 §6)") @@ -1244,12 +1597,16 @@ func TestRefreshTokenWithScopes(t *testing.T) { []byte(`{"access_token":"new","refresh_token":"rotated-r","token_type":"bearer","expires_in":3600}`)) cfg := newConfig(t, nil) - ctx := context.WithValue(context.Background(), oauth2.HTTPClient, client) - link, err := cfg.RefreshToken(ctx, mockDBPassthrough(t), database.ExternalAuthLink{ + ctx := context.WithValue(testutil.Context(t, testutil.WaitLong), oauth2.HTTPClient, client) + link := database.ExternalAuthLink{ OAuthAccessToken: "old", OAuthRefreshToken: "prior-r", OAuthExpiry: expired, - }) + } + mDB := mockDB(t, + withLease(link), + withUpdatePassthrough()) + link, err := cfg.RefreshToken(ctx, mDB, link) require.NoError(t, err) require.Equal(t, "rotated-r", link.OAuthRefreshToken, "rotated refresh_token from AS must be persisted") @@ -1341,7 +1698,7 @@ func TestValidateToken(t *testing.T) { t.Helper() tp := &http.Transport{} t.Cleanup(tp.CloseIdleConnections) - return oidc.ClientContext(context.Background(), &http.Client{Transport: tp}) + return oidc.ClientContext(testutil.Context(t, testutil.WaitLong), &http.Client{Transport: tp}) } // RateLimitRemaining: 403 with X-RateLimit-Remaining: 0 should be @@ -1703,7 +2060,8 @@ func TestExchangeWithClientSecret(t *testing.T) { }), } - _, err = config.Exchange(context.WithValue(context.Background(), oauth2.HTTPClient, client), "code") + ctx := testutil.Context(t, testutil.WaitLong) + _, err = config.Exchange(context.WithValue(ctx, oauth2.HTTPClient, client), "code") require.NoError(t, err) } @@ -1941,6 +2299,9 @@ type testConfig struct { ExternalAuthOpt func(cfg *externalauth.Config) // If DB is passed in, the link will be inserted into the DB. DB database.Store + // ExternalAuthLinkOpts can be used to manipulate the link inserted into the + // DB when the DB is provided. + ExternalAuthLinkOpts func(link *database.ExternalAuthLink) } // setupTest will configure a fake IDP and a externalauth.Config for testing. @@ -1992,6 +2353,9 @@ func setupOauth2Test(t *testing.T, settings testConfig) (*oidctest.FakeIDP, *ext // The caller can manually expire this if they want. OAuthExpiry: now.Add(time.Hour), } + if settings.ExternalAuthLinkOpts != nil { + settings.ExternalAuthLinkOpts(&link) + } if settings.DB != nil { // Feel free to insert additional things like the user, etc if required. @@ -2010,6 +2374,84 @@ func setupOauth2Test(t *testing.T, settings testConfig) (*oidctest.FakeIDP, *ext return fake, config, link } +type mockDBOption func(*dbmock.MockStore) + +// withLease wraps withLeaseErrors with nil errors. +func withLease(link database.ExternalAuthLink) mockDBOption { + return withLeaseErrors(link, nil, nil) +} + +// withLeaseErrors expects that a lease is acquired and that same lease is then +// released. If acquireErr is non-nil then a release is not expected and +// releaseErr goes unused. +func withLeaseErrors(link database.ExternalAuthLink, acquireErr, releaseErr error) mockDBOption { + var lease sql.NullTime + return func(mDB *dbmock.MockStore) { + mDB.EXPECT().AcquireExternalAuthLinkRefreshLease( + gomock.Any(), + gomock.Cond(func(params database.AcquireExternalAuthLinkRefreshLeaseParams) bool { + return params.ProviderID == link.ProviderID && + params.UserID == link.UserID && + params.TimeoutMs > 0 + })).DoAndReturn(func(_ context.Context, params database.AcquireExternalAuthLinkRefreshLeaseParams) (database.ExternalAuthLink, error) { + if acquireErr == nil { + // Return the same link but now with a lease attached. + lease = sql.NullTime{Valid: true, Time: dbtime.Now().Add(time.Duration(params.TimeoutMs) * time.Millisecond)} + leasedLink := link + leasedLink.RefreshLeaseExpiresAt = lease + return leasedLink, nil + } + return database.ExternalAuthLink{}, acquireErr + }).Times(1) + if acquireErr == nil { + mDB.EXPECT().ReleaseExternalAuthLinkRefreshLease(gomock.Any(), gomock.Cond(func(params database.ReleaseExternalAuthLinkRefreshLeaseParams) bool { + return params.ProviderID == link.ProviderID && + params.UserID == link.UserID && + params.RefreshLeaseExpiresAt.Valid && + // Should have passed the same lease back in. + params.RefreshLeaseExpiresAt.Time.Equal(lease.Time) + })).Return(releaseErr).Times(1) + } + } +} + +// withUpdateError expects an update to be attempted and returns an error for +// that update. +func withUpdateError(err error) mockDBOption { + return func(mDB *dbmock.MockStore) { + mDB.EXPECT().UpdateExternalAuthLink(gomock.Any(), gomock.Any()). + Return(database.ExternalAuthLink{}, err).Times(1) + } +} + +// withUpdatePassthrough echoes the UpdateExternalAuthLink params back as a +// populated ExternalAuthLink, letting the test read what RefreshToken decided +// to persist, if anything. +func withUpdatePassthrough() mockDBOption { + return func(mDB *dbmock.MockStore) { + mDB.EXPECT().UpdateExternalAuthLink(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, p database.UpdateExternalAuthLinkParams) (database.ExternalAuthLink, error) { + return database.ExternalAuthLink{ + ProviderID: p.ProviderID, + UserID: p.UserID, + OAuthAccessToken: p.OAuthAccessToken, + OAuthRefreshToken: p.OAuthRefreshToken, + OAuthExpiry: p.OAuthExpiry, + }, nil + }).Times(1) + } +} + +func mockDB(t *testing.T, opts ...mockDBOption) *dbmock.MockStore { + t.Helper() + ctrl := gomock.NewController(t) + mDB := dbmock.NewMockStore(ctrl) + for _, opt := range opts { + opt(mDB) + } + return mDB +} + func TestApplyDefaultsToConfig_CaseInsensitive(t *testing.T) { t.Parallel() diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go index c5eb376c8d313..fb28ad2ecd135 100644 --- a/coderd/provisionerdserver/provisionerdserver_test.go +++ b/coderd/provisionerdserver/provisionerdserver_test.go @@ -522,7 +522,7 @@ func TestAcquireJob(t *testing.T) { OAuthExpiry: dbtime.Now().Add(time.Hour), OAuthAccessToken: "access-token", }) - dbgen.ExternalAuthLink(t, db, database.ExternalAuthLink{ + ealink := dbgen.ExternalAuthLink(t, db, database.ExternalAuthLink{ ProviderID: gitAuthProvider.Id, UserID: user.ID, }) @@ -768,7 +768,7 @@ func TestAcquireJob(t *testing.T) { }, ExternalAuthProviders: []*sdkproto.ExternalAuthProvider{{ Id: gitAuthProvider.Id, - AccessToken: "access_token", + AccessToken: ealink.OAuthAccessToken, }}, Metadata: wantedMetadata, }, diff --git a/enterprise/dbcrypt/cliutil.go b/enterprise/dbcrypt/cliutil.go index 086c9ba6b4f97..a6656356ccc6c 100644 --- a/enterprise/dbcrypt/cliutil.go +++ b/enterprise/dbcrypt/cliutil.go @@ -60,15 +60,17 @@ func Rotate(ctx context.Context, log slog.Logger, sqlDB *sql.DB, ciphers []Ciphe continue } if _, err := cryptTx.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ - ProviderID: externalAuthLink.ProviderID, - UserID: uid, - UpdatedAt: externalAuthLink.UpdatedAt, - OAuthAccessToken: externalAuthLink.OAuthAccessToken, - OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required - OAuthRefreshToken: externalAuthLink.OAuthRefreshToken, - OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required - OAuthExpiry: externalAuthLink.OAuthExpiry, - OAuthExtra: externalAuthLink.OAuthExtra, + ProviderID: externalAuthLink.ProviderID, + UserID: uid, + UpdatedAt: externalAuthLink.UpdatedAt, + OAuthAccessToken: externalAuthLink.OAuthAccessToken, + OAuthAccessTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthRefreshToken: externalAuthLink.OAuthRefreshToken, + OAuthRefreshTokenKeyID: sql.NullString{}, // dbcrypt will update as required + OAuthExpiry: externalAuthLink.OAuthExpiry, + OAuthExtra: externalAuthLink.OAuthExtra, + OauthRefreshFailureReason: "", + RefreshLeaseExpiresAt: sql.NullTime{}, }); err != nil { return xerrors.Errorf("update external auth link user_id=%s provider_id=%s: %w", externalAuthLink.UserID, externalAuthLink.ProviderID, err) } @@ -274,15 +276,17 @@ func Decrypt(ctx context.Context, log slog.Logger, sqlDB *sql.DB, ciphers []Ciph continue } if _, err := tx.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ - ProviderID: externalAuthLink.ProviderID, - UserID: uid, - UpdatedAt: externalAuthLink.UpdatedAt, - OAuthAccessToken: externalAuthLink.OAuthAccessToken, - OAuthAccessTokenKeyID: sql.NullString{}, // we explicitly want to clear the key id - OAuthRefreshToken: externalAuthLink.OAuthRefreshToken, - OAuthRefreshTokenKeyID: sql.NullString{}, // we explicitly want to clear the key id - OAuthExpiry: externalAuthLink.OAuthExpiry, - OAuthExtra: externalAuthLink.OAuthExtra, + ProviderID: externalAuthLink.ProviderID, + UserID: uid, + UpdatedAt: externalAuthLink.UpdatedAt, + OAuthAccessToken: externalAuthLink.OAuthAccessToken, + OAuthAccessTokenKeyID: sql.NullString{}, // we explicitly want to clear the key id + OAuthRefreshToken: externalAuthLink.OAuthRefreshToken, + OAuthRefreshTokenKeyID: sql.NullString{}, // we explicitly want to clear the key id + OAuthExpiry: externalAuthLink.OAuthExpiry, + OAuthExtra: externalAuthLink.OAuthExtra, + OauthRefreshFailureReason: "", + RefreshLeaseExpiresAt: sql.NullTime{}, }); err != nil { return xerrors.Errorf("update external auth link user_id=%s provider_id=%s: %w", externalAuthLink.UserID, externalAuthLink.ProviderID, err) } diff --git a/enterprise/dbcrypt/dbcrypt.go b/enterprise/dbcrypt/dbcrypt.go index 72b9003b2e98a..44b109d83cadd 100644 --- a/enterprise/dbcrypt/dbcrypt.go +++ b/enterprise/dbcrypt/dbcrypt.go @@ -243,6 +243,20 @@ func (db *dbCrypt) GetExternalAuthLinksByUserID(ctx context.Context, userID uuid return links, nil } +func (db *dbCrypt) AcquireExternalAuthLinkRefreshLease(ctx context.Context, params database.AcquireExternalAuthLinkRefreshLeaseParams) (database.ExternalAuthLink, error) { + link, err := db.Store.AcquireExternalAuthLinkRefreshLease(ctx, params) + if err != nil { + return database.ExternalAuthLink{}, err + } + if err := db.decryptField(&link.OAuthAccessToken, link.OAuthAccessTokenKeyID); err != nil { + return database.ExternalAuthLink{}, err + } + if err := db.decryptField(&link.OAuthRefreshToken, link.OAuthRefreshTokenKeyID); err != nil { + return database.ExternalAuthLink{}, err + } + return link, nil +} + func (db *dbCrypt) UpdateExternalAuthLink(ctx context.Context, params database.UpdateExternalAuthLinkParams) (database.ExternalAuthLink, error) { if err := db.encryptField(¶ms.OAuthAccessToken, ¶ms.OAuthAccessTokenKeyID); err != nil { return database.ExternalAuthLink{}, err @@ -263,54 +277,6 @@ func (db *dbCrypt) UpdateExternalAuthLink(ctx context.Context, params database.U return link, nil } -func (db *dbCrypt) UpdateExternalAuthLinkRefreshToken(ctx context.Context, params database.UpdateExternalAuthLinkRefreshTokenParams) error { - // The SQL query uses an optimistic lock: - // WHERE oauth_refresh_token = @old_oauth_refresh_token - // The caller supplies the plaintext old token (since dbcrypt - // decrypts on read), but the DB stores the encrypted value. - // Because AES-GCM is non-deterministic, we cannot simply - // re-encrypt the old token — the ciphertext would differ. - // Instead, read the current row from the inner (raw) store - // and use the actual encrypted value for the WHERE clause. - if params.OldOauthRefreshToken != "" && db.ciphers != nil && db.primaryCipherDigest != "" { - raw, err := db.Store.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ - ProviderID: params.ProviderID, - UserID: params.UserID, - }) - if err != nil { - return err - } - // Decrypt the stored token so we can compare with the - // caller-supplied plaintext. - decrypted := raw.OAuthRefreshToken - if err := db.decryptField(&decrypted, raw.OAuthRefreshTokenKeyID); err != nil { - return err - } - if decrypted != params.OldOauthRefreshToken { - // The token has changed since the caller read it; - // the optimistic lock should fail (no rows updated). - // Return nil to match the :exec semantics of the SQL - // query, which silently updates zero rows. - return nil - } - // Use the raw encrypted value so the WHERE clause matches. - params.OldOauthRefreshToken = raw.OAuthRefreshToken - } - - // We would normally use a sql.NullString here, but sqlc does not want to make - // a params struct with a nullable string. - var digest sql.NullString - if params.OAuthRefreshTokenKeyID != "" { - digest.String = params.OAuthRefreshTokenKeyID - digest.Valid = true - } - if err := db.encryptField(¶ms.OAuthRefreshToken, &digest); err != nil { - return err - } - - return db.Store.UpdateExternalAuthLinkRefreshToken(ctx, params) -} - func (db *dbCrypt) GetCryptoKeys(ctx context.Context) ([]database.CryptoKey, error) { keys, err := db.Store.GetCryptoKeys(ctx) if err != nil { diff --git a/enterprise/dbcrypt/dbcrypt_internal_test.go b/enterprise/dbcrypt/dbcrypt_internal_test.go index 996a384ab6dd6..47fe8d943919a 100644 --- a/enterprise/dbcrypt/dbcrypt_internal_test.go +++ b/enterprise/dbcrypt/dbcrypt_internal_test.go @@ -100,32 +100,6 @@ func TestUserLinks(t *testing.T) { require.EqualValues(t, expectedClaims, rawLink.Claims) }) - t.Run("UpdateExternalAuthLinkRefreshToken", func(t *testing.T) { - t.Parallel() - db, crypt, ciphers := setup(t) - user := dbgen.User(t, crypt, database.User{}) - link := dbgen.ExternalAuthLink(t, crypt, database.ExternalAuthLink{ - UserID: user.ID, - }) - - err := crypt.UpdateExternalAuthLinkRefreshToken(ctx, database.UpdateExternalAuthLinkRefreshTokenParams{ - OAuthRefreshToken: "", - OAuthRefreshTokenKeyID: link.OAuthRefreshTokenKeyID.String, - OldOauthRefreshToken: link.OAuthRefreshToken, - UpdatedAt: dbtime.Now(), - ProviderID: link.ProviderID, - UserID: link.UserID, - }) - require.NoError(t, err) - - rawLink, err := db.GetExternalAuthLink(ctx, database.GetExternalAuthLinkParams{ - ProviderID: link.ProviderID, - UserID: link.UserID, - }) - require.NoError(t, err) - requireEncryptedEquals(t, ciphers[0], rawLink.OAuthRefreshToken, "") - }) - t.Run("GetUserLinkByLinkedID", func(t *testing.T) { t.Parallel() t.Run("OK", func(t *testing.T) { @@ -360,6 +334,61 @@ func TestExternalAuthLinks(t *testing.T) { }) }) + t.Run("AcquireExternalAuthLinkRefreshLease", func(t *testing.T) { + t.Run("OK", func(t *testing.T) { + t.Parallel() + db, crypt, ciphers := setup(t) + link := dbgen.ExternalAuthLink(t, crypt, database.ExternalAuthLink{ + OAuthAccessToken: "access", + OAuthRefreshToken: "refresh", + }) + link, err := db.AcquireExternalAuthLinkRefreshLease(ctx, database.AcquireExternalAuthLinkRefreshLeaseParams{ + UserID: link.UserID, + ProviderID: link.ProviderID, + TimeoutMs: 10, + }) + require.NoError(t, err) + requireEncryptedEquals(t, ciphers[0], link.OAuthAccessToken, "access") + requireEncryptedEquals(t, ciphers[0], link.OAuthRefreshToken, "refresh") + }) + t.Run("Decrypt", func(t *testing.T) { + t.Parallel() + _, crypt, ciphers := setup(t) + link := dbgen.ExternalAuthLink(t, crypt, database.ExternalAuthLink{ + OAuthAccessToken: "access", + OAuthRefreshToken: "refresh", + }) + link, err := crypt.AcquireExternalAuthLinkRefreshLease(ctx, database.AcquireExternalAuthLinkRefreshLeaseParams{ + UserID: link.UserID, + ProviderID: link.ProviderID, + TimeoutMs: 10, + }) + require.NoError(t, err) + require.Equal(t, "access", link.OAuthAccessToken) + require.Equal(t, "refresh", link.OAuthRefreshToken) + require.Equal(t, ciphers[0].HexDigest(), link.OAuthAccessTokenKeyID.String) + require.Equal(t, ciphers[0].HexDigest(), link.OAuthRefreshTokenKeyID.String) + }) + t.Run("DecryptErr", func(t *testing.T) { + t.Parallel() + db, crypt, ciphers := setup(t) + link := dbgen.ExternalAuthLink(t, db, database.ExternalAuthLink{ + OAuthAccessToken: fakeBase64RandomData(t, 32), + OAuthRefreshToken: fakeBase64RandomData(t, 32), + OAuthAccessTokenKeyID: sql.NullString{String: ciphers[0].HexDigest(), Valid: true}, + OAuthRefreshTokenKeyID: sql.NullString{String: ciphers[0].HexDigest(), Valid: true}, + }) + link, err := crypt.AcquireExternalAuthLinkRefreshLease(ctx, database.AcquireExternalAuthLinkRefreshLeaseParams{ + UserID: link.UserID, + ProviderID: link.ProviderID, + TimeoutMs: 10, + }) + require.Error(t, err, "expected an error") + var derr *DecryptFailedError + require.ErrorAs(t, err, &derr, "expected a decrypt error") + }) + }) + t.Run("GetExternalAuthLinksByUserID", func(t *testing.T) { t.Parallel()