diff --git a/client/server/login_account_test.go b/client/server/login_account_test.go index 88a009fee..c113c8bfe 100644 --- a/client/server/login_account_test.go +++ b/client/server/login_account_test.go @@ -165,6 +165,29 @@ func TestReuseOAuthFlow_ForcedPromptRefusesTheCachedFlow(t *testing.T) { require.Equal(t, "code", resp.UserCode) } +func TestReplaceOAuthFlow_CancelsTheDisplacedWait(t *testing.T) { + s := New(internal.CtxInitState(context.Background()), "console", "", false, false, false, false) + cancelled := false + s.oauthAuthFlow = oauthAuthFlow{ + flow: &stubOAuthFlow{}, + info: auth.AuthFlowInfo{UserCode: "code"}, + expiresAt: time.Now().Add(time.Hour), + hint: "user@example.com", + accountPrompted: true, + waitCancel: func() { cancelled = true }, + } + + next := &stubOAuthFlow{} + s.replaceOAuthFlow(oauthAuthFlow{flow: next, info: auth.AuthFlowInfo{UserCode: "next"}}) + + require.True(t, cancelled, "the displaced wait was left without an owner") + require.Equal(t, next, s.oauthAuthFlow.flow) + require.Equal(t, "next", s.oauthAuthFlow.info.UserCode) + require.Empty(t, s.oauthAuthFlow.hint, "the previous flow's hint survived the replacement") + require.False(t, s.oauthAuthFlow.accountPrompted) + require.Nil(t, s.oauthAuthFlow.waitCancel, "the consumed cancel stayed on the record") +} + func newSSOTestServer(t *testing.T, hint string, accountPrompted bool, tokenEmail string) *Server { t.Helper() s := New(internal.CtxInitState(context.Background()), "console", "", false, false, false, false) diff --git a/client/server/server.go b/client/server/server.go index a121e1a98..d9bda1bca 100644 --- a/client/server/server.go +++ b/client/server/server.go @@ -743,15 +743,13 @@ func (s *Server) startSSOLogin(ctx context.Context, msg *proto.LoginRequest, con return nil, err } - s.mutex.Lock() - s.oauthAuthFlow = oauthAuthFlow{ + s.replaceOAuthFlow(oauthAuthFlow{ flow: oAuthFlow, info: authInfo, expiresAt: time.Now().Add(time.Duration(authInfo.ExpiresIn) * time.Second), hint: hint, accountPrompted: promptForAccount, - } - s.mutex.Unlock() + }) state.Set(internal.StatusNeedsLogin) @@ -763,6 +761,25 @@ func (s *Server) startSSOLogin(ctx context.Context, msg *proto.LoginRequest, con }, nil } +// replaceOAuthFlow installs next as the shared OAuth flow record and takes over +// the wait it displaces, so a WaitSSOLogin still parked on the old flow is not +// left without an owner: nothing would preempt it, and it could go on to run +// attemptLogin or mutate the record behind the new flow. +// +// The displaced cancel is read in the same critical section that replaces the +// record, so two callers racing here cannot both take the same predecessor. The +// cancel runs after the unlock — the displaced wait takes s.mutex as it unwinds. +func (s *Server) replaceOAuthFlow(next oauthAuthFlow) { + s.mutex.Lock() + staleCancel := s.oauthAuthFlow.waitCancel + s.oauthAuthFlow = next + s.mutex.Unlock() + + if staleCancel != nil { + staleCancel() + } +} + // reuseOAuthFlow returns the cached auth info when the previous flow targets // the same client and still has enough life left, and otherwise cancels the // stale wait and returns nil so the caller requests a fresh flow. @@ -1851,16 +1868,14 @@ func (s *Server) RequestJWTAuth( return nil, gstatus.Errorf(codes.Internal, "failed to request auth info: %v", err) } - // Replace the whole record: this flow carries no profile hint, and leaving - // the previous login's hint and accountPrompted in place would have - // WaitSSOLogin judge a later token against them. - s.mutex.Lock() - s.oauthAuthFlow = oauthAuthFlow{ + // This flow carries no profile hint: leaving the previous login's hint and + // accountPrompted in place would have WaitSSOLogin judge a later token + // against them. + s.replaceOAuthFlow(oauthAuthFlow{ flow: oAuthFlow, info: authInfo, expiresAt: time.Now().Add(time.Duration(authInfo.ExpiresIn) * time.Second), - } - s.mutex.Unlock() + }) return &proto.RequestJWTAuthResponse{ VerificationURI: authInfo.VerificationURI,