From 651d4fd014cf8c82e384226f14e616ba9e6ef2bd Mon Sep 17 00:00:00 2001
From: Prince Rathi <77000634+prathi-eightfold@users.noreply.github.com>
Date: Tue, 11 Oct 2022 14:19:50 +0530
Subject: [PATCH] fix: dialog focus trap and restore (#412)

* fix: dialog focus trap and restore

* fix: add prop types

* fix: graceful handling

* fix: handle zero focusable elements case

Co-authored-by: Prince Rathi <prathi@eighfold.ai>
---
 .../Dialog/BaseDialog/BaseDialog.tsx          | 142 +++++++++---------
 .../Dialog/BaseDialog/BaseDialog.types.ts     |   5 +
 src/shared/FocusTrap/hooks/useFocusTrap.ts    |  35 ++++-
 3 files changed, 105 insertions(+), 77 deletions(-)

diff --git a/src/components/Dialog/BaseDialog/BaseDialog.tsx b/src/components/Dialog/BaseDialog/BaseDialog.tsx
index b677c4edb..fe66352eb 100644
--- a/src/components/Dialog/BaseDialog/BaseDialog.tsx
+++ b/src/components/Dialog/BaseDialog/BaseDialog.tsx
@@ -9,6 +9,7 @@ import {
 import { IconName } from '../../Icon';
 import { ButtonShape, NeutralButton } from '../../Button';
 import { useScrollLock } from '../../../hooks/useScrollLock';
+import { FocusTrap } from '../../../shared/FocusTrap';
 import { NoFormStyle } from '../../Form/Context';
 import { useCanvasDirection } from '../../../hooks/useCanvasDirection';
 
@@ -47,6 +48,7 @@ export const BaseDialog: FC<BaseDialogProps> = React.forwardRef(
             visible,
             width,
             zIndex,
+            focusTrap = true,
             ...rest
         },
         ref: Ref<HTMLDivElement>
@@ -114,79 +116,81 @@ export const BaseDialog: FC<BaseDialogProps> = React.forwardRef(
 
         const getDialog = (): JSX.Element => (
             <NoFormStyle status override>
-                <div
-                    {...rest}
-                    ref={ref}
-                    role="dialog"
-                    aria-modal={true}
-                    aria-labelledby={labelId}
-                    style={dialogBackdropStyle}
-                    className={dialogBackdropClasses}
-                    onClick={(e: React.MouseEvent<HTMLDivElement>) => {
-                        maskClosable && onClose?.(e);
-                    }}
-                >
+                <FocusTrap trap={visible && focusTrap}>
                     <div
-                        className={dialogClasses}
-                        style={dialogStyle}
-                        onClick={stopPropagation}
+                        {...rest}
+                        ref={ref}
+                        role="dialog"
+                        aria-modal={true}
+                        aria-labelledby={labelId}
+                        style={dialogBackdropStyle}
+                        className={dialogBackdropClasses}
+                        onClick={(e: React.MouseEvent<HTMLDivElement>) => {
+                            maskClosable && onClose?.(e);
+                        }}
                     >
-                        <div className={headerClasses}>
-                            <span id={labelId}>
-                                {headerButtonProps && (
-                                    <NeutralButton
-                                        classNames={styles.headerButton}
-                                        shape={ButtonShape.Round}
-                                        iconProps={{ path: headerIcon }}
-                                        style={{
-                                            transform:
-                                                htmlDir === 'rtl'
-                                                    ? 'rotate(180deg)'
-                                                    : 'none',
-                                        }}
-                                        {...headerButtonProps}
-                                    />
-                                )}
-                                {header}
-                            </span>
-                            <span className={styles.headerButtons}>
-                                {actionButtonThreeProps && (
-                                    <NeutralButton
-                                        shape={ButtonShape.Round}
-                                        {...actionButtonThreeProps}
-                                    />
-                                )}
-                                {actionButtonTwoProps && (
-                                    <NeutralButton
-                                        shape={ButtonShape.Round}
-                                        {...actionButtonTwoProps}
-                                    />
-                                )}
-                                {actionButtonOneProps && (
-                                    <NeutralButton
-                                        shape={ButtonShape.Round}
-                                        {...actionButtonOneProps}
-                                    />
-                                )}
-                                {closable && (
-                                    <NeutralButton
-                                        ariaLabel={closeButtonAriaLabelText}
-                                        iconProps={{ path: closeIcon }}
-                                        shape={ButtonShape.Round}
-                                        onClick={onClose}
-                                        {...closeButtonProps}
-                                    />
-                                )}
-                            </span>
+                        <div
+                            className={dialogClasses}
+                            style={dialogStyle}
+                            onClick={stopPropagation}
+                        >
+                            <div className={headerClasses}>
+                                <span id={labelId}>
+                                    {headerButtonProps && (
+                                        <NeutralButton
+                                            classNames={styles.headerButton}
+                                            shape={ButtonShape.Round}
+                                            iconProps={{ path: headerIcon }}
+                                            style={{
+                                                transform:
+                                                    htmlDir === 'rtl'
+                                                        ? 'rotate(180deg)'
+                                                        : 'none',
+                                            }}
+                                            {...headerButtonProps}
+                                        />
+                                    )}
+                                    {header}
+                                </span>
+                                <span className={styles.headerButtons}>
+                                    {actionButtonThreeProps && (
+                                        <NeutralButton
+                                            shape={ButtonShape.Round}
+                                            {...actionButtonThreeProps}
+                                        />
+                                    )}
+                                    {actionButtonTwoProps && (
+                                        <NeutralButton
+                                            shape={ButtonShape.Round}
+                                            {...actionButtonTwoProps}
+                                        />
+                                    )}
+                                    {actionButtonOneProps && (
+                                        <NeutralButton
+                                            shape={ButtonShape.Round}
+                                            {...actionButtonOneProps}
+                                        />
+                                    )}
+                                    {closable && (
+                                        <NeutralButton
+                                            ariaLabel={closeButtonAriaLabelText}
+                                            iconProps={{ path: closeIcon }}
+                                            shape={ButtonShape.Round}
+                                            onClick={onClose}
+                                            {...closeButtonProps}
+                                        />
+                                    )}
+                                </span>
+                            </div>
+                            <div ref={scrollRef} className={bodyClasses}>
+                                {body}
+                            </div>
+                            {actions && (
+                                <div className={actionsClasses}>{actions}</div>
+                            )}
                         </div>
-                        <div ref={scrollRef} className={bodyClasses}>
-                            {body}
-                        </div>
-                        {actions && (
-                            <div className={actionsClasses}>{actions}</div>
-                        )}
                     </div>
-                </div>
+                </FocusTrap>
             </NoFormStyle>
         );
         return <Portal getContainer={() => parent}>{getDialog()}</Portal>;
diff --git a/src/components/Dialog/BaseDialog/BaseDialog.types.ts b/src/components/Dialog/BaseDialog/BaseDialog.types.ts
index 9b6d8fbf1..63f076ff6 100644
--- a/src/components/Dialog/BaseDialog/BaseDialog.types.ts
+++ b/src/components/Dialog/BaseDialog/BaseDialog.types.ts
@@ -100,6 +100,11 @@ export interface BaseDialogProps
      * Custom classes for the dialog wrapper
      */
     dialogWrapperClassNames?: string;
+    /**
+     * Unset this to disable focus trap
+     * @default true
+     */
+    focusTrap?: boolean;
     /**
      * The header of the dialog
      */
diff --git a/src/shared/FocusTrap/hooks/useFocusTrap.ts b/src/shared/FocusTrap/hooks/useFocusTrap.ts
index c43d4f224..afe729f87 100644
--- a/src/shared/FocusTrap/hooks/useFocusTrap.ts
+++ b/src/shared/FocusTrap/hooks/useFocusTrap.ts
@@ -5,28 +5,34 @@ import { eventKeys } from '../../utilities/eventKeys';
 const SELECTORS =
     'a[href], button, input, textarea, select, details, [tabindex]:not([tabindex="-1"], iframe, object, embed';
 
+const INITIAL_FOCUS_DELAY = 100;
+
 export function useFocusTrap(
     visible = true
 ): React.MutableRefObject<HTMLDivElement> {
     const elRef = useRef<any>(null);
-    const handleFocus = (e: React.KeyboardEvent) => {
-        const focusableEls = [
-            ...elRef.current.querySelectorAll(SELECTORS),
-        ].filter(
+    const restoreFocusRef = useRef<any>(null);
+
+    const getFocusableElements = () => {
+        return [...elRef.current.querySelectorAll(SELECTORS)].filter(
             (el) =>
                 !el.hasAttribute('disabled') && !el.getAttribute('aria-hidden')
         );
+    };
 
-        const firstFocusableEl = focusableEls[0];
-        const lastFocusableEl = focusableEls[focusableEls.length - 1];
-
+    const handleFocus = (e: React.KeyboardEvent) => {
         const isTabPressed = e.key === eventKeys.TAB;
         const isShiftPressed = e.key === eventKeys.SHIFTLEFT;
 
-        if (!isTabPressed) {
+        const focusableEls = getFocusableElements();
+
+        if (!isTabPressed || !focusableEls.length) {
             return;
         }
 
+        const firstFocusableEl = focusableEls[0];
+        const lastFocusableEl = focusableEls[focusableEls.length - 1];
+
         if (isShiftPressed) {
             if (document.activeElement === firstFocusableEl) {
                 /* shift + tab */
@@ -42,11 +48,24 @@ export function useFocusTrap(
         }
     };
 
+    const setUpFocus = () => {
+        if (!elRef.current) {
+            return;
+        }
+        restoreFocusRef.current = document.activeElement;
+        const elementToFocus = getFocusableElements()?.[0];
+        setTimeout(() => {
+            elementToFocus?.focus();
+        }, INITIAL_FOCUS_DELAY);
+    };
+
     useEffect(() => {
         if (visible) {
+            setUpFocus();
             elRef.current?.addEventListener('keydown', handleFocus);
         }
         return () => {
+            restoreFocusRef.current?.focus();
             elRef.current?.removeEventListener('keydown', handleFocus);
         };
     }, [visible]);