From 7ba2e6e22af3e71331d4c135195c6f9518c95e17 Mon Sep 17 00:00:00 2001 From: Asher Date: Wed, 19 Aug 2026 14:00:31 -0700 Subject: [PATCH] fix: add distributed protection for external auth refresh (#27479) Before refreshing, the caller has to put a lease on the row to prevent any other replicas from trying to refresh using the same token. Callers unable to get a lease enter a holding pattern waiting for the results. Also use this same lease query to re-read the link. This more or less functions the same as the re-read we used to have after a refresh failure, except it also allows us to avoid making the failed request in the first place. The downside is that we have to refetch the link every time even if the caller just fetched it, but since we have to make a database call to set the lease anyway... --- coderd/database/dbauthz/dbauthz.go | 21 +- coderd/database/dbauthz/dbauthz_test.go | 24 +- coderd/database/dbmetrics/querymetrics.go | 24 +- coderd/database/dbmock/dbmock.go | 43 +- coderd/database/dump.sql | 74 +- .../000577_external_auth_lock.down.sql | 2 + .../000577_external_auth_lock.up.sql | 32 + coderd/database/models.go | 2 + coderd/database/querier.go | 11 +- coderd/database/queries.sql.go | 160 ++-- coderd/database/queries/externalauth.sql | 55 +- coderd/externalauth.go | 40 +- coderd/externalauth/externalauth.go | 358 +++++--- coderd/externalauth/externalauth_test.go | 842 +++++++++++++----- .../provisionerdserver_test.go | 4 +- enterprise/dbcrypt/cliutil.go | 40 +- enterprise/dbcrypt/dbcrypt.go | 62 +- enterprise/dbcrypt/dbcrypt_internal_test.go | 81 +- 18 files changed, 1272 insertions(+), 603 deletions(-) create mode 100644 coderd/database/migrations/000577_external_auth_lock.down.sql create mode 100644 coderd/database/migrations/000577_external_auth_lock.up.sql diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 77ab7dc5ae5..632efc4c8e1 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 fa784ded25b..738b4a2de99 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 ae63a286768..492686145ad 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 6864efcb86e..0c08caa669e 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 e32e7513112..343491f73f0 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 00000000000..82f935fabd6 --- /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 00000000000..56ee05ae576 --- /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 0e914ec66f5..68130280eb1 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 58c197125b5..ac1945f20f0 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 58492959776..e364a9400ea 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 e5d0ec548bf..aeaf714420f 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: 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: 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: 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 51b7727c00d..04ab62d3af6 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 c6d87e34b4f..e4e9809fe2f 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. +// 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, + )) +} + +// 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 { - refreshToken = "" - } - - existingToken := &oauth2.Token{ - AccessToken: externalAuthLink.OAuthAccessToken, - RefreshToken: refreshToken, - Expiry: externalAuthLink.OAuthExpiry, + 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 - } - - 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) + // 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 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, xerrors.Errorf("persist refreshed token: %w", 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 { + user, err := c.validateWithRetry(ctx, token) + if err != nil { + 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 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 9e954219d82..1001668dce3 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,11 +1163,152 @@ func TestRefreshToken(t *testing.T) { "DB errors should not be treated as invalid token") }) - // OptimisticLockPreventsStaleOverwrite verifies that the - // UpdateExternalAuthLinkRefreshToken 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.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) @@ -1059,6 +1318,90 @@ func TestRefreshToken(t *testing.T) { 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 + // 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) { return jwt.MapClaims{}, nil }), @@ -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 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") + }() - // 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{ + // 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 c5eb376c8d3..fb28ad2ecd1 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 086c9ba6b4f..a6656356ccc 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 72b9003b2e9..44b109d83ca 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 996a384ab6d..47fe8d94391 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()