[focus trap] Fix focus leaking in shadow DOM

packages/mui-material/src/Unstable_TrapFocus/FocusTrap.test.tsx

@@ -1,7 +1,15 @@
 import * as React from 'react';
 import * as ReactDOM from 'react-dom';
 import { expect } from 'chai';
-import { act, createRenderer, fireEvent, reactMajor, screen } from '@mui/internal-test-utils';
+import {
+  act,
+  createRenderer,
+  fireEvent,
+  isJsdom,
+  reactMajor,
+  screen,
+  within,
+} from '@mui/internal-test-utils';
 import FocusTrap from '@mui/material/Unstable_TrapFocus';
 import Portal from '@mui/material/Portal';
 import getActiveElement from '../utils/getActiveElement';
@@ -11,6 +19,27 @@ interface GenericProps {
   [index: string]: any;
 }
 
+/**
+ * This file runs in both node and browser projects, so `vitest/browser` must be
+ * imported lazily from browser-only tests.
+ *
+ * Use the Vitest Browser userEvent provider for shadow-root keyboard interactions.
+ * The `user` returned by `render()` comes from @testing-library/user-event.
+ * Its Tab implementation computes the tab order from `document.querySelectorAll()`,
+ * which doesn't include focusable elements inside shadow roots.
+ *
+ * Keep the import specifier widened to `string` so package TypeScript checks
+ * don't load Vitest Browser's ambient matcher types.
+ */
+async function setupBrowser() {
+  const { page, userEvent: user } = await import('vitest/browser' as string);
+
+  return {
+    page,
+    user,
+  };
+}
+
 describe('<FocusTrap />', () => {
   const { clock, render } = createRenderer();
 
@@ -366,6 +395,177 @@ describe('<FocusTrap />', () => {
     expect(screen.getByRole('textbox')).toHaveFocus();
   });
 
+  describe.skipIf(isJsdom())('shadow DOM', () => {
+    it('loops focus when rendered in a shadow root', async () => {
+      const shadowHost = document.createElement('div');
+      document.body.appendChild(shadowHost);
+      const shadowRoot = shadowHost.attachShadow({ mode: 'open' });
+      const shadowContainer = document.createElement('div');
+      shadowRoot.appendChild(shadowContainer);
+
+      let unmount: (() => void) | undefined;
+
+      try {
+        ({ unmount } = render(
+          <FocusTrap open>
+            <div tabIndex={-1}>
+              <button type="button">first</button>
+              <button type="button">last</button>
+            </div>
+          </FocusTrap>,
+          { container: shadowContainer },
+        ));
+
+        const { user } = await setupBrowser();
+        const firstButton = within(shadowContainer).getByRole('button', { name: 'first' });
+        const lastButton = within(shadowContainer).getByRole('button', { name: 'last' });
+
+        await user.click(lastButton);
+        expect(shadowRoot.activeElement).to.equal(lastButton);
+
+        await user.tab();
+        expect(shadowRoot.activeElement).to.equal(firstButton);
+
+        await user.tab({ shift: true });
+        expect(shadowRoot.activeElement).to.equal(lastButton);
+      } finally {
+        unmount?.();
+        document.body.removeChild(shadowHost);
+      }
+    });
+
+    it('loops backward when Shift+Tab starts from the shadow-root trap container', async () => {
+      const shadowHost = document.createElement('div');
+      document.body.appendChild(shadowHost);
+      const shadowRoot = shadowHost.attachShadow({ mode: 'open' });
+      const shadowContainer = document.createElement('div');
+      shadowRoot.appendChild(shadowContainer);
+
+      let unmount: (() => void) | undefined;
+
+      try {
+        ({ unmount } = render(
+          <FocusTrap open>
+            <div tabIndex={-1} data-testid="root">
+              <button type="button">first</button>
+              <button type="button">last</button>
+            </div>
+          </FocusTrap>,
+          { container: shadowContainer },
+        ));
+
+        const { user } = await setupBrowser();
+        const root = within(shadowContainer).getByTestId('root');
+        const lastButton = within(shadowContainer).getByRole('button', { name: 'last' });
+
+        expect(shadowRoot.activeElement).to.equal(root);
+
+        await user.tab({ shift: true });
+        expect(shadowRoot.activeElement).to.equal(lastButton);
+      } finally {
+        unmount?.();
+        document.body.removeChild(shadowHost);
+      }
+    });
+
+    it('loops backward when Shift+Tab enters a lazy shadow-root trap from the document', async () => {
+      const shadowHost = document.createElement('div');
+      const outsideAfter = document.createElement('input');
+      outsideAfter.setAttribute('aria-label', 'after');
+      document.body.appendChild(shadowHost);
+      document.body.appendChild(outsideAfter);
+      const shadowRoot = shadowHost.attachShadow({ mode: 'open' });
+      const shadowContainer = document.createElement('div');
+      shadowRoot.appendChild(shadowContainer);
+
+      let unmount: (() => void) | undefined;
+
+      try {
+        ({ unmount } = render(
+          <FocusTrap open disableAutoFocus>
+            <div tabIndex={-1}>
+              <button type="button">first</button>
+              <button type="button">last</button>
+            </div>
+          </FocusTrap>,
+          { container: shadowContainer },
+        ));
+
+        const { page, user } = await setupBrowser();
+        const lastButton = within(shadowContainer).getByRole('button', { name: 'last' });
+
+        // In Firefox, filling through a provider locator is the reliable way
+        // to move focus to this dynamically appended input before pressing Shift+Tab.
+        // eslint-disable-next-line testing-library/prefer-screen-queries -- `page` is a Vitest Browser locator provider, not a render result.
+        await user.fill(page.getByRole('textbox', { name: 'after' }), 'focused');
+        expect(document.activeElement).to.equal(outsideAfter);
+
+        await user.tab({ shift: true });
+        expect(shadowRoot.activeElement).to.equal(lastButton);
+      } finally {
+        unmount?.();
+        document.body.removeChild(shadowHost);
+        document.body.removeChild(outsideAfter);
+      }
+    });
+
+    it('loops backward when Shift+Tab enters a shadow-root trap from an iframe document', async () => {
+      const frame = document.createElement('iframe');
+      frame.setAttribute('data-testid', 'focus-trap-shadow-root-iframe');
+      // Firefox in the full no-isolate browser run needs the iframe to load
+      // a real same-origin document before Vitest Browser can query inside it reliably.
+      frame.srcdoc = '<!doctype html><html><body></body></html>';
+      const frameLoaded = new Promise<void>((resolve) => {
+        frame.addEventListener('load', () => resolve(), { once: true });
+      });
+      document.body.appendChild(frame);
+      await frameLoaded;
+      const frameDocument = frame.contentDocument!;
+      const shadowHost = frameDocument.createElement('div');
+      const outsideAfter = frameDocument.createElement('input');
+      outsideAfter.setAttribute('aria-label', 'after');
+      frameDocument.body.appendChild(shadowHost);
+      frameDocument.body.appendChild(outsideAfter);
+      const shadowRoot = shadowHost.attachShadow({ mode: 'open' });
+      const shadowContainer = frameDocument.createElement('div');
+      shadowRoot.appendChild(shadowContainer);
+
+      let unmount: (() => void) | undefined;
+
+      try {
+        ({ unmount } = render(
+          <FocusTrap open disableAutoFocus>
+            <div tabIndex={-1}>
+              <button type="button">first</button>
+              <button type="button">last</button>
+            </div>
+          </FocusTrap>,
+          { container: shadowContainer },
+        ));
+
+        const { page, user } = await setupBrowser();
+        // Use a unique page locator instead of `page.elementLocator(frame)`.
+        // Vitest Browser serializes element locators for iframes to a broad `iframe`
+        // selector, which can bind to a stale iframe in the full no-isolate Firefox run.
+        // eslint-disable-next-line testing-library/prefer-screen-queries -- `page` is a Vitest Browser locator provider, not a render result.
+        const frameLocator = page.frameLocator(page.getByTestId('focus-trap-shadow-root-iframe'));
+        const lastButton = within(shadowContainer).getByRole('button', { name: 'last' });
+
+        // Passing the iframe-owned DOM element directly to `user.fill()` does
+        // not focus it reliably in Firefox.
+        // eslint-disable-next-line testing-library/prefer-screen-queries -- `frameLocator` is a Vitest Browser locator, not a render result.
+        await user.fill(frameLocator.getByLabelText('after'), 'focused');
+        expect(frameDocument.activeElement).to.equal(outsideAfter);
+
+        await user.tab({ shift: true });
+        expect(shadowRoot.activeElement).to.equal(lastButton);
+      } finally {
+        unmount?.();
+        document.body.removeChild(frame);
+      }
+    });
+  });
+
   it('restores focus when closed', () => {
     function Test(props: GenericProps) {
       return (

packages/mui-material/src/Unstable_TrapFocus/FocusTrap.tsx

@@ -123,6 +123,26 @@ function defaultIsEnabled(): boolean {
   return true;
 }
 
+const DOCUMENT_FRAGMENT_NODE = 11;
+
+function isShadowRoot(node: Node): node is ShadowRoot {
+  return node.nodeType === DOCUMENT_FRAGMENT_NODE && 'host' in node && 'activeElement' in node;
+}
+
+function isKeyboardEvent(event: Event): event is KeyboardEvent {
+  return 'key' in event && typeof event.key === 'string';
+}
+
+function getFocusRoot(node: HTMLElement): Document | ShadowRoot {
+  const rootNode = node.getRootNode();
+
+  if (isShadowRoot(rootNode)) {
+    return rootNode;
+  }
+
+  return ownerDocument(node);
+}
+
 /**
  * @ignore - internal component.
  */
@@ -169,7 +189,8 @@ function FocusTrap(props: FocusTrapProps): React.JSX.Element {
     }
 
     const doc = ownerDocument(rootRef.current);
-    const activeElement = getActiveElement(doc);
+    const focusRoot = getFocusRoot(rootRef.current);
+    const activeElement = getActiveElement(focusRoot) ?? getActiveElement(doc);
 
     // Prefer the explicitly marked focusable element. Fall back to the root
     // element for generic FocusTrap usage.
@@ -214,16 +235,22 @@ function FocusTrap(props: FocusTrapProps): React.JSX.Element {
     }
 
     const doc = ownerDocument(rootRef.current);
+    const focusRoot = getFocusRoot(rootRef.current);
+    const getRootActiveElement = () => getActiveElement(focusRoot) ?? getActiveElement(doc);
+
+    const handleLoopFocus: EventListener = (nativeEvent) => {
+      if (!isKeyboardEvent(nativeEvent)) {
+        return;
+      }
 
-    const loopFocus = (nativeEvent: KeyboardEvent) => {
       lastKeydown.current = nativeEvent;
 
       if (disableEnforceFocus || !isEnabled() || nativeEvent.key !== 'Tab') {
         return;
       }
 
       const rootElement = rootRef.current;
-      const activeElement = getActiveElement(doc);
+      const activeElement = getRootActiveElement();
 
       if (rootElement === null) {
         return;
@@ -280,16 +307,18 @@ function FocusTrap(props: FocusTrapProps): React.JSX.Element {
       }
     };
 
-    const contain = () => {
+    // Enforce the focus trap when focus moves outside the root or lands on
+    // one of the sentinel elements.
+    const handleFocusIn = () => {
       const rootElement = 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.
+      // The focus-in handler can be called between the component being unmounted and its cleanup function being run.
       if (rootElement === null) {
         return;
       }
 
-      const activeEl = getActiveElement(doc);
+      const activeEl = getRootActiveElement();
 
       if (!doc.hasFocus() || !isEnabled() || ignoreNextEnforceFocus.current) {
         ignoreNextEnforceFocus.current = false;
@@ -349,8 +378,13 @@ function FocusTrap(props: FocusTrapProps): React.JSX.Element {
       }
     };
 
-    doc.addEventListener('focusin', contain);
-    doc.addEventListener('keydown', loopFocus, true);
+    focusRoot.addEventListener('focusin', handleFocusIn);
+    focusRoot.addEventListener('keydown', handleLoopFocus, true);
+
+    if (focusRoot !== doc) {
+      doc.addEventListener('focusin', handleFocusIn);
+      doc.addEventListener('keydown', handleLoopFocus, true);
+    }
 
     // With Edge, Safari and Firefox, no focus related events are fired when the focused area stops being a focused area.
     // for example https://bugzilla.mozilla.org/show_bug.cgi?id=559561.
@@ -359,17 +393,22 @@ function FocusTrap(props: FocusTrapProps): React.JSX.Element {
     // The whatwg spec defines how the browser should behave but does not explicitly mention any events:
     // https://html.spec.whatwg.org/multipage/interaction.html#focus-fixup-rule.
     const interval = setInterval(() => {
-      const activeEl = getActiveElement(doc);
+      const activeEl = getRootActiveElement();
       if (activeEl && activeEl.tagName === 'BODY') {
-        contain();
+        handleFocusIn();
       }
     }, 50);
 
     return () => {
       clearInterval(interval);
 
-      doc.removeEventListener('focusin', contain);
-      doc.removeEventListener('keydown', loopFocus, true);
+      focusRoot.removeEventListener('focusin', handleFocusIn);
+      focusRoot.removeEventListener('keydown', handleLoopFocus, true);
+
+      if (focusRoot !== doc) {
+        doc.removeEventListener('focusin', handleFocusIn);
+        doc.removeEventListener('keydown', handleLoopFocus, true);
+      }
     };
   }, [disableAutoFocus, disableEnforceFocus, disableRestoreFocus, isEnabled, open, getTabbable]);