Fix FocusTrap escape due to strange tabindex values (#2093)

* sort DOM nodes using tabIndex first

It will still keep the same DOM order if tabIndex matches, thanks to
stable sorts!

* refactor `focusIn` API

All the arguments resulted in usage like `focusIn(container,
Focus.First, true, null)`, and to make things worse, we need to add
something else to this list in the future.

Instead, let's keep the `container` and the type of `Focus` as known
params, all the other things can sit in an options object.

* 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

* update changelog

Author: RobinMalfait

packages/@headlessui-react/CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fix regression where `displayValue` crashes ([#2087](https://github.com/tailwindlabs/headlessui/pull/2087))
 - Fix `displayValue` syncing when `Combobox.Input` is unmounted and re-mounted in different trees ([#2090](https://github.com/tailwindlabs/headlessui/pull/2090))
+- Fix FocusTrap escape due to strange tabindex values ([#2093](https://github.com/tailwindlabs/headlessui/pull/2093))
 
 ## [1.7.5] - 2022-12-08
 

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/components/popover/popover.tsx

@@ -792,7 +792,7 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
             }
           }
 
-          focusIn(combined, Focus.First, false)
+          focusIn(combined, Focus.First, { sorted: false })
         },
         [TabDirection.Backwards]: () => focusIn(el, Focus.Last),
       })