From 36997a7d8aafb6a2475811dbdfeea985da2ddf13 Mon Sep 17 00:00:00 2001 From: Asher Date: Fri, 26 Jun 2026 13:39:09 -0800 Subject: [PATCH 1/9] fix: prevent concurrent external auth token refreshes This can cause bad refresh token errors, since it can only be used once. Only covers a single instance. --- coderd/externalauth/externalauth.go | 32 ++++-- coderd/externalauth/externalauth_test.go | 104 +++++++++++------- coderd/promoauth/oauth2.go | 21 +++- coderd/util/singleflight/singleflight.go | 53 +++++++++ coderd/util/singleflight/singleflight_test.go | 85 ++++++++++++++ testutil/oauth2.go | 12 ++ 6 files changed, 255 insertions(+), 52 deletions(-) create mode 100644 coderd/util/singleflight/singleflight.go create mode 100644 coderd/util/singleflight/singleflight_test.go diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go index 6633ab1936e6c..f75b99f55139e 100644 --- a/coderd/externalauth/externalauth.go +++ b/coderd/externalauth/externalauth.go @@ -190,6 +190,22 @@ func IsInvalidTokenError(err error) bool { // RefreshToken automatically refreshes the token if expired and permitted. func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAuthLink database.ExternalAuthLink) (database.ExternalAuthLink, error) { + // 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. + key := externalAuthLink.OAuthAccessToken + link, err := c.Group().Do(key, func() (any, error) { + return c.innerRefreshToken(ctx, db, externalAuthLink) + }) + if newlink, ok := link.(database.ExternalAuthLink); ok { + return newlink, err + } else if err == nil { + err = xerrors.Errorf("got invalid type from token refresh: %T", link) + } + return externalAuthLink, err +} + +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 && @@ -237,21 +253,17 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu // // The error message is saved for debugging purposes. if isFailedRefresh(existingToken, err) { - // Before caching the failure, re-read the external auth link - // from the database. A 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. + // 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 { - // Another caller won the refresh race and stored a new - // refresh token. Return their updated link instead of - // caching a failure. return currentLink, nil } diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go index 9951e06dc96e1..4b32b1d7cf6aa 100644 --- a/coderd/externalauth/externalauth_test.go +++ b/coderd/externalauth/externalauth_test.go @@ -21,6 +21,7 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/oauth2" + "golang.org/x/sync/errgroup" "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd" @@ -336,55 +337,82 @@ func TestRefreshToken(t *testing.T) { "permanent failures should not be retried") }) - // ConcurrentRefreshRace tests that when multiple concurrent requests - // race to refresh the same token, the loser does not poison the - // database with a cached "bad_refresh_token" failure. This - // reproduces the issue described in coder/coder#17069 where - // providers with single-use refresh tokens (e.g., GitHub Apps) - // reject the second refresh attempt, and the resulting error was - // incorrectly cached. + // ConcurrentRefreshRace tests that when multiple concurrent requests race to + // refresh the same token, they share the same request instead of only one + // succeeding and the others using a now-invalid refresh token. 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", + parallelRequests := 5 + ch := make(chan string) + refreshedToken := &oauth2.Token{ + AccessToken: "winner-access-token", + RefreshToken: "winner-refresh-token", + Expiry: time.Now().Add(time.Hour), + } + + var refreshCalls atomic.Int64 + config := &externalauth.Config{ + InstrumentedOAuth2Config: &testutil.OAuth2Config{ + // The first call to refresh well succeed and all others will fail. The + // first will wait for all callers to join the group before returning. + TokenSourceFunc: func() (*oauth2.Token, error) { + if refreshCalls.Add(1) == 1 { + // Wait for all the other calls to be subscribed, to prevent + // the test from flaking. + subscribed := 1 + for { + <-ch + subscribed++ + if subscribed >= parallelRequests { + return refreshedToken, nil + } + } } - }), + return nil, xerrors.New("bad_refresh_token") + }, + Notifier: ch, }, - ExternalAuthOpt: func(cfg *externalauth.Config) {}, - }) + } - ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) - link.OAuthExpiry = time.Now().Add(time.Hour * -1) - - // 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) + link := database.ExternalAuthLink{OAuthExpiry: expired} + refreshedLink := database.ExternalAuthLink{ + OAuthAccessToken: refreshedToken.AccessToken, + OAuthRefreshToken: refreshedToken.RefreshToken, + OAuthExpiry: refreshedToken.Expiry, + } - // UpdateExternalAuthLinkRefreshToken should NOT be called - // because the re-read detected the concurrent refresh. + // 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) + + // When we fire off all requests in parallel... + ctx := testutil.Context(t, testutil.WaitLong) + var eg errgroup.Group + results := make([]database.ExternalAuthLink, parallelRequests) + for i := range parallelRequests { + eg.Go(func() error { + result, err := config.RefreshToken(ctx, mDB, link) + results[i] = result + return err + }) + } + + // No call should error. + err := eg.Wait() + require.NoError(t, err) - 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) + // All calls should have picked up the winning token. + for i := range parallelRequests { + require.Equal(t, refreshedLink, results[i]) + } + + // Only one refresh call should have actually been made. + require.Equal(t, int64(1), refreshCalls.Load()) }) // ValidateFailure tests if the token is no longer valid with a 401 response. diff --git a/coderd/promoauth/oauth2.go b/coderd/promoauth/oauth2.go index 91b34dbd95019..7d42779aa37e0 100644 --- a/coderd/promoauth/oauth2.go +++ b/coderd/promoauth/oauth2.go @@ -9,6 +9,8 @@ import ( "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promauto" "golang.org/x/oauth2" + + "github.com/coder/coder/v2/coderd/util/singleflight" ) type Oauth2PKCEChallengeMethod string @@ -42,16 +44,21 @@ type OAuth2Config interface { TokenSource(context.Context, *oauth2.Token) oauth2.TokenSource } -// InstrumentedOAuth2Config extends OAuth2Config with a `Do` method that allows -// external oauth related calls to be instrumented. This is to support -// "ValidateToken" which is not an oauth2 specified method. -// These calls still count against the api rate limit, and should be instrumented. +// InstrumentedOAuth2Config extends OAuth2Config with: +// - A `Do` method that allows external oauth related calls to be +// instrumented. This is to support "ValidateToken" which is not an oauth2 +// specified method. These calls still count against the api rate limit, and +// should be instrumented. +// - A `Group` method that allows calls to be deduplicated. type InstrumentedOAuth2Config interface { OAuth2Config // Do is provided as a convenience method to make a request with the oauth2 client. // It mirrors `http.Client.Do`. Do(ctx context.Context, source Oauth2Source, req *http.Request) (*http.Response, error) + + // Group returns a singleflight group for deduplicating concurrent requests. + Group() *singleflight.Group } var _ OAuth2Config = (*Config)(nil) @@ -152,6 +159,7 @@ func (f *Factory) New(name string, under OAuth2Config) *Config { name: name, underlying: under, metrics: f.metrics, + group: singleflight.NewGroup(nil), } } @@ -202,6 +210,7 @@ type Config struct { metrics *metrics // interceptors are called after every request made by the oauth2 client. interceptors []func(resp *http.Response, err error) + group *singleflight.Group } func (c *Config) Do(ctx context.Context, source Oauth2Source, req *http.Request) (*http.Response, error) { @@ -209,6 +218,10 @@ func (c *Config) Do(ctx context.Context, source Oauth2Source, req *http.Request) return cli.Do(req) } +func (c *Config) Group() *singleflight.Group { + return c.group +} + func (c *Config) AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string { // No external requests are made when constructing the auth code url. return c.underlying.AuthCodeURL(state, opts...) diff --git a/coderd/util/singleflight/singleflight.go b/coderd/util/singleflight/singleflight.go new file mode 100644 index 0000000000000..73dd2364933ac --- /dev/null +++ b/coderd/util/singleflight/singleflight.go @@ -0,0 +1,53 @@ +package singleflight + +import ( + "sync" +) + +type call struct { + wg sync.WaitGroup + val any + err error +} + +// Group protects concurrent calls. +type Group struct { + mu sync.Mutex // protects m + m map[string]*call // lazily initialized + notify chan string +} + +func NewGroup(notifier chan string) *Group { + group := new(Group) + group.notify = notifier + return group +} + +// Do ensures there is only one call to fn in flight at a time. Any calls that +// come in while it is in flight wait for the original call and get the same +// results. +func (g *Group) Do(key string, fn func() (any, error)) (v any, err error) { + g.mu.Lock() + if g.m == nil { + g.m = make(map[string]*call) + } + if c, ok := g.m[key]; ok { + if g.notify != nil { + g.notify <- key + } + g.mu.Unlock() + c.wg.Wait() + return c.val, c.err + } + c := new(call) + c.wg.Add(1) + g.m[key] = c + g.mu.Unlock() + + c.val, c.err = fn() + g.mu.Lock() + defer g.mu.Unlock() + c.wg.Done() + delete(g.m, key) + return c.val, c.err +} diff --git a/coderd/util/singleflight/singleflight_test.go b/coderd/util/singleflight/singleflight_test.go new file mode 100644 index 0000000000000..91dcc7699312b --- /dev/null +++ b/coderd/util/singleflight/singleflight_test.go @@ -0,0 +1,85 @@ +package singleflight_test + +import ( + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/sync/errgroup" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/util/singleflight" +) + +func TestSingleflightGroup(t *testing.T) { + t.Parallel() + + t.Run("Sequential", func(t *testing.T) { + t.Parallel() + + group := singleflight.NewGroup(nil) + + var refreshCalls atomic.Int64 + fn := func() (any, error) { + return refreshCalls.Add(1), nil + } + + calls := 5 + for i := range calls { + result, err := group.Do("sequential", fn) + require.NoError(t, err) + require.Equal(t, int64(i+1), result) + } + + // Should have been called each time. + require.Equal(t, int64(calls), refreshCalls.Load()) + }) + + t.Run("Parallel", func(t *testing.T) { + t.Parallel() + + calls := 5 + + ch := make(chan string) + group := singleflight.NewGroup(ch) + + var refreshCalls atomic.Int64 + fn := func() (any, error) { + // Wait for calls to have joined the group before returning, otherwise it + // might return before all have joined and the test will flake. + if refreshCalls.Add(1) == 1 { + subscribed := 1 + for { + <-ch + subscribed++ + if subscribed >= calls { + return 1, nil + } + } + } + return 0, xerrors.New("should not be called") + } + + var eg errgroup.Group + results := make([]int, calls) + for i := range calls { + eg.Go(func() error { + result, err := group.Do("parallel", fn) + results[i] = result.(int) + return err + }) + } + + // No call should error. + err := eg.Wait() + require.NoError(t, err) + + // First group of calls should have a one. + for i := range calls { + require.Equal(t, 1, results[i]) + } + + // Should only have called once. + require.Equal(t, int64(1), refreshCalls.Load()) + }) +} diff --git a/testutil/oauth2.go b/testutil/oauth2.go index 1bdfdcb854a1a..b6cfc3b42cce8 100644 --- a/testutil/oauth2.go +++ b/testutil/oauth2.go @@ -10,6 +10,7 @@ import ( "golang.org/x/oauth2" "github.com/coder/coder/v2/coderd/promoauth" + "github.com/coder/coder/v2/coderd/util/singleflight" ) type OAuth2Config struct { @@ -18,6 +19,10 @@ type OAuth2Config struct { httpClientOnce sync.Once httpClient *http.Client + + Notifier chan string + groupOnce sync.Once + group *singleflight.Group } // Do issues req using a dedicated http.Client per OAuth2Config so a @@ -30,6 +35,13 @@ func (c *OAuth2Config) Do(_ context.Context, _ promoauth.Oauth2Source, req *http return c.httpClient.Do(req) } +func (c *OAuth2Config) Group() *singleflight.Group { + c.groupOnce.Do(func() { + c.group = singleflight.NewGroup(c.Notifier) + }) + return c.group +} + func (*OAuth2Config) AuthCodeURL(state string, _ ...oauth2.AuthCodeOption) string { return "/?state=" + url.QueryEscape(state) } From a836d1cd4f37c96bd498d55cd095e044bd159790 Mon Sep 17 00:00:00 2001 From: Asher Date: Tue, 14 Jul 2026 08:31:31 -0800 Subject: [PATCH 2/9] Ensure we always remove the singleflight group --- coderd/util/singleflight/singleflight.go | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/coderd/util/singleflight/singleflight.go b/coderd/util/singleflight/singleflight.go index 73dd2364933ac..4179109c7ed32 100644 --- a/coderd/util/singleflight/singleflight.go +++ b/coderd/util/singleflight/singleflight.go @@ -44,10 +44,13 @@ func (g *Group) Do(key string, fn func() (any, error)) (v any, err error) { g.m[key] = c g.mu.Unlock() + defer func() { + g.mu.Lock() + defer g.mu.Unlock() + c.wg.Done() + delete(g.m, key) + }() + c.val, c.err = fn() - g.mu.Lock() - defer g.mu.Unlock() - c.wg.Done() - delete(g.m, key) return c.val, c.err } From c385807cf70ddcc7220be1929ecbc9cd1b8426e3 Mon Sep 17 00:00:00 2001 From: Asher Date: Tue, 14 Jul 2026 09:14:24 -0800 Subject: [PATCH 3/9] Preserve re-read token test --- coderd/externalauth/externalauth_test.go | 63 ++++++++++++++++++++++-- 1 file changed, 58 insertions(+), 5 deletions(-) diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go index 4b32b1d7cf6aa..28daed1db287f 100644 --- a/coderd/externalauth/externalauth_test.go +++ b/coderd/externalauth/externalauth_test.go @@ -337,10 +337,10 @@ func TestRefreshToken(t *testing.T) { "permanent failures should not be retried") }) - // ConcurrentRefreshRace tests that when multiple concurrent requests race to - // refresh the same token, they share the same request instead of only one - // succeeding and the others using a now-invalid refresh token. - t.Run("ConcurrentRefreshRace", func(t *testing.T) { + // ConcurrentRefreshGroup tests that when requests try to refresh a token + // while another request is pending, they wait on the first caller and share + // the result instead of all attempting to perform the refresh. + t.Run("ConcurrentRefreshGroup", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) @@ -357,7 +357,7 @@ func TestRefreshToken(t *testing.T) { var refreshCalls atomic.Int64 config := &externalauth.Config{ InstrumentedOAuth2Config: &testutil.OAuth2Config{ - // The first call to refresh well succeed and all others will fail. The + // The first call to refresh will succeed and all others will fail. The // first will wait for all callers to join the group before returning. TokenSourceFunc: func() (*oauth2.Token, error) { if refreshCalls.Add(1) == 1 { @@ -415,6 +415,59 @@ func TestRefreshToken(t *testing.T) { require.Equal(t, int64(1), refreshCalls.Load()) }) + // 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. + // + // 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. + 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", + } + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) {}, + }) + + ctx := oidc.ClientContext(context.Background(), fake.HTTPClient(nil)) + link.OAuthExpiry = time.Now().Add(time.Hour * -1) + + // 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. + + 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) + }) + // ValidateFailure tests if the token is no longer valid with a 401 response. t.Run("ValidateFailure", func(t *testing.T) { t.Parallel() From 86163765e0521584adbf0d82b947dbceb77ee467 Mon Sep 17 00:00:00 2001 From: Asher Date: Tue, 14 Jul 2026 09:17:14 -0800 Subject: [PATCH 4/9] Better group ID --- coderd/externalauth/externalauth.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go index f75b99f55139e..269d35ee5ecce 100644 --- a/coderd/externalauth/externalauth.go +++ b/coderd/externalauth/externalauth.go @@ -193,7 +193,7 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu // 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. - key := externalAuthLink.OAuthAccessToken + key := c.ID + ":" + externalAuthLink.UserID.String() link, err := c.Group().Do(key, func() (any, error) { return c.innerRefreshToken(ctx, db, externalAuthLink) }) From 7a8a9bfccd044713c5b2e3f76815471eb47ce77d Mon Sep 17 00:00:00 2001 From: Asher Date: Tue, 14 Jul 2026 09:55:57 -0800 Subject: [PATCH 5/9] Use interface for singleflight group --- cli/create_test.go | 2 + coderd/aitasks_test.go | 5 ++ coderd/coderdtest/oidctest/idp.go | 2 + coderd/externalauth/externalauth.go | 13 ++- coderd/externalauth/externalauth_test.go | 58 ++++++++++++- coderd/externalauth_test.go | 17 +++- coderd/promoauth/oauth2.go | 21 +---- coderd/promoauth/oauth2_test.go | 2 + .../provisionerdserver_test.go | 2 + coderd/templateversions_test.go | 2 + coderd/util/singleflight/singleflight.go | 56 ------------ coderd/util/singleflight/singleflight_test.go | 85 ------------------- coderd/workspaceagents_test.go | 2 + coderd/workspaces_test.go | 6 ++ enterprise/aibridged_integration_test.go | 2 + testutil/oauth2.go | 12 --- 16 files changed, 111 insertions(+), 176 deletions(-) delete mode 100644 coderd/util/singleflight/singleflight.go delete mode 100644 coderd/util/singleflight/singleflight_test.go diff --git a/cli/create_test.go b/cli/create_test.go index 73778be1d63d6..b8fb1b4e64de6 100644 --- a/cli/create_test.go +++ b/cli/create_test.go @@ -11,6 +11,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/sync/singleflight" "github.com/coder/coder/v2/cli" "github.com/coder/coder/v2/cli/clitest" @@ -2117,6 +2118,7 @@ func TestCreateWithGitAuth(t *testing.T) { Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", + RefreshGroup: new(singleflight.Group), }}, IncludeProvisionerDaemon: true, }) diff --git a/coderd/aitasks_test.go b/coderd/aitasks_test.go index a5425eba62b13..9ee6df9e1ad40 100644 --- a/coderd/aitasks_test.go +++ b/coderd/aitasks_test.go @@ -17,6 +17,7 @@ import ( "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/sync/singleflight" "golang.org/x/xerrors" agentapisdk "github.com/coder/agentapi-sdk-go" @@ -1524,6 +1525,7 @@ func TestCreateTaskExternalAuth(t *testing.T) { Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) @@ -1582,6 +1584,7 @@ func TestCreateTaskExternalAuth(t *testing.T) { Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) @@ -1633,6 +1636,7 @@ func TestCreateTaskExternalAuth(t *testing.T) { Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) @@ -1665,6 +1669,7 @@ func TestCreateTaskExternalAuth(t *testing.T) { Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) diff --git a/coderd/coderdtest/oidctest/idp.go b/coderd/coderdtest/oidctest/idp.go index a7f608c632cfd..4bf6d0287dac1 100644 --- a/coderd/coderdtest/oidctest/idp.go +++ b/coderd/coderdtest/oidctest/idp.go @@ -33,6 +33,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" + "golang.org/x/sync/singleflight" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -1641,6 +1642,7 @@ func (f *FakeIDP) ExternalAuthConfig(t testing.TB, id string, custom *ExternalAu Scopes: []string{}, CodeURL: f.locked.Provider().DeviceCodeURL, }, + RefreshGroup: new(singleflight.Group), } if !custom.UseDeviceAuth { diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go index 269d35ee5ecce..14a22dfd71e0e 100644 --- a/coderd/externalauth/externalauth.go +++ b/coderd/externalauth/externalauth.go @@ -19,6 +19,7 @@ import ( "github.com/sqlc-dev/pqtype" "golang.org/x/oauth2" xgithub "golang.org/x/oauth2/github" + "golang.org/x/sync/singleflight" "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/database" @@ -53,6 +54,12 @@ const ( defaultRefreshRetryTimeout = 10 * time.Second ) +// SingleflightGroup exposes a subset of singleflight.Group for easier testing. +// singleflight.Group should be used instead of implementing this in production. +type SingleflightGroup interface { + Do(key string, fn func() (any, error)) (v any, err error, shared bool) +} + // Config is used for authentication for Git operations. type Config struct { promoauth.InstrumentedOAuth2Config @@ -142,6 +149,9 @@ type Config struct { // defaultRefreshRetryTimeout. A negative value disables transient-failure // retries entirely, so exactly one refresh attempt is made. RefreshRetryTimeout time.Duration + + // RefreshGroup deduplicates concurrent requests. + RefreshGroup SingleflightGroup } // Git returns a Provider for this config if the provider type is a @@ -194,7 +204,7 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu // in-flight refresh. Otherwise, the parallel calls will fail with a bad // refresh token error as they can only be used once. key := c.ID + ":" + externalAuthLink.UserID.String() - link, err := c.Group().Do(key, func() (any, error) { + link, err, _ := c.RefreshGroup.Do(key, func() (any, error) { return c.innerRefreshToken(ctx, db, externalAuthLink) }) if newlink, ok := link.(database.ExternalAuthLink); ok { @@ -936,6 +946,7 @@ func ConvertConfig(instrument *promoauth.Factory, entries []codersdk.ExternalAut MCPToolAllowRegex: mcpToolAllow, MCPToolDenyRegex: mcpToolDeny, CodeChallengeMethodsSupported: slice.StringEnums[promoauth.Oauth2PKCEChallengeMethod](entry.CodeChallengeMethodsSupported), + RefreshGroup: new(singleflight.Group), } if entry.DeviceFlow { diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go index 28daed1db287f..2253246f43883 100644 --- a/coderd/externalauth/externalauth_test.go +++ b/coderd/externalauth/externalauth_test.go @@ -9,6 +9,7 @@ import ( "net/http/httptest" "net/url" "strings" + "sync" "sync/atomic" "testing" "time" @@ -22,6 +23,7 @@ import ( "go.uber.org/mock/gomock" "golang.org/x/oauth2" "golang.org/x/sync/errgroup" + "golang.org/x/sync/singleflight" "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd" @@ -110,6 +112,7 @@ func TestRefreshToken(t *testing.T) { return nil, xerrors.New("failure") }, }, + RefreshGroup: new(singleflight.Group), } _, err := config.RefreshToken(context.Background(), nil, database.ExternalAuthLink{ @@ -374,7 +377,9 @@ func TestRefreshToken(t *testing.T) { } return nil, xerrors.New("bad_refresh_token") }, - Notifier: ch, + }, + RefreshGroup: &group{ + notify: ch, }, } @@ -1090,6 +1095,7 @@ func TestValidateToken(t *testing.T) { ID: "test-validate", Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), ValidateURL: validateURL, + RefreshGroup: new(singleflight.Group), } } @@ -1694,6 +1700,7 @@ func setupOauth2Test(t *testing.T, settings testConfig) (*oidctest.FakeIDP, *ext RevokeURL: fake.WellknownConfig().RevokeURL, RevokeTimeout: 1 * time.Second, CodeChallengeMethodsSupported: []promoauth.Oauth2PKCEChallengeMethod{promoauth.PKCEChallengeMethodSha256}, + RefreshGroup: new(singleflight.Group), } settings.ExternalAuthOpt(config) @@ -1770,3 +1777,52 @@ type roundTripper func(req *http.Request) (*http.Response, error) func (r roundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return r(req) } + +var _ externalauth.SingleflightGroup = (*group)(nil) + +type call struct { + wg sync.WaitGroup + val any + err error +} + +// group is like singleflight.Group +type group struct { + mu sync.Mutex // protects m + m map[string]*call // lazily initialized + notify chan string +} + +// Do ensures there is only one call to fn in flight at a time. Any calls that +// come in while it is in flight wait for the original call and get the same +// results. +// +// Blocks if the notifier blocks or is absent. +func (g *group) Do(key string, fn func() (any, error)) (v any, err error, shared bool) { + g.mu.Lock() + if g.m == nil { + g.m = make(map[string]*call) + } + if c, ok := g.m[key]; ok { + if g.notify != nil { + g.notify <- key + } + g.mu.Unlock() + c.wg.Wait() + return c.val, c.err, true + } + c := new(call) + c.wg.Add(1) + g.m[key] = c + g.mu.Unlock() + + defer func() { + g.mu.Lock() + defer g.mu.Unlock() + c.wg.Done() + delete(g.m, key) + }() + + c.val, c.err = fn() + return c.val, c.err, false +} diff --git a/coderd/externalauth_test.go b/coderd/externalauth_test.go index 4aa327313b10f..e30a81f861264 100644 --- a/coderd/externalauth_test.go +++ b/coderd/externalauth_test.go @@ -17,6 +17,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" + "golang.org/x/sync/singleflight" "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/coderdtest" @@ -519,6 +520,7 @@ func TestExternalAuthCallback(t *testing.T) { ID: "github", Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) user := coderdtest.CreateFirstUser(t, client) @@ -549,6 +551,7 @@ func TestExternalAuthCallback(t *testing.T) { ID: "github", Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) resp := coderdtest.RequestExternalAuthCallback(t, "github", client) @@ -563,6 +566,7 @@ func TestExternalAuthCallback(t *testing.T) { ID: "github", Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) _ = coderdtest.CreateFirstUser(t, client) @@ -586,6 +590,7 @@ func TestExternalAuthCallback(t *testing.T) { ID: "github", Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) maliciousHost := "https://malicious.com" @@ -619,6 +624,7 @@ func TestExternalAuthCallback(t *testing.T) { ID: "github", Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) user := coderdtest.CreateFirstUser(t, client) @@ -676,10 +682,11 @@ func TestExternalAuthCallback(t *testing.T) { Expiry: dbtime.Now().Add(-time.Hour), }, }, - ID: "github", - Regex: regexp.MustCompile(`github\.com`), - Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), - NoRefresh: true, + ID: "github", + Regex: regexp.MustCompile(`github\.com`), + Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + NoRefresh: true, + RefreshGroup: new(singleflight.Group), }}, }) user := coderdtest.CreateFirstUser(t, client) @@ -726,6 +733,7 @@ func TestExternalAuthCallback(t *testing.T) { ID: "github", Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) user := coderdtest.CreateFirstUser(t, client) @@ -791,6 +799,7 @@ func TestExternalAuthCallback(t *testing.T) { ID: "github", Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) user := coderdtest.CreateFirstUser(t, client) diff --git a/coderd/promoauth/oauth2.go b/coderd/promoauth/oauth2.go index 7d42779aa37e0..91b34dbd95019 100644 --- a/coderd/promoauth/oauth2.go +++ b/coderd/promoauth/oauth2.go @@ -9,8 +9,6 @@ import ( "github.com/prometheus/client_golang/prometheus" "github.com/prometheus/client_golang/prometheus/promauto" "golang.org/x/oauth2" - - "github.com/coder/coder/v2/coderd/util/singleflight" ) type Oauth2PKCEChallengeMethod string @@ -44,21 +42,16 @@ type OAuth2Config interface { TokenSource(context.Context, *oauth2.Token) oauth2.TokenSource } -// InstrumentedOAuth2Config extends OAuth2Config with: -// - A `Do` method that allows external oauth related calls to be -// instrumented. This is to support "ValidateToken" which is not an oauth2 -// specified method. These calls still count against the api rate limit, and -// should be instrumented. -// - A `Group` method that allows calls to be deduplicated. +// InstrumentedOAuth2Config extends OAuth2Config with a `Do` method that allows +// external oauth related calls to be instrumented. This is to support +// "ValidateToken" which is not an oauth2 specified method. +// These calls still count against the api rate limit, and should be instrumented. type InstrumentedOAuth2Config interface { OAuth2Config // Do is provided as a convenience method to make a request with the oauth2 client. // It mirrors `http.Client.Do`. Do(ctx context.Context, source Oauth2Source, req *http.Request) (*http.Response, error) - - // Group returns a singleflight group for deduplicating concurrent requests. - Group() *singleflight.Group } var _ OAuth2Config = (*Config)(nil) @@ -159,7 +152,6 @@ func (f *Factory) New(name string, under OAuth2Config) *Config { name: name, underlying: under, metrics: f.metrics, - group: singleflight.NewGroup(nil), } } @@ -210,7 +202,6 @@ type Config struct { metrics *metrics // interceptors are called after every request made by the oauth2 client. interceptors []func(resp *http.Response, err error) - group *singleflight.Group } func (c *Config) Do(ctx context.Context, source Oauth2Source, req *http.Request) (*http.Response, error) { @@ -218,10 +209,6 @@ func (c *Config) Do(ctx context.Context, source Oauth2Source, req *http.Request) return cli.Do(req) } -func (c *Config) Group() *singleflight.Group { - return c.group -} - func (c *Config) AuthCodeURL(state string, opts ...oauth2.AuthCodeOption) string { // No external requests are made when constructing the auth code url. return c.underlying.AuthCodeURL(state, opts...) diff --git a/coderd/promoauth/oauth2_test.go b/coderd/promoauth/oauth2_test.go index a2cb6f9bc4069..f2cd9dd83e7fa 100644 --- a/coderd/promoauth/oauth2_test.go +++ b/coderd/promoauth/oauth2_test.go @@ -13,6 +13,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/oauth2" + "golang.org/x/sync/singleflight" "github.com/coder/coder/v2/coderd/coderdtest/oidctest" "github.com/coder/coder/v2/coderd/coderdtest/promhelp" @@ -50,6 +51,7 @@ func TestInstrument(t *testing.T) { InstrumentedOAuth2Config: factory.New(id, idp.OIDCConfig(t, []string{})), ID: "test", ValidateURL: must[*url.URL](t)(idp.IssuerURL().Parse("/oauth2/userinfo")).String(), + RefreshGroup: new(singleflight.Group), } // 0 Requests before we start diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go index 8f0112decf732..80ad75a493dc7 100644 --- a/coderd/provisionerdserver/provisionerdserver_test.go +++ b/coderd/provisionerdserver/provisionerdserver_test.go @@ -23,6 +23,7 @@ import ( "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/trace" "golang.org/x/oauth2" + "golang.org/x/sync/singleflight" "golang.org/x/xerrors" "google.golang.org/protobuf/types/known/timestamppb" "storj.io/drpc" @@ -381,6 +382,7 @@ func TestAcquireJob(t *testing.T) { externalAuthConfigs: []*externalauth.Config{{ ID: gitAuthProvider.Id, InstrumentedOAuth2Config: &testutil.OAuth2Config{}, + RefreshGroup: new(singleflight.Group), }}, }) ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort) diff --git a/coderd/templateversions_test.go b/coderd/templateversions_test.go index c3d2153f3421e..9f5e494464fc4 100644 --- a/coderd/templateversions_test.go +++ b/coderd/templateversions_test.go @@ -15,6 +15,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/sync/errgroup" + "golang.org/x/sync/singleflight" "github.com/coder/coder/v2/coderd/audit" "github.com/coder/coder/v2/coderd/coderdtest" @@ -1009,6 +1010,7 @@ func TestTemplateVersionsExternalAuth(t *testing.T) { ID: "github", Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) user := coderdtest.CreateFirstUser(t, client) diff --git a/coderd/util/singleflight/singleflight.go b/coderd/util/singleflight/singleflight.go deleted file mode 100644 index 4179109c7ed32..0000000000000 --- a/coderd/util/singleflight/singleflight.go +++ /dev/null @@ -1,56 +0,0 @@ -package singleflight - -import ( - "sync" -) - -type call struct { - wg sync.WaitGroup - val any - err error -} - -// Group protects concurrent calls. -type Group struct { - mu sync.Mutex // protects m - m map[string]*call // lazily initialized - notify chan string -} - -func NewGroup(notifier chan string) *Group { - group := new(Group) - group.notify = notifier - return group -} - -// Do ensures there is only one call to fn in flight at a time. Any calls that -// come in while it is in flight wait for the original call and get the same -// results. -func (g *Group) Do(key string, fn func() (any, error)) (v any, err error) { - g.mu.Lock() - if g.m == nil { - g.m = make(map[string]*call) - } - if c, ok := g.m[key]; ok { - if g.notify != nil { - g.notify <- key - } - g.mu.Unlock() - c.wg.Wait() - return c.val, c.err - } - c := new(call) - c.wg.Add(1) - g.m[key] = c - g.mu.Unlock() - - defer func() { - g.mu.Lock() - defer g.mu.Unlock() - c.wg.Done() - delete(g.m, key) - }() - - c.val, c.err = fn() - return c.val, c.err -} diff --git a/coderd/util/singleflight/singleflight_test.go b/coderd/util/singleflight/singleflight_test.go deleted file mode 100644 index 91dcc7699312b..0000000000000 --- a/coderd/util/singleflight/singleflight_test.go +++ /dev/null @@ -1,85 +0,0 @@ -package singleflight_test - -import ( - "sync/atomic" - "testing" - - "github.com/stretchr/testify/require" - "golang.org/x/sync/errgroup" - "golang.org/x/xerrors" - - "github.com/coder/coder/v2/coderd/util/singleflight" -) - -func TestSingleflightGroup(t *testing.T) { - t.Parallel() - - t.Run("Sequential", func(t *testing.T) { - t.Parallel() - - group := singleflight.NewGroup(nil) - - var refreshCalls atomic.Int64 - fn := func() (any, error) { - return refreshCalls.Add(1), nil - } - - calls := 5 - for i := range calls { - result, err := group.Do("sequential", fn) - require.NoError(t, err) - require.Equal(t, int64(i+1), result) - } - - // Should have been called each time. - require.Equal(t, int64(calls), refreshCalls.Load()) - }) - - t.Run("Parallel", func(t *testing.T) { - t.Parallel() - - calls := 5 - - ch := make(chan string) - group := singleflight.NewGroup(ch) - - var refreshCalls atomic.Int64 - fn := func() (any, error) { - // Wait for calls to have joined the group before returning, otherwise it - // might return before all have joined and the test will flake. - if refreshCalls.Add(1) == 1 { - subscribed := 1 - for { - <-ch - subscribed++ - if subscribed >= calls { - return 1, nil - } - } - } - return 0, xerrors.New("should not be called") - } - - var eg errgroup.Group - results := make([]int, calls) - for i := range calls { - eg.Go(func() error { - result, err := group.Do("parallel", fn) - results[i] = result.(int) - return err - }) - } - - // No call should error. - err := eg.Wait() - require.NoError(t, err) - - // First group of calls should have a one. - for i := range calls { - require.Equal(t, 1, results[i]) - } - - // Should only have called once. - require.Equal(t, int64(1), refreshCalls.Load()) - }) -} diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 65c7eb1dbdf6d..f41921c7bb5f2 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -27,6 +27,7 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/oauth2" + "golang.org/x/sync/singleflight" "golang.org/x/xerrors" "google.golang.org/protobuf/types/known/timestamppb" "tailscale.com/tailcfg" @@ -3736,6 +3737,7 @@ func TestWorkspaceAgentsExternalAuthExpiresAt(t *testing.T) { Regex: regexp.MustCompile(`.*`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), // ValidateURL intentionally omitted: tokens are always valid. + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, ownerClient) diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go index 2c4627d3662ff..5140e7e91fea1 100644 --- a/coderd/workspaces_test.go +++ b/coderd/workspaces_test.go @@ -19,6 +19,7 @@ import ( "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/sync/singleflight" "cdr.dev/slog/v3" "github.com/coder/coder/v2/agent/agenttest" @@ -1501,6 +1502,7 @@ func TestCreateWorkspaceExternalAuth(t *testing.T) { Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) @@ -1553,6 +1555,7 @@ func TestCreateWorkspaceExternalAuth(t *testing.T) { Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) @@ -1601,6 +1604,7 @@ func TestCreateWorkspaceExternalAuth(t *testing.T) { Regex: regexp.MustCompile(`github\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) @@ -1640,6 +1644,7 @@ func TestCreateWorkspaceExternalAuth(t *testing.T) { Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), DisplayName: "GitHub", ValidateURL: validateSrv.URL, + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) @@ -1680,6 +1685,7 @@ func TestCreateWorkspaceExternalAuth(t *testing.T) { ID: "fallback-provider", Regex: regexp.MustCompile(`fallback\.example\.com`), Type: codersdk.EnhancedExternalAuthProviderGitHub.String(), + RefreshGroup: new(singleflight.Group), }}, }) first := coderdtest.CreateFirstUser(t, client) diff --git a/enterprise/aibridged_integration_test.go b/enterprise/aibridged_integration_test.go index 3b66f9e36d1df..03f71996e7edf 100644 --- a/enterprise/aibridged_integration_test.go +++ b/enterprise/aibridged_integration_test.go @@ -18,6 +18,7 @@ import ( "go.opentelemetry.io/otel/attribute" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" + "golang.org/x/sync/singleflight" "github.com/coder/coder/v2/aibridge" "github.com/coder/coder/v2/aibridge/aibridgetest" @@ -162,6 +163,7 @@ func TestIntegration(t *testing.T) { Type: "mock", DisplayName: "Mock", MCPURL: mockMCPServer.URL, + RefreshGroup: new(singleflight.Group), }, }, }, diff --git a/testutil/oauth2.go b/testutil/oauth2.go index b6cfc3b42cce8..1bdfdcb854a1a 100644 --- a/testutil/oauth2.go +++ b/testutil/oauth2.go @@ -10,7 +10,6 @@ import ( "golang.org/x/oauth2" "github.com/coder/coder/v2/coderd/promoauth" - "github.com/coder/coder/v2/coderd/util/singleflight" ) type OAuth2Config struct { @@ -19,10 +18,6 @@ type OAuth2Config struct { httpClientOnce sync.Once httpClient *http.Client - - Notifier chan string - groupOnce sync.Once - group *singleflight.Group } // Do issues req using a dedicated http.Client per OAuth2Config so a @@ -35,13 +30,6 @@ func (c *OAuth2Config) Do(_ context.Context, _ promoauth.Oauth2Source, req *http return c.httpClient.Do(req) } -func (c *OAuth2Config) Group() *singleflight.Group { - c.groupOnce.Do(func() { - c.group = singleflight.NewGroup(c.Notifier) - }) - return c.group -} - func (*OAuth2Config) AuthCodeURL(state string, _ ...oauth2.AuthCodeOption) string { return "/?state=" + url.QueryEscape(state) } From e411d18b28ffd78d4a25a1501501a7e39321b3b4 Mon Sep 17 00:00:00 2001 From: Asher Date: Tue, 14 Jul 2026 09:58:13 -0800 Subject: [PATCH 6/9] Improve request cancellation --- coderd/externalauth/externalauth.go | 19 ++- coderd/externalauth/externalauth_test.go | 169 +++++++++++++++++++---- 2 files changed, 155 insertions(+), 33 deletions(-) diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go index 14a22dfd71e0e..330ea72ebeb04 100644 --- a/coderd/externalauth/externalauth.go +++ b/coderd/externalauth/externalauth.go @@ -57,7 +57,7 @@ const ( // SingleflightGroup exposes a subset of singleflight.Group for easier testing. // singleflight.Group should be used instead of implementing this in production. type SingleflightGroup interface { - Do(key string, fn func() (any, error)) (v any, err error, shared bool) + DoChan(key string, fn func() (any, error)) <-chan singleflight.Result } // Config is used for authentication for Git operations. @@ -204,15 +204,20 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu // in-flight refresh. Otherwise, the parallel calls will fail with a bad // refresh token error as they can only be used once. key := c.ID + ":" + externalAuthLink.UserID.String() - link, err, _ := c.RefreshGroup.Do(key, func() (any, error) { + ch := c.RefreshGroup.DoChan(key, func() (any, error) { return c.innerRefreshToken(ctx, db, externalAuthLink) }) - if newlink, ok := link.(database.ExternalAuthLink); ok { - return newlink, err - } else if err == nil { - err = xerrors.Errorf("got invalid type from token refresh: %T", link) + select { + case results := <-ch: + if newlink, ok := results.Val.(database.ExternalAuthLink); ok { + return newlink, results.Err + } else if results.Err == nil { + return externalAuthLink, xerrors.Errorf("got invalid type from token refresh: %T", results.Val) + } + return externalAuthLink, results.Err + case <-ctx.Done(): + return externalAuthLink, ctx.Err() } - return externalAuthLink, err } func (c *Config) innerRefreshToken(ctx context.Context, db database.Store, externalAuthLink database.ExternalAuthLink) (database.ExternalAuthLink, error) { diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go index 2253246f43883..98e1f58fd4949 100644 --- a/coderd/externalauth/externalauth_test.go +++ b/coderd/externalauth/externalauth_test.go @@ -1,6 +1,7 @@ package externalauth_test import ( + "bytes" "context" "encoding/json" "fmt" @@ -8,6 +9,7 @@ import ( "net/http" "net/http/httptest" "net/url" + "runtime/debug" "strings" "sync" "sync/atomic" @@ -752,18 +754,21 @@ func TestRefreshToken(t *testing.T) { link.OAuthExpiry = expired _, err := config.RefreshToken(ctx, db, link) - require.NoError(t, err) + require.ErrorIs(t, err, context.Canceled) require.Equal(t, int64(1), refreshCalls.Load()) - dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ - ProviderID: link.ProviderID, - UserID: link.UserID, - }) - require.NoError(t, err) - require.NotEqual(t, oldAccessToken, dbLink.OAuthAccessToken, - "DB should have the new access token despite context cancellation") - require.NotEqual(t, oldRefreshToken, dbLink.OAuthRefreshToken, - "DB should have the new refresh token despite context cancellation") + require.Eventually(t, func() bool { + dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ + ProviderID: link.ProviderID, + UserID: link.UserID, + }) + if err != nil { + return false + } + return err == nil && + dbLink.OAuthAccessToken != oldAccessToken && + dbLink.OAuthRefreshToken != oldRefreshToken + }, testutil.WaitShort, testutil.IntervalFast, "never saw refresh token db updated") }) // SaveBeforeValidate_RateLimited tests the full path: refresh @@ -1780,49 +1785,161 @@ func (r roundTripper) RoundTrip(req *http.Request) (*http.Response, error) { var _ externalauth.SingleflightGroup = (*group)(nil) +// The following has been copied from x/sync/singleflight but has been modified +// to notify when callers join the group so the tests can be deterministic. + +// Copyright 2013 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// errGoexit indicates runtime.Goexit was called in +// the user-given function. +var errGoexit = xerrors.New("runtime.Goexit was called") + +// A panicError is an arbitrary value recovered from a panic +// with the stack trace during the execution of the given function. +type panicError struct { + value any + stack []byte +} + +// Error implements error interface. +func (p *panicError) Error() string { + return fmt.Sprintf("%v\n\n%s", p.value, p.stack) +} + +func (p *panicError) Unwrap() error { + err, ok := p.value.(error) + if !ok { + return nil + } + + return err +} + +func newPanicError(v any) error { + stack := debug.Stack() + + // The first line of the stack trace is of the form "goroutine N [status]:" + // but by the time the panic reaches Do the goroutine may no longer exist + // and its status will have changed. Trim out the misleading line. + if line := bytes.IndexByte(stack, '\n'); line >= 0 { + stack = stack[line+1:] + } + return &panicError{value: v, stack: stack} +} + +// call is an in-flight or completed singleflight.Do call type call struct { - wg sync.WaitGroup + wg sync.WaitGroup + + // These fields are written once before the WaitGroup is done + // and are only read after the WaitGroup is done. val any err error + + // These fields are read and written with the singleflight + // mutex held before the WaitGroup is done, and are read but + // not written after the WaitGroup is done. + dups int + chans []chan<- singleflight.Result } -// group is like singleflight.Group +// group represents a class of work and forms a namespace in +// which units of work can be executed with duplicate suppression. type group struct { mu sync.Mutex // protects m m map[string]*call // lazily initialized notify chan string } -// Do ensures there is only one call to fn in flight at a time. Any calls that -// come in while it is in flight wait for the original call and get the same -// results. +// DoChan is like Do but returns a channel that will receive the +// results when they are ready. // -// Blocks if the notifier blocks or is absent. -func (g *group) Do(key string, fn func() (any, error)) (v any, err error, shared bool) { +// The returned channel will not be closed. +func (g *group) DoChan(key string, fn func() (any, error)) <-chan singleflight.Result { + ch := make(chan singleflight.Result, 1) g.mu.Lock() if g.m == nil { g.m = make(map[string]*call) } if c, ok := g.m[key]; ok { - if g.notify != nil { - g.notify <- key - } + c.dups++ + c.chans = append(c.chans, ch) + g.notify <- key g.mu.Unlock() - c.wg.Wait() - return c.val, c.err, true + return ch } - c := new(call) + c := &call{chans: []chan<- singleflight.Result{ch}} c.wg.Add(1) g.m[key] = c g.mu.Unlock() + go g.doCall(c, key, fn) + + return ch +} + +// doCall handles the single call for a key. +func (g *group) doCall(c *call, key string, fn func() (any, error)) { + normalReturn := false + recovered := false + + // use double-defer to distinguish panic from runtime.Goexit, + // more details see https://golang.org/cl/134395 defer func() { + // the given function invoked runtime.Goexit + if !normalReturn && !recovered { + c.err = errGoexit + } + g.mu.Lock() defer g.mu.Unlock() c.wg.Done() - delete(g.m, key) + if g.m[key] == c { + delete(g.m, key) + } + + if e, ok := c.err.(*panicError); ok { + // In order to prevent the waiting channels from being blocked forever, + // needs to ensure that this panic cannot be recovered. + if len(c.chans) > 0 { + go panic(e) + select {} // Keep this goroutine around so that it will appear in the crash dump. + } else { + panic(e) + } + } else if c.err == errGoexit { + // Already in the process of goexit, no need to call again + } else { + // Normal return + for _, ch := range c.chans { + ch <- singleflight.Result{Val: c.val, Err: c.err, Shared: c.dups > 0} + } + } }() - c.val, c.err = fn() - return c.val, c.err, false + func() { + defer func() { + if !normalReturn { + // Ideally, we would wait to take a stack trace until we've determined + // whether this is a panic or a runtime.Goexit. + // + // Unfortunately, the only way we can distinguish the two is to see + // whether the recover stopped the goroutine from terminating, and by + // the time we know that, the part of the stack trace relevant to the + // panic has been discarded. + if r := recover(); r != nil { + c.err = newPanicError(r) + } + } + }() + + c.val, c.err = fn() + normalReturn = true + }() + + if !normalReturn { + recovered = true + } } From b21f399dc8ef302b199417b197ce3d533b117c55 Mon Sep 17 00:00:00 2001 From: Asher Date: Tue, 14 Jul 2026 11:14:27 -0800 Subject: [PATCH 7/9] nolint on imported singleflight --- coderd/externalauth/externalauth_test.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go index 98e1f58fd4949..34ce633555c6d 100644 --- a/coderd/externalauth/externalauth_test.go +++ b/coderd/externalauth/externalauth_test.go @@ -1900,16 +1900,18 @@ func (g *group) doCall(c *call, key string, fn func() (any, error)) { delete(g.m, key) } + //nolint:errorlint // Avoid changing the original code. if e, ok := c.err.(*panicError); ok { // In order to prevent the waiting channels from being blocked forever, // needs to ensure that this panic cannot be recovered. + //nolint:revive // Avoid changing the original code. if len(c.chans) > 0 { go panic(e) select {} // Keep this goroutine around so that it will appear in the crash dump. } else { panic(e) } - } else if c.err == errGoexit { + } else if c.err == errGoexit { //nolint:revive // Avoid changing the original code. // Already in the process of goexit, no need to call again } else { // Normal return From 785b08889f89f8c5cbee0c9460342d7f6963b0ca Mon Sep 17 00:00:00 2001 From: Asher Date: Tue, 14 Jul 2026 12:13:03 -0800 Subject: [PATCH 8/9] Handle concurrently canceled requests --- coderd/externalauth/externalauth.go | 11 ++- coderd/externalauth/externalauth_test.go | 100 +++++++++++++++++++++++ 2 files changed, 110 insertions(+), 1 deletion(-) diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go index 330ea72ebeb04..621f8c2be3655 100644 --- a/coderd/externalauth/externalauth.go +++ b/coderd/externalauth/externalauth.go @@ -205,7 +205,16 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu // refresh token error as they can only be used once. key := c.ID + ":" + externalAuthLink.UserID.String() ch := c.RefreshGroup.DoChan(key, func() (any, error) { - return c.innerRefreshToken(ctx, db, externalAuthLink) + // Use a detached context so if a request is canceled it does not cancel all + // the other requests as well. Preserve any original deadline. + rctx := context.WithoutCancel(ctx) + deadline, ok := ctx.Deadline() + if ok { + var cancel context.CancelFunc + rctx, cancel = context.WithDeadline(rctx, deadline) + defer cancel() + } + return c.innerRefreshToken(rctx, db, externalAuthLink) }) select { case results := <-ch: diff --git a/coderd/externalauth/externalauth_test.go b/coderd/externalauth/externalauth_test.go index 34ce633555c6d..8d91f6df9e2fb 100644 --- a/coderd/externalauth/externalauth_test.go +++ b/coderd/externalauth/externalauth_test.go @@ -475,6 +475,106 @@ func TestRefreshToken(t *testing.T) { require.Equal(t, "winner-refresh-token", result.OAuthRefreshToken) }) + // ConcurrentContextCancel tests that if one request is canceled, it does not + // cancel other requests waiting on it. + t.Run("ConcurrentContextCanceled", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + parallelRequests := 5 + ch := make(chan string) + + var refreshCalls atomic.Int64 + ctx := testutil.Context(t, testutil.WaitLong) + cancelOnRefresh, cancel := context.WithCancel(ctx) + defer cancel() + + // Use to know when the first call has started the group, so we know which + // context we can cancel. + listening := make(chan struct{}) + + fake, config, link := setupOauth2Test(t, testConfig{ + FakeIDPOpts: []oidctest.FakeIDPOpt{ + oidctest.WithRefresh(func(_ string) error { + if refreshCalls.Add(1) == 1 { + close(listening) + // Wait for all the other calls to be subscribed, to prevent + // the test from flaking. + subscribed := 1 + for { + <-ch + subscribed++ + if subscribed >= parallelRequests { + // Cancel the parent context after refresh succeeds + // but before the DB save and validation. + cancel() + return nil + } + } + } + // Should never reach here. + return xerrors.New("bad_refresh_token") + }), + oidctest.WithDynamicUserInfo(func(_ string) (jwt.MapClaims, error) { + return jwt.MapClaims{}, nil + }), + }, + ExternalAuthOpt: func(cfg *externalauth.Config) { + cfg.Type = codersdk.EnhancedExternalAuthProviderGitHub.String() + cfg.RefreshGroup = &group{notify: ch} + }, + DB: db, + }) + + oldAccessToken := link.OAuthAccessToken + oldRefreshToken := link.OAuthRefreshToken + link.OAuthExpiry = expired + + var wg sync.WaitGroup + // Start the first call with the cancelable context. + wg.Add(1) + go func() { + defer wg.Done() + ctx := oidc.ClientContext(cancelOnRefresh, fake.HTTPClient(nil)) + _, err := config.RefreshToken(ctx, db, link) + assert.ErrorIs(t, err, context.Canceled) + }() + + // Wait for it to start the group, to make sure the callback above is + // canceling the right context (if we fire them all at once, any one of them + // could start the group). + <-listening + + // Now we can fire off the remaining requests. + for range parallelRequests - 1 { + wg.Add(1) + go func() { + defer wg.Done() + ctx := oidc.ClientContext(ctx, fake.HTTPClient(nil)) + result, err := config.RefreshToken(ctx, db, link) + assert.NoError(t, err) + assert.NotEqual(t, oldAccessToken, result.OAuthAccessToken) + assert.NotEqual(t, oldRefreshToken, result.OAuthRefreshToken) + }() + } + + wg.Wait() + + // DB link should have been updated. + dbLink, err := db.GetExternalAuthLink(context.Background(), database.GetExternalAuthLinkParams{ + ProviderID: link.ProviderID, + UserID: link.UserID, + }) + require.NoError(t, err) + require.NotEqual(t, oldAccessToken, dbLink.OAuthAccessToken, + "DB should have the new access token despite context cancellation") + require.NotEqual(t, oldRefreshToken, dbLink.OAuthRefreshToken, + "DB should have the new refresh token despite context cancellation") + + // Only one refresh call should have actually been made. + require.Equal(t, int64(1), refreshCalls.Load()) + }) + // ValidateFailure tests if the token is no longer valid with a 401 response. t.Run("ValidateFailure", func(t *testing.T) { t.Parallel() From b484a70a35e47123b0332b16c8f0746f50ca6a8f Mon Sep 17 00:00:00 2001 From: Asher Date: Wed, 15 Jul 2026 10:38:57 -0800 Subject: [PATCH 9/9] Use a separate deadline --- coderd/externalauth/externalauth.go | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go index 621f8c2be3655..ed88a4843dd25 100644 --- a/coderd/externalauth/externalauth.go +++ b/coderd/externalauth/externalauth.go @@ -205,15 +205,16 @@ func (c *Config) RefreshToken(ctx context.Context, db database.Store, externalAu // refresh token error as they can only be used once. key := c.ID + ":" + externalAuthLink.UserID.String() ch := c.RefreshGroup.DoChan(key, func() (any, error) { - // Use a detached context so if a request is canceled it does not cancel all - // the other requests as well. Preserve any original deadline. - rctx := context.WithoutCancel(ctx) - deadline, ok := ctx.Deadline() - if ok { - var cancel context.CancelFunc - rctx, cancel = context.WithDeadline(rctx, deadline) - defer cancel() + // Use a detached context so if a request is canceled or times out it does + // not cancel all the other requests as well. The deadline is arbitrary but + // we give at least enough time for the refresh timeout then another 10 + // seconds for updating the database and validating the link. + timeout := 10 * time.Second + if c.RefreshRetryTimeout > 0 { + timeout += c.RefreshRetryTimeout } + rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout) + defer cancel() return c.innerRefreshToken(rctx, db, externalAuthLink) }) select { @@ -358,15 +359,9 @@ func (c *Config) innerRefreshToken(ctx context.Context, db database.Store, exter // 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. - // Use a detached context for the DB write only. The IDP already - // consumed the old refresh token, so if the caller's request - // context is canceled mid-save, the new token would be lost. - persistCtx, persistCancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) - defer persistCancel() - originalAccessToken := externalAuthLink.OAuthAccessToken if token.AccessToken != originalAccessToken { - updatedAuthLink, err := db.UpdateExternalAuthLink(persistCtx, database.UpdateExternalAuthLinkParams{ + updatedAuthLink, err := db.UpdateExternalAuthLink(ctx, database.UpdateExternalAuthLinkParams{ ProviderID: c.ID, UserID: externalAuthLink.UserID, UpdatedAt: dbtime.Now(),