[client] Cancel the wait displaced by an OAuth flow replacement

Replacing the shared record with a whole struct value dropped the previous
flow's waitCancel, so an SSO browser wait still parked on it lost its
cancel: nothing could preempt it, and it could go on to run attemptLogin
or mutate the record behind the new flow. Both replacement sites now take
the displaced cancel over in the same critical section, via a shared
replaceOAuthFlow, and invoke it after the unlock.

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)

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,