fix FocusTrap escape due to strange tabindex values

This code will now ensure that we can't escape the FocusTrap if you use
`<tab>` and you happen to tab to an element outside of the FocusTrap
because the next item in line happens to be outside of the FocusTrap and
we never hit any of the focus guard elements.

How it works is as follows:

1. The `onBlur` is implemented on the `FocusTrap` itself, this will give
   us some information in the event itself.
   - `e.target` is the element that is being blurred (think of it as `from`)
   - `e.currentTarget` is the element with the event listener (the dialog)
   - `e.relatedTarget` is the element we are going to (think of it as `to`)
2. If the blur happened due to a `<tab>` or `<shift>+<tab>`, then we
   will move focus back inside the FocusTrap, and go from the `e.target`
   to the next or previous value.
3. If the blur happened programmatically (so no tab keys are involved,
   aka no direction is known), then the focus is restored to the
   `e.target` value.

Fixes: #1656

packages/@headlessui-react/src/components/dialog/dialog.test.tsx

@@ -218,7 +218,6 @@ describe('Rendering', () => {
     })
 
     it('should be possible to use a different render strategy for the Dialog', async () => {
-      let focusCounter = jest.fn()
       function Example() {
         let [isOpen, setIsOpen] = useState(false)
 
@@ -228,7 +227,7 @@ describe('Rendering', () => {
               Trigger
             </button>
             <Dialog open={isOpen} onClose={setIsOpen} unmount={false}>
-              <input onFocus={focusCounter} />
+              <input />
             </Dialog>
           </>
         )
@@ -239,17 +238,14 @@ describe('Rendering', () => {
       await nextFrame()
 
       assertDialog({ state: DialogState.InvisibleHidden })
-      expect(focusCounter).toHaveBeenCalledTimes(0)
 
       // Let's open the Dialog, to see if it is not hidden anymore
       await click(document.getElementById('trigger'))
-      expect(focusCounter).toHaveBeenCalledTimes(1)
 
       assertDialog({ state: DialogState.Visible })
 
       // Let's close the Dialog
       await press(Keys.Escape)
-      expect(focusCounter).toHaveBeenCalledTimes(1)
 
       assertDialog({ state: DialogState.InvisibleHidden })
     })

packages/@headlessui-react/src/components/focus-trap/focus-trap.test.tsx

@@ -1,11 +1,18 @@
-import React, { useState, useRef, FocusEvent } from 'react'
+import React, { useState, useRef } from 'react'
 import { render, screen } from '@testing-library/react'
 
 import { FocusTrap } from './focus-trap'
 import { assertActiveElement } from '../../test-utils/accessibility-assertions'
 import { suppressConsoleLogs } from '../../test-utils/suppress-console-logs'
 import { click, press, shift, Keys } from '../../test-utils/interactions'
 
+beforeAll(() => {
+  jest.spyOn(window, 'requestAnimationFrame').mockImplementation(setImmediate as any)
+  jest.spyOn(window, 'cancelAnimationFrame').mockImplementation(clearImmediate as any)
+})
+
+afterAll(() => jest.restoreAllMocks())
+
 function nextFrame() {
   return new Promise<void>((resolve) => {
     requestAnimationFrame(() => {
@@ -365,76 +372,134 @@ it('should be possible skip disabled elements within the focus trap', async () =
   assertActiveElement(document.getElementById('item-a'))
 })
 
-it('should try to focus all focusable items (and fail)', async () => {
-  let spy = jest.spyOn(console, 'warn').mockImplementation(jest.fn())
-  let focusHandler = jest.fn()
-  function handleFocus(e: FocusEvent) {
-    let target = e.target as HTMLElement
-    focusHandler(target.id)
-    screen.getByText('After')?.focus()
-  }
+it(
+  'should not be possible to programmatically escape the focus trap',
+  suppressConsoleLogs(async () => {
+    function Example() {
+      return (
+        <>
+          <input id="a" autoFocus />
 
-  render(
-    <>
-      <button id="before">Before</button>
-      <FocusTrap>
-        <button id="item-a" onFocus={handleFocus}>
-          Item A
-        </button>
-        <button id="item-b" onFocus={handleFocus}>
-          Item B
-        </button>
-        <button id="item-c" onFocus={handleFocus}>
-          Item C
-        </button>
-        <button id="item-d" onFocus={handleFocus}>
-          Item D
-        </button>
-      </FocusTrap>
-      <button>After</button>
-    </>
-  )
+          <FocusTrap>
+            <input id="b" />
+            <input id="c" />
+            <input id="d" />
+          </FocusTrap>
+        </>
+      )
+    }
 
-  await nextFrame()
+    render(<Example />)
 
-  expect(focusHandler.mock.calls).toEqual([['item-a'], ['item-b'], ['item-c'], ['item-d']])
-  expect(spy).toHaveBeenCalledWith('There are no focusable elements inside the <FocusTrap />')
-  spy.mockReset()
-})
+    await nextFrame()
 
-it('should end up at the last focusable element', async () => {
-  let spy = jest.spyOn(console, 'warn').mockImplementation(jest.fn())
+    let [a, b, c, d] = Array.from(document.querySelectorAll('input'))
 
-  let focusHandler = jest.fn()
-  function handleFocus(e: FocusEvent) {
-    let target = e.target as HTMLElement
-    focusHandler(target.id)
-    screen.getByText('After')?.focus()
-  }
+    // Ensure that input-b is the active element
+    assertActiveElement(b)
 
-  render(
-    <>
-      <button id="before">Before</button>
-      <FocusTrap>
-        <button id="item-a" onFocus={handleFocus}>
-          Item A
-        </button>
-        <button id="item-b" onFocus={handleFocus}>
-          Item B
-        </button>
-        <button id="item-c" onFocus={handleFocus}>
-          Item C
-        </button>
-        <button id="item-d">Item D</button>
-      </FocusTrap>
-      <button>After</button>
-    </>
-  )
+    // Tab to the next item
+    await press(Keys.Tab)
 
-  await nextFrame()
+    // Ensure that input-c is the active element
+    assertActiveElement(c)
 
-  expect(focusHandler.mock.calls).toEqual([['item-a'], ['item-b'], ['item-c']])
-  assertActiveElement(screen.getByText('Item D'))
-  expect(spy).not.toHaveBeenCalled()
-  spy.mockReset()
-})
+    // Try to move focus
+    a?.focus()
+
+    // Ensure that input-c is still the active element
+    assertActiveElement(c)
+
+    // Click on an element within the FocusTrap
+    await click(b)
+
+    // Ensure that input-b is the active element
+    assertActiveElement(b)
+
+    // Try to move focus again
+    a?.focus()
+
+    // Ensure that input-b is still the active element
+    assertActiveElement(b)
+
+    // Focus on an element within the FocusTrap
+    d?.focus()
+
+    // Ensure that input-d is the active element
+    assertActiveElement(d)
+
+    // Try to move focus again
+    a?.focus()
+
+    // Ensure that input-d is still the active element
+    assertActiveElement(d)
+  })
+)
+
+it(
+  'should not be possible to escape the FocusTrap due to strange tabIndex usage',
+  suppressConsoleLogs(async () => {
+    function Example() {
+      return (
+        <>
+          <div tabIndex={-1}>
+            <input tabIndex={2} id="a" />
+            <input tabIndex={1} id="b" />
+          </div>
+
+          <FocusTrap>
+            <input tabIndex={1} id="c" />
+            <input id="d" />
+          </FocusTrap>
+        </>
+      )
+    }
+
+    render(<Example />)
+
+    await nextFrame()
+
+    let [_a, _b, c, d] = Array.from(document.querySelectorAll('input'))
+
+    // First item in the FocusTrap should be the active one
+    assertActiveElement(c)
+
+    // Tab to the next item
+    await press(Keys.Tab)
+
+    // Ensure that input-d is the active element
+    assertActiveElement(d)
+
+    // Tab to the next item
+    await press(Keys.Tab)
+
+    // Ensure that input-c is the active element
+    assertActiveElement(c)
+
+    // Tab to the next item
+    await press(Keys.Tab)
+
+    // Ensure that input-d is the active element
+    assertActiveElement(d)
+
+    // Let's go the other way
+
+    // Tab to the previous item
+    await press(shift(Keys.Tab))
+
+    // Ensure that input-c is the active element
+    assertActiveElement(c)
+
+    // Tab to the previous item
+    await press(shift(Keys.Tab))
+
+    // Ensure that input-d is the active element
+    assertActiveElement(d)
+
+    // Tab to the previous item
+    await press(shift(Keys.Tab))
+
+    // Ensure that input-c is the active element
+    assertActiveElement(c)
+  })
+)

packages/@headlessui-react/src/components/focus-trap/focus-trap.tsx

@@ -6,6 +6,7 @@ import React, {
   ElementType,
   MutableRefObject,
   Ref,
+  FocusEvent as ReactFocusEvent,
 } from 'react'
 
 import { Props } from '../../types'
@@ -22,6 +23,7 @@ import { useOwnerDocument } from '../../hooks/use-owner'
 import { useEventListener } from '../../hooks/use-event-listener'
 import { microTask } from '../../utils/micro-task'
 import { useWatch } from '../../hooks/use-watch'
+import { useDisposables } from '../../hooks/use-disposables'
 
 let DEFAULT_FOCUS_TRAP_TAG = 'div' as const
 
@@ -75,34 +77,77 @@ export let FocusTrap = Object.assign(
     )
 
     let direction = useTabDirection()
-    let handleFocus = useEvent(() => {
+    let handleFocus = useEvent((e: ReactFocusEvent) => {
       let el = container.current as HTMLElement
       if (!el) return
 
       // TODO: Cleanup once we are using real browser tests
-      if (process.env.NODE_ENV === 'test') {
-        microTask(() => {
-          match(direction.current, {
-            [TabDirection.Forwards]: () => focusIn(el, Focus.First),
-            [TabDirection.Backwards]: () => focusIn(el, Focus.Last),
-          })
-        })
-      } else {
+      let wrapper = process.env.NODE_ENV === 'test' ? microTask : (cb: Function) => cb()
+      wrapper(() => {
         match(direction.current, {
-          [TabDirection.Forwards]: () => focusIn(el, Focus.First),
-          [TabDirection.Backwards]: () => focusIn(el, Focus.Last),
+          [TabDirection.Forwards]: () =>
+            focusIn(el, Focus.First, { skipElements: [e.relatedTarget as HTMLElement] }),
+          [TabDirection.Backwards]: () =>
+            focusIn(el, Focus.Last, { skipElements: [e.relatedTarget as HTMLElement] }),
         })
-      }
+      })
     })
 
-    let ourProps = { ref: focusTrapRef }
+    let d = useDisposables()
+    let recentlyUsedTabKey = useRef(false)
+    let ourProps = {
+      ref: focusTrapRef,
+      onKeyDown(e: KeyboardEvent) {
+        if (e.key == 'Tab') {
+          recentlyUsedTabKey.current = true
+          d.requestAnimationFrame(() => {
+            recentlyUsedTabKey.current = false
+          })
+        }
+      },
+      onBlur(e: ReactFocusEvent) {
+        let allContainers = new Set(containers?.current)
+        allContainers.add(container)
+
+        let relatedTarget = e.relatedTarget as HTMLElement | null
+        if (!relatedTarget) return
+
+        // Known guards, leave them alone!
+        if (relatedTarget.dataset.headlessuiFocusGuard === 'true') {
+          return
+        }
+
+        // Blur is triggered due to focus on relatedTarget, and the relatedTarget is not inside any
+        // of the dialog containers. In other words, let's move focus back in!
+        if (!contains(allContainers, relatedTarget)) {
+          // Was the blur invoke via the keyboard? Redirect to the next in line.
+          if (recentlyUsedTabKey.current) {
+            focusIn(
+              container.current as HTMLElement,
+              match(direction.current, {
+                [TabDirection.Forwards]: () => Focus.Next,
+                [TabDirection.Backwards]: () => Focus.Previous,
+              }) | Focus.WrapAround,
+              { relativeTo: e.target as HTMLElement }
+            )
+          }
+
+          // It was invoke via something else (e.g.: click, programmatically, ...). Redirect to the
+          // previous active item in the FocusTrap
+          else if (e.target instanceof HTMLElement) {
+            focusElement(e.target)
+          }
+        }
+      },
+    }
 
     return (
       <>
         {Boolean(features & Features.TabLock) && (
           <Hidden
             as="button"
             type="button"
+            data-headlessui-focus-guard
             onFocus={handleFocus}
             features={HiddenFeatures.Focusable}
           />
@@ -117,6 +162,7 @@ export let FocusTrap = Object.assign(
           <Hidden
             as="button"
             type="button"
+            data-headlessui-focus-guard
             onFocus={handleFocus}
             features={HiddenFeatures.Focusable}
           />

packages/@headlessui-react/src/utils/focus-management.ts

@@ -154,9 +154,10 @@ export function focusIn(
   container: HTMLElement | HTMLElement[],
   focus: Focus,
   {
-  sorted = true,
+    sorted = true,
     relativeTo = null,
-  }: Partial<{ sorted: boolean; relativeTo: HTMLElement | null }> = {}
+    skipElements = [],
+  }: Partial<{ sorted: boolean; relativeTo: HTMLElement | null; skipElements: HTMLElement[] }> = {}
 ) {
   let ownerDocument = Array.isArray(container)
     ? container.length > 0
@@ -170,6 +171,10 @@ export function focusIn(
       : container
     : getFocusableElements(container)
 
+  if (skipElements.length > 0) {
+    elements = elements.filter((x) => !skipElements.includes(x))
+  }
+
   relativeTo = relativeTo ?? (ownerDocument.activeElement as HTMLElement)
 
   let direction = (() => {