[FocusTrap] Fix disableEnforceFocus behavior (mui#38816)

packages/mui-base/src/FocusTrap/FocusTrap.tsx

@@ -212,76 +212,90 @@ function FocusTrap(props: FocusTrapProps): JSX.Element {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [open]);
 
-  React.useEffect(() => {
-    // We might render an empty child.
-    if (!open || !rootRef.current) {
-      return;
-    }
-
-    const doc = ownerDocument(rootRef.current);
-
-    const contain = (nativeEvent: FocusEvent | null) => {
-      const { current: rootElement } = rootRef;
+  const contain = React.useCallback(
+    (nativeEvent: FocusEvent | null) => {
+      const rootElement = rootRef.current;
+      const doc = ownerDocument(rootRef.current);
 
       // Cleanup functions are executed lazily in React 17.
       // Contain can be called between the component being unmounted and its cleanup function being run.
       if (rootElement === null) {
         return;
       }
 
-      if (
-        !doc.hasFocus() ||
-        disableEnforceFocus ||
-        !isEnabled() ||
-        ignoreNextEnforceFocus.current
-      ) {
+      if (!doc.hasFocus() || !isEnabled() || ignoreNextEnforceFocus.current) {
         ignoreNextEnforceFocus.current = false;
         return;
       }
 
-      if (!rootElement.contains(doc.activeElement)) {
-        // if the focus event is not coming from inside the children's react tree, reset the refs
-        if (
-          (nativeEvent && reactFocusEventTarget.current !== nativeEvent.target) ||
-          doc.activeElement !== reactFocusEventTarget.current
-        ) {
-          reactFocusEventTarget.current = null;
-        } else if (reactFocusEventTarget.current !== null) {
-          return;
-        }
+      // The focus is already inside
+      if (rootElement.contains(doc.activeElement)) {
+        return;
+      }
 
-        if (!activated.current) {
-          return;
-        }
+      // The disableEnforceFocus is set and the focus is outside of the focus trap (and sentinel nodes)
+      if (
+        disableEnforceFocus &&
+        doc.activeElement !== sentinelStart.current &&
+        doc.activeElement !== sentinelEnd.current
+      ) {
+        return;
+      }
 
-        let tabbable: string[] | HTMLElement[] = [];
-        if (
-          doc.activeElement === sentinelStart.current ||
-          doc.activeElement === sentinelEnd.current
-        ) {
-          tabbable = getTabbable(rootRef.current as HTMLElement);
-        }
+      // if the focus event is not coming from inside the children's react tree, reset the refs
+      if (
+        (nativeEvent && reactFocusEventTarget.current !== nativeEvent.target) ||
+        doc.activeElement !== reactFocusEventTarget.current
+      ) {
+        reactFocusEventTarget.current = null;
+      } else if (reactFocusEventTarget.current !== null) {
+        return;
+      }
 
-        if (tabbable.length > 0) {
-          const isShiftTab = Boolean(
-            lastKeydown.current?.shiftKey && lastKeydown.current?.key === 'Tab',
-          );
+      if (!activated.current) {
+        return;
+      }
 
-          const focusNext = tabbable[0];
-          const focusPrevious = tabbable[tabbable.length - 1];
+      let tabbable: string[] | HTMLElement[] = [];
+      if (
+        doc.activeElement === sentinelStart.current ||
+        doc.activeElement === sentinelEnd.current
+      ) {
+        tabbable = getTabbable(rootRef.current as HTMLElement);
+      }
 
-          if (typeof focusNext !== 'string' && typeof focusPrevious !== 'string') {
-            if (isShiftTab) {
-              focusPrevious.focus();
-            } else {
-              focusNext.focus();
-            }
+      // one of the sentinel nodes was focused, so move the focus
+      // to the first/last tabbable element inside the focus trap
+      if (tabbable.length > 0) {
+        const isShiftTab = Boolean(
+          lastKeydown.current?.shiftKey && lastKeydown.current?.key === 'Tab',
+        );
+
+        const focusNext = tabbable[0];
+        const focusPrevious = tabbable[tabbable.length - 1];
+
+        if (typeof focusNext !== 'string' && typeof focusPrevious !== 'string') {
+          if (isShiftTab) {
+            focusPrevious.focus();
+          } else {
+            focusNext.focus();
           }
-        } else {
-          rootElement.focus();
         }
+        // no tabbable elements in the trap focus or the focus was outside of the focus trap
+      } else {
+        rootElement.focus();
       }
-    };
+    },
+    [disableEnforceFocus, isEnabled, getTabbable],
+  );
+
+  React.useEffect(() => {
+    // We might render an empty child.
+    if (!open || !rootRef.current) {
+      return;
+    }
+
+    const doc = ownerDocument(rootRef.current);
 
     const loopFocus = (nativeEvent: KeyboardEvent) => {
       lastKeydown.current = nativeEvent;
@@ -323,7 +337,15 @@ function FocusTrap(props: FocusTrapProps): JSX.Element {
       doc.removeEventListener('focusin', contain);
       doc.removeEventListener('keydown', loopFocus, true);
     };
-  }, [disableAutoFocus, disableEnforceFocus, disableRestoreFocus, isEnabled, open, getTabbable]);
+  }, [
+    disableAutoFocus,
+    disableEnforceFocus,
+    disableRestoreFocus,
+    isEnabled,
+    open,
+    getTabbable,
+    contain,
+  ]);
 
   const onFocus = (event: FocusEvent) => {
     if (nodeToRestore.current === null) {

test/e2e/fixtures/FocusTrap/DisableEnforceFocusFocusTrap.tsx

@@ -0,0 +1,19 @@
+import * as React from 'react';
+import { FocusTrap } from '@mui/base/FocusTrap';
+
+export default function disableEnforceFocusFocusTrap() {
+  return (
+    <React.Fragment>
+      <button data-testid="initial-focus" type="button" autoFocus>
+        initial focus
+      </button>
+      <FocusTrap open disableEnforceFocus disableAutoFocus>
+        <div data-testid="root">
+          <button data-testid="inside-trap-focus" type="button">
+            inside focusable
+          </button>
+        </div>
+      </FocusTrap>
+    </React.Fragment>
+  );
+}

test/e2e/index.test.ts

@@ -172,6 +172,26 @@ describe('e2e', () => {
       await page.keyboard.press('Tab');
       await expect(screen.getByText('final tab target')).toHaveFocus();
     });
+
+    it('should not trap focus when clicking outside when disableEnforceFocus is set', async () => {
+      await renderFixture('FocusTrap/DisableEnforceFocusFocusTrap');
+
+      // initial focus is on the button outside of the trap focus
+      await expect(screen.getByTestId('initial-focus')).toHaveFocus();
+
+      // focus the button inside the trap focus
+      await page.keyboard.press('Tab');
+      await expect(screen.getByTestId('inside-trap-focus')).toHaveFocus();
+
+      // the focus is now trapped inside
+      await page.keyboard.press('Tab');
+      await expect(screen.getByTestId('inside-trap-focus')).toHaveFocus();
+
+      const initialFocus = (await screen.getByTestId('initial-focus'))!;
+      await initialFocus.click();
+
+      await expect(screen.getByTestId('initial-focus')).toHaveFocus();
+    });
   });
 
   describe('<Rating />', () => {