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.importantModalSemaphore.Acquire(ctx); err != nil {
-		return trace.Wrap(err)
-	}
-	defer s.importantModalSemaphore.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.importantModalSemaphore = newWaitSemaphore(maxConcurrentImportantModals, imporantModalWaitDuraiton)
+	s.headlessAuthSemaphore = newWaitSemaphore(maxConcurrentImportantModals, imporantModalWaitDuraiton)
 
 	return nil
 }
@@ -1191,20 +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 only display one important Modal at a time. tshd events
-	// that trigger an important modal (relogin, headless login) should use this
-	// lock to ensure it doesn't overwrite existing tshd-initiated important modals.
-	//
-	// We use a semaphore instead of a mutex in order to cancel important modals that
-	// are no longer relevant before acquisition.
+
+	// The Electron App can display multiple important modals by showing the latest one and hiding the others.
+	// 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.
 	//
-	// We use a waitSemaphore in order to make sure there is a clear transition between modals.
-	importantModalSemaphore *waitSemaphore
+	// 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.
+	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.importantModalSemaphore.Acquire(ctx); err != nil {
+	if err := s.headlessAuthSemaphore.Acquire(ctx); err != nil {
 		return trace.Wrap(err)
 	}
-	defer s.importantModalSemaphore.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.importantModalSemaphore.waitDuration = customWaitDuration
-	err = daemon.importantModalSemaphore.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.importantModalSemaphore.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,6 +30,22 @@ import (
 
 // NewHardwareKeyPromptConstructor returns a new hardware key prompt constructor
 // for this service and the given root cluster URI.
+//
+// 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):
+// 1. We don't offer a choice which key should be used on the initial login.
+// 2. Keys are cached per slot, not per physical key - it's not possible to use different keys with the same slot.
+//
+// Additionally, using the same hardware key for two clusters is not ideal too.
+// Since we cache the keys per slot, if two clusters specify the same one,
+// the user will always see the prompt for the same cluster URI. For example, if you are logged into both
+// cluster-a and cluster-b, the prompt will always say "Unlock hardware key to access cluster-b."
+// It seems that the better option would be to have a prompt per physical key, not per cluster.
+// 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}
 }
@@ -41,10 +57,6 @@ type hardwareKeyPrompter struct {
 
 // Touch prompts the user to touch the hardware key.
 func (h *hardwareKeyPrompter) Touch(ctx context.Context) error {
-	if err := h.s.importantModalSemaphore.Acquire(ctx); err != nil {
-		return trace.Wrap(err)
-	}
-	defer h.s.importantModalSemaphore.Release()
 	_, err := h.s.tshdEventsClient.PromptHardwareKeyTouch(ctx, &api.PromptHardwareKeyTouchRequest{
 		RootClusterUri: h.rootClusterURI.String(),
 	})
@@ -56,10 +68,6 @@ func (h *hardwareKeyPrompter) Touch(ctx context.Context) error {
 
 // AskPIN prompts the user for a PIN.
 func (h *hardwareKeyPrompter) AskPIN(ctx context.Context, requirement keys.PINPromptRequirement) (string, error) {
-	if err := h.s.importantModalSemaphore.Acquire(ctx); err != nil {
-		return "", trace.Wrap(err)
-	}
-	defer h.s.importantModalSemaphore.Release()
 	res, err := h.s.tshdEventsClient.PromptHardwareKeyPIN(ctx, &api.PromptHardwareKeyPINRequest{
 		RootClusterUri: h.rootClusterURI.String(),
 		PinOptional:    requirement == keys.PINOptional,
@@ -74,10 +82,6 @@ func (h *hardwareKeyPrompter) AskPIN(ctx context.Context, requirement keys.PINPr
 // The Electron app prompt must handle default values for PIN and PUK,
 // preventing the user from submitting empty/default values.
 func (h *hardwareKeyPrompter) ChangePIN(ctx context.Context) (*keys.PINAndPUK, error) {
-	if err := h.s.importantModalSemaphore.Acquire(ctx); err != nil {
-		return nil, trace.Wrap(err)
-	}
-	defer h.s.importantModalSemaphore.Release()
 	res, err := h.s.tshdEventsClient.PromptHardwareKeyPINChange(ctx, &api.PromptHardwareKeyPINChangeRequest{
 		RootClusterUri: h.rootClusterURI.String(),
 	})
@@ -93,10 +97,6 @@ func (h *hardwareKeyPrompter) ChangePIN(ctx context.Context) (*keys.PINAndPUK, e
 
 // ConfirmSlotOverwrite asks the user if the slot's private key and certificate can be overridden.
 func (h *hardwareKeyPrompter) ConfirmSlotOverwrite(ctx context.Context, message string) (bool, error) {
-	if err := h.s.importantModalSemaphore.Acquire(ctx); err != nil {
-		return false, trace.Wrap(err)
-	}
-	defer h.s.importantModalSemaphore.Release()
 	res, err := h.s.tshdEventsClient.ConfirmHardwareKeySlotOverwrite(ctx, &api.ConfirmHardwareKeySlotOverwriteRequest{
 		RootClusterUri: h.rootClusterURI.String(),
 		Message:        message,

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.importantModalSemaphore.Acquire(ctx); err != nil {
-		return nil, trace.Wrap(err)
-	}
-	defer s.importantModalSemaphore.Release()
+	s.mfaMu.Lock()
+	defer s.mfaMu.Unlock()
 
 	return s.tshdEventsClient.PromptMFA(ctx, in)
 }

web/packages/design/src/DialogConfirmation/DialogConfirmation.story.jsx → web/packages/design/src/DialogConfirmation/DialogConfirmation.story.tsx

@@ -16,8 +16,6 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import React from 'react';
-
 import { ButtonPrimary } from '../Button';
 
 import DialogConfirmation, {

web/packages/design/src/DialogConfirmation/DialogConfirmation.test.jsx → web/packages/design/src/DialogConfirmation/DialogConfirmation.test.tsx

@@ -16,11 +16,9 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import React from 'react';
-
 import { render, fireEvent } from 'design/utils/testing';
 
-import DialogConfirmation from './DialogConfirmation';
+import { DialogConfirmation } from './DialogConfirmation';
 
 test('onClose is respected', () => {
   const onClose = jest.fn();

web/packages/design/src/DialogConfirmation/DialogConfirmation.jsx → web/packages/design/src/DialogConfirmation/DialogConfirmation.tsx

@@ -16,22 +16,36 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import React from 'react';
+import { ReactNode } from 'react';
+import { StyleFunction } from 'styled-components';
 
 import Dialog from 'design/Dialog';
 
-function DialogConfirmation(props) {
-  const { children, open, onClose, dialogCss } = props;
+export function DialogConfirmation(props: {
+  open: boolean;
+  /**
+   * Prevent unmounting the component and its children from the DOM when closed.
+   * Instead, hides it with CSS.
+   */
+  keepInDOMAfterClose?: boolean;
+  /** @deprecated This props has no effect, it was never passed down to `Dialog`. */
+  disableEscapeKeyDown?: boolean;
+  children?: ReactNode;
+  onClose?: (
+    event: KeyboardEvent | React.MouseEvent,
+    reason: 'escapeKeyDown' | 'backdropClick'
+  ) => void;
+  dialogCss?: StyleFunction<any>;
+}) {
   return (
     <Dialog
-      dialogCss={dialogCss}
+      dialogCss={props.dialogCss}
       disableEscapeKeyDown={false}
-      onClose={onClose}
-      open={open}
+      onClose={props.onClose}
+      open={props.open}
+      keepInDOMAfterClose={props.keepInDOMAfterClose}
     >
-      {children}
+      {props.children}
     </Dialog>
   );
 }
-
-export default DialogConfirmation;

web/packages/design/src/DialogConfirmation/index.js → web/packages/design/src/DialogConfirmation/index.ts

@@ -16,7 +16,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import DialogConfirmation from './DialogConfirmation';
+import { DialogConfirmation } from './DialogConfirmation';
 import {
   DialogTitle,
   DialogContent,