Allow displaying hardware keys prompts when relogin is in progress

lib/teleterm/daemon/daemon.go

@@ -114,11 +114,6 @@ func (s *Service) relogin(ctx context.Context, req *api.ReloginRequest) error {
 	}
 	defer s.reloginMu.Unlock()
 
-	if err := s.singleImportantModalSemaphore.Acquire(ctx); err != nil {
-		return trace.Wrap(err)
-	}
-	defer s.singleImportantModalSemaphore.Release()
-
 	const reloginUserTimeout = time.Minute
 	timeoutCtx, cancelTshdEventsCtx := context.WithTimeout(ctx, reloginUserTimeout)
 	defer cancelTshdEventsCtx()
@@ -892,7 +887,7 @@ func (s *Service) UpdateAndDialTshdEventsServerAddress(serverAddress string) err
 	client := api.NewTshdEventsServiceClient(conn)
 
 	s.tshdEventsClient = client
-	s.singleImportantModalSemaphore = newWaitSemaphore(maxConcurrentImportantModals, imporantModalWaitDuraiton)
+	s.headlessAuthSemaphore = newWaitSemaphore(maxConcurrentImportantModals, imporantModalWaitDuraiton)
 
 	return nil
 }
@@ -1191,26 +1186,32 @@ type Service struct {
 	gateways map[string]gateway.Gateway
 	// tshdEventsClient is a client to send events to the Electron App.
 	tshdEventsClient api.TshdEventsServiceClient
+
 	// The Electron App can display multiple important modals by showing the latest one and hiding the others.
-	// However, we should use this feature only when necessary (as of now, we use it to show hardware key prompts
-	// when relogin is in progress). In any other case, we should open one modal at a time.
-	// The special thing about hardware key prompts is that it's the only case where we need to complete
-	// an action started in one modal in another modal.
+	// However, we should be careful with it, and generally try to limit the number of prompts on the tshd side,
+	// to avoid flooding the app.
+	// Multiple prompts of the same type may also conflict with each other. This is currently possible with
+	// MFA prompts (see the mfaMu comment for details).
+	// Generally, only one prompt of each type (e.g., re-login, MFA) should be allowed at the same time.
 	//
-	// This semaphore also prevents concurrent MFA prompts when using VNet with per-session MFA.
+	// But why do we allow multiple important modals at all? It is necessary in specific scenarios where one
+	// modal action requires completing in another. Currently, there are two cases:
+	// 1. A hardware key prompt may appear during a re-login process.
+	// 2. An MFA prompt may appear during a re-login process.
+
+	// We use a waitSemaphore to make sure there is a clear transition between modals.
+	// We allow a single headless auth prompt at a time.
+	headlessAuthSemaphore *waitSemaphore
+	// mfaMu prevents concurrent MFA prompts. These can happen when using VNet with per-session MFA.
 	// Issuing an MFA prompt starts the Webauthn goroutine which prompts for key touch on hardware level.
 	// Webauthn does not support concurrent prompts, so without this semaphore, one of the prompts would fail immediately.
-	//
-	// We use a semaphore instead of a mutex in order to cancel important modals that
-	// are no longer relevant before acquisition.
-	//
-	// We use a waitSemaphore in order to make sure there is a clear transition between modals.
-	singleImportantModalSemaphore *waitSemaphore
+	mfaMu sync.Mutex
+	// reloginMu is used when a goroutine needs to request a relogin from the Electron app.
+	// We allow a single relogin prompt at a time.
+	reloginMu sync.Mutex
+
 	// usageReporter batches the events and sends them to prehog
 	usageReporter *usagereporter.UsageReporter
-	// reloginMu is used when a goroutine needs to request a relogin from the Electron app. Since the
-	// app can show only one login modal at a time, we need to submit only one request at a time.
-	reloginMu sync.Mutex
 	// headlessWatcherClosers holds a map of root cluster URIs to headless watchers.
 	headlessWatcherClosers   map[string]context.CancelFunc
 	headlessWatcherClosersMu sync.Mutex

lib/teleterm/daemon/daemon_headless.go

@@ -272,10 +272,10 @@ func (s *Service) sendPendingHeadlessAuthentication(ctx context.Context, ha *typ
 		HeadlessAuthenticationClientIp: ha.ClientIpAddress,
 	}
 
-	if err := s.singleImportantModalSemaphore.Acquire(ctx); err != nil {
+	if err := s.headlessAuthSemaphore.Acquire(ctx); err != nil {
 		return trace.Wrap(err)
 	}
-	defer s.singleImportantModalSemaphore.Release()
+	defer s.headlessAuthSemaphore.Release()
 
 	_, err := s.tshdEventsClient.SendPendingHeadlessAuthentication(ctx, req)
 	return trace.Wrap(err)

lib/teleterm/daemon/daemon_test.go

@@ -489,7 +489,7 @@ func TestRetryWithRelogin(t *testing.T) {
 	}
 }
 
-func TestImportantModalSemaphore(t *testing.T) {
+func TestConcurrentHeadlessAuthPrompts(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
@@ -522,54 +522,54 @@ func TestImportantModalSemaphore(t *testing.T) {
 	// Claim the important modal semaphore.
 
 	customWaitDuration := 10 * time.Millisecond
-	daemon.singleImportantModalSemaphore.waitDuration = customWaitDuration
-	err = daemon.singleImportantModalSemaphore.Acquire(ctx)
+	daemon.headlessAuthSemaphore.waitDuration = customWaitDuration
+	err = daemon.headlessAuthSemaphore.Acquire(ctx)
 	require.NoError(t, err)
 
-	// relogin and sending pending headless authentications should be blocked.
+	// Pending headless authentications should be blocked.
 
-	reloginErrC := make(chan error)
+	headlessPromptErr1 := make(chan error)
 	go func() {
-		reloginErrC <- daemon.relogin(ctx, &api.ReloginRequest{})
+		headlessPromptErr1 <- daemon.sendPendingHeadlessAuthentication(ctx, &types.HeadlessAuthentication{}, "")
 	}()
 
-	sphaErrC := make(chan error)
+	headlessPromptErr2 := make(chan error)
 	go func() {
-		sphaErrC <- daemon.sendPendingHeadlessAuthentication(ctx, &types.HeadlessAuthentication{}, "")
+		headlessPromptErr2 <- daemon.sendPendingHeadlessAuthentication(ctx, &types.HeadlessAuthentication{}, "")
 	}()
 
 	select {
-	case <-reloginErrC:
-		t.Error("relogin completed successfully without acquiring the important modal semaphore")
-	case <-sphaErrC:
-		t.Error("sendPendingHeadlessAuthentication completed successfully without acquiring the important modal semaphore")
+	case <-headlessPromptErr1:
+		t.Error("sendPendingHeadlessAuthentication for the first prompt completed successfully without acquiring the semaphore")
+	case <-headlessPromptErr2:
+		t.Error("sendPendingHeadlessAuthentication for the second prompt completed successfully without acquiring the semaphore")
 	case <-time.After(100 * time.Millisecond):
 	}
 
-	// if the request's ctx is canceled, they will unblock and return an error instead.
+	// If the request's ctx is canceled, they will unblock and return an error instead.
 
 	cancelCtx, cancel := context.WithCancel(ctx)
 	cancel()
 
-	err = daemon.relogin(cancelCtx, &api.ReloginRequest{})
+	err = daemon.sendPendingHeadlessAuthentication(cancelCtx, &types.HeadlessAuthentication{}, "")
 	require.Error(t, err)
 	err = daemon.sendPendingHeadlessAuthentication(cancelCtx, &types.HeadlessAuthentication{}, "")
 	require.Error(t, err)
 
-	// Release the semaphore. relogin and sending pending headless authentication should
+	// Release the semaphore. Pending headless authentication should
 	// complete successfully after a short delay between each semaphore release.
 
 	releaseTime := time.Now()
-	daemon.singleImportantModalSemaphore.Release()
+	daemon.headlessAuthSemaphore.Release()
 
 	var otherC chan error
 	select {
-	case err := <-reloginErrC:
+	case err := <-headlessPromptErr1:
 		require.NoError(t, err)
-		otherC = sphaErrC
-	case err := <-sphaErrC:
+		otherC = headlessPromptErr2
+	case err := <-headlessPromptErr2:
 		require.NoError(t, err)
-		otherC = reloginErrC
+		otherC = headlessPromptErr1
 	case <-time.After(time.Second):
 		t.Error("important modal operations failed to acquire unclaimed semaphore")
 	}
@@ -589,8 +589,7 @@ func TestImportantModalSemaphore(t *testing.T) {
 		t.Error("important modal semaphore should not be acquired before waiting the specified duration")
 	}
 
-	require.EqualValues(t, 1, service.reloginCount.Load(), "Unexpected number of calls to service.Relogin")
-	require.EqualValues(t, 1, service.sendPendingHeadlessAuthenticationCount.Load(), "Unexpected number of calls to service.SendPendingHeadlessAuthentication")
+	require.EqualValues(t, 2, service.sendPendingHeadlessAuthenticationCount.Load(), "Unexpected number of calls to service.SendPendingHeadlessAuthentication")
 }
 
 type mockTSHDEventsService struct {

lib/teleterm/daemon/hardwarekeyprompt.go

@@ -30,8 +30,6 @@ import (
 
 // NewHardwareKeyPromptConstructor returns a new hardware key prompt constructor
 // for this service and the given root cluster URI.
-// Unlike other modals triggered by tshd events, we do not acquire the singleImportantModalSemaphore here.
-// This allows these prompts to show up while a relogin modal is still opened.
 //
 // TODO(gzdunek): Improve multi-cluster and multi-hardware keys support.
 // The code in yubikey.go doesn't really support using multiple hardware keys (like one per cluster):
@@ -46,6 +44,7 @@ import (
 // But I will leave that for the future, it's hard to say how common these scenarios will be in Connect.
 //
 // Because the code in yubikey.go assumes you use a single key, we don't have any mutex here.
+// (unlike other modals triggered by tshd).
 // We don't expect receiving prompts from different hardware keys.
 func (s *Service) NewHardwareKeyPromptConstructor(rootClusterURI uri.ResourceURI) keys.HardwareKeyPrompt {
 	return &hardwareKeyPrompter{s: s, rootClusterURI: rootClusterURI}

lib/teleterm/daemon/mfaprompt.go

@@ -63,10 +63,8 @@ func (s *Service) NewMFAPrompt(resourceURI uri.ResourceURI, cfg *libmfa.PromptCo
 }
 
 func (s *Service) promptAppMFA(ctx context.Context, in *api.PromptMFARequest) (*api.PromptMFAResponse, error) {
-	if err := s.singleImportantModalSemaphore.Acquire(ctx); err != nil {
-		return nil, trace.Wrap(err)
-	}
-	defer s.singleImportantModalSemaphore.Release()
+	s.mfaMu.Lock()
+	defer s.mfaMu.Unlock()
 
 	return s.tshdEventsClient.PromptMFA(ctx, in)
 }