[v16] Allow displaying hardware keys prompts when relogin is in progr…

…ess (#49702)

* Allow displaying hardware keys prompts when relogin is in progress (#48813)

* Convert `DialogConfirmation` to TS

* Allow `DialogConfirmation` and `Modal` to be hidden using CSS while closed

* Allow hiding all dialogs that are displayed as important

* Allow displaying multiple important dialogs, separate regular and important dialogs and get rid of `DialogNone`

* Pass `hidden` prop to important dialogs

* Rename `importantModalSemaphore` to `singleImportantModalSemaphore`

* Remove semaphores from hardware key prompts

* `keepMounted` -> `keepInDOMAfterClose`

* Remove the explicit value from `keepInDOMAfterClose`

* Revert splitting dialogs into regular and important ones, pass `hidden` to all of them, hide regular modal when important one is visible

* Use random id as a modal key

* Improve `singleImportantModalSemaphore` comment

* Improve `NewHardwareKeyPromptConstructor` comment

* Do not acquire important modal semaphore in MFA prompt and relogin, give each prompt its own mutex/semaphore

(cherry picked from commit 7428128)

* Re-add `keepInDOMAfterClose` to `Modal.jsx`

* Fix a mistake causing document reopen dialog to be always hidden (#49061)

(cherry picked from commit 7b520cc)

* Change `dialogCss` type to `any` since there is no `StyleFunction` available in `styled-components` in branch/v16

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()
@@ -912,7 +907,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
 }
@@ -1217,20 +1212,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

@@ -61,10 +61,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,9 +16,7 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import React from 'react';
-
-import { ButtonPrimary } from './../Button';
+import { ButtonPrimary } from '../Button';
 
 import DialogConfirmation, {
   DialogHeader,

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,35 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-import React from 'react';
+import { ReactNode } from 'react';
 
 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?: 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,

web/packages/design/src/Modal/Modal.jsx

@@ -179,11 +179,12 @@ export default class Modal extends React.Component {
       hideBackdrop,
       open,
       className,
+      keepInDOMAfterClose,
     } = this.props;
 
     const childProps = {};
 
-    if (!open) {
+    if (!open && !keepInDOMAfterClose) {
       return null;
     }
 
@@ -196,6 +197,7 @@ export default class Modal extends React.Component {
         data-testid="portal"
       >
         <StyledModal
+          hiddenInDom={!open}
           modalCss={modalCss}
           data-testid="Modal"
           ref={this.handleModalRef}
@@ -225,6 +227,11 @@ Modal.propTypes = {
    * A single child content element.
    */
   children: PropTypes.element,
+  /**
+   * Prevent unmounting the component and its children from the DOM when closed.
+   * Instead, hides it with CSS.
+   */
+  keepInDOMAfterClose: PropTypes.bool,
   /**
    * A node, component instance, or function that returns either.
    * The `container` will have the portal children appended to it.
@@ -298,6 +305,7 @@ Modal.propTypes = {
 };
 
 Modal.defaultProps = {
+  keepInDOMAfterClose: false,
   disableAutoFocus: false,
   disableBackdropClick: false,
   disableEnforceFocus: false,
@@ -339,5 +347,6 @@ const StyledModal = styled.div`
   bottom: 0;
   top: 0;
   left: 0;
-  ${props => props.modalCss && props.modalCss(props)}
+  ${props => props.hiddenInDom && `display: none;`};
+  ${props => props.modalCss?.(props)}
 `;