use ownerDocument instead of document

This should ensure that in iframes and new windows the correct document
is being used.

packages/@headlessui-react/src/components/listbox/listbox.tsx

@@ -33,6 +33,7 @@ import { isFocusableElement, FocusableMode } from '../../utils/focus-management'
 import { useWindowEvent } from '../../hooks/use-window-event'
 import { useOpenClosed, State, OpenClosedProvider } from '../../internal/open-closed'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
+import { getOwnerDocument } from '../../utils/owner-document'
 
 enum ListboxStates {
   Open,
@@ -498,7 +499,7 @@ let Options = forwardRefWithAs(function Options<
     let container = state.optionsRef.current
     if (!container) return
     if (state.listboxState !== ListboxStates.Open) return
-    if (container === document.activeElement) return
+    if (container === getOwnerDocument(container).activeElement) return
 
     container.focus({ preventScroll: true })
   }, [state.listboxState, state.optionsRef])

packages/@headlessui-react/src/components/menu/menu.tsx

@@ -35,6 +35,7 @@ import { useWindowEvent } from '../../hooks/use-window-event'
 import { useTreeWalker } from '../../hooks/use-tree-walker'
 import { useOpenClosed, State, OpenClosedProvider } from '../../internal/open-closed'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
+import { getOwnerDocument } from '../../utils/owner-document'
 
 enum MenuStates {
   Open,
@@ -382,7 +383,7 @@ let Items = forwardRefWithAs(function Items<TTag extends ElementType = typeof DE
     let container = state.itemsRef.current
     if (!container) return
     if (state.menuState !== MenuStates.Open) return
-    if (container === document.activeElement) return
+    if (container === getOwnerDocument(container).activeElement) return
 
     container.focus({ preventScroll: true })
   }, [state.menuState, state.itemsRef])

packages/@headlessui-react/src/components/popover/popover.tsx

@@ -36,6 +36,7 @@ import {
 import { useWindowEvent } from '../../hooks/use-window-event'
 import { OpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
+import { getOwnerDocument } from '../../utils/owner-document'
 
 enum PopoverStates {
   Open,
@@ -197,7 +198,8 @@ let PopoverRoot = forwardRefWithAs(function Popover<
   let isFocusWithinPopoverGroup = useCallback(() => {
     return (
       groupContext?.isFocusWithinPopoverGroup() ??
-      (button?.contains(document.activeElement) || panel?.contains(document.activeElement))
+      (button?.contains(getOwnerDocument(button).activeElement) ||
+        panel?.contains(getOwnerDocument(panel).activeElement))
     )
   }, [groupContext, button, panel])
 
@@ -322,7 +324,7 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
     'focus',
     () => {
       previousActiveElementRef.current = activeElementRef.current
-      activeElementRef.current = document.activeElement
+      activeElementRef.current = getOwnerDocument(internalButtonRef).activeElement as HTMLElement
     },
     true
   )
@@ -354,7 +356,10 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
           case Keys.Escape:
             if (state.popoverState !== PopoverStates.Open) return closeOthers?.(state.buttonId)
             if (!internalButtonRef.current) return
-            if (!internalButtonRef.current.contains(document.activeElement)) return
+            if (
+              !internalButtonRef.current.contains(getOwnerDocument(internalButtonRef).activeElement)
+            )
+              return
             event.preventDefault()
             event.stopPropagation()
             dispatch({ type: ActionTypes.ClosePopover })
@@ -606,7 +611,8 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
         case Keys.Escape:
           if (state.popoverState !== PopoverStates.Open) return
           if (!internalPanelRef.current) return
-          if (!internalPanelRef.current.contains(document.activeElement)) return
+          if (!internalPanelRef.current.contains(getOwnerDocument(internalPanelRef).activeElement))
+            return
           event.preventDefault()
           event.stopPropagation()
           dispatch({ type: ActionTypes.ClosePopover })
@@ -635,7 +641,7 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
     if (state.popoverState !== PopoverStates.Open) return
     if (!internalPanelRef.current) return
 
-    let activeElement = document.activeElement as HTMLElement
+    let activeElement = getOwnerDocument(internalPanelRef).activeElement as HTMLElement
     if (internalPanelRef.current.contains(activeElement)) return // Already focused within Dialog
 
     focusIn(internalPanelRef.current, Focus.First)
@@ -646,9 +652,9 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
     if (state.popoverState !== PopoverStates.Open) return
     if (!internalPanelRef.current) return
     if (event.key !== Keys.Tab) return
-    if (!document.activeElement) return
+    if (!getOwnerDocument(internalPanelRef).activeElement) return
     if (!internalPanelRef.current) return
-    if (!internalPanelRef.current.contains(document.activeElement)) return
+    if (!internalPanelRef.current.contains(getOwnerDocument(internalPanelRef).activeElement)) return
 
     // We will take-over the default tab behaviour so that we have a bit
     // control over what is focused next. It will behave exactly the same,
@@ -689,7 +695,8 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
       if (state.popoverState !== PopoverStates.Open) return
       if (!internalPanelRef.current) return
 
-      if (internalPanelRef.current?.contains(document.activeElement as HTMLElement)) return
+      if (internalPanelRef.current?.contains(getOwnerDocument(internalPanelRef).activeElement))
+        return
       dispatch({ type: ActionTypes.ClosePopover })
     },
     true
@@ -757,7 +764,7 @@ let Group = forwardRefWithAs(function Group<TTag extends ElementType = typeof DE
   )
 
   let isFocusWithinPopoverGroup = useCallback(() => {
-    let element = document.activeElement as HTMLElement
+    let element = getOwnerDocument(internalGroupRef).activeElement
 
     if (internalGroupRef.current?.contains(element)) return true
 

packages/@headlessui-react/src/components/radio-group/radio-group.tsx

@@ -26,6 +26,7 @@ import { Label, useLabels } from '../../components/label/label'
 import { Description, useDescriptions } from '../../components/description/description'
 import { useTreeWalker } from '../../hooks/use-tree-walker'
 import { useSyncRefs } from '../../hooks/use-sync-refs'
+import { getOwnerDocument } from '../../utils/owner-document'
 
 interface Option {
   id: string
@@ -184,7 +185,8 @@ let RadioGroupRoot = forwardRefWithAs(function RadioGroup<
 
             if (result === FocusResult.Success) {
               let activeOption = options.find(
-                (option) => option.element.current === document.activeElement
+                (option) =>
+                  option.element.current === getOwnerDocument(option.element).activeElement
               )
               if (activeOption) triggerChange(activeOption.propsRef.current.value)
             }
@@ -201,7 +203,8 @@ let RadioGroupRoot = forwardRefWithAs(function RadioGroup<
 
             if (result === FocusResult.Success) {
               let activeOption = options.find(
-                (option) => option.element.current === document.activeElement
+                (option) =>
+                  option.element.current === getOwnerDocument(option.element).activeElement
               )
               if (activeOption) triggerChange(activeOption.propsRef.current.value)
             }
@@ -214,7 +217,7 @@ let RadioGroupRoot = forwardRefWithAs(function RadioGroup<
             event.stopPropagation()
 
             let activeOption = options.find(
-              (option) => option.element.current === document.activeElement
+              (option) => option.element.current === getOwnerDocument(option.element).activeElement
             )
             if (activeOption) triggerChange(activeOption.propsRef.current.value)
           }

packages/@headlessui-react/src/hooks/use-focus-trap.ts

@@ -9,6 +9,7 @@ import { Keys } from '../components/keyboard'
 import { focusElement, focusIn, Focus, FocusResult } from '../utils/focus-management'
 import { useWindowEvent } from './use-window-event'
 import { useIsMounted } from './use-is-mounted'
+import { getOwnerDocument } from '../utils/owner-document'
 
 export enum Features {
   /** No features enabled for the `useFocusTrap` hook. */
@@ -42,7 +43,9 @@ export function useFocusTrap(
   } = {}
 ) {
   let restoreElement = useRef<HTMLElement | null>(
-    typeof window !== 'undefined' ? (document.activeElement as HTMLElement) : null
+    typeof window !== 'undefined'
+      ? (getOwnerDocument(container).activeElement as HTMLElement | null)
+      : null
   )
   let previousActiveElement = useRef<HTMLElement | null>(null)
   let mounted = useIsMounted()
@@ -54,7 +57,7 @@ export function useFocusTrap(
   useEffect(() => {
     if (!featuresRestoreFocus) return
 
-    restoreElement.current = document.activeElement as HTMLElement
+    restoreElement.current = getOwnerDocument(container).activeElement as HTMLElement
   }, [featuresRestoreFocus])
 
   // Restore the focus when we unmount the component.
@@ -72,7 +75,7 @@ export function useFocusTrap(
     if (!featuresInitialFocus) return
     if (!container.current) return
 
-    let activeElement = document.activeElement as HTMLElement
+    let activeElement = getOwnerDocument(container).activeElement as HTMLElement
 
     if (initialFocus?.current) {
       if (initialFocus?.current === activeElement) {
@@ -93,7 +96,7 @@ export function useFocusTrap(
       }
     }
 
-    previousActiveElement.current = document.activeElement as HTMLElement
+    previousActiveElement.current = getOwnerDocument(container).activeElement as HTMLElement
   }, [container, initialFocus, featuresInitialFocus])
 
   // Handle `Tab` & `Shift+Tab` keyboard events
@@ -111,7 +114,7 @@ export function useFocusTrap(
         (event.shiftKey ? Focus.Previous : Focus.Next) | Focus.WrapAround
       ) === FocusResult.Success
     ) {
-      previousActiveElement.current = document.activeElement as HTMLElement
+      previousActiveElement.current = getOwnerDocument(container).activeElement as HTMLElement
     }
   })
 

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

@@ -103,6 +103,12 @@ export function focusElement(element: HTMLElement | null) {
 }
 
 export function focusIn(container: HTMLElement | HTMLElement[], focus: Focus) {
+  let ownerDocument = Array.isArray(container)
+    ? container.length > 0
+      ? container[0].ownerDocument
+      : document
+    : container.ownerDocument
+
   let elements = Array.isArray(container)
     ? container.slice().sort((a, z) => {
         let position = a.compareDocumentPosition(z)
@@ -112,7 +118,7 @@ export function focusIn(container: HTMLElement | HTMLElement[], focus: Focus) {
         return 0
       })
     : getFocusableElements(container)
-  let active = document.activeElement as HTMLElement
+  let active = ownerDocument.activeElement as HTMLElement
 
   let direction = (() => {
     if (focus & (Focus.First | Focus.Next)) return Direction.Next
@@ -155,7 +161,7 @@ export function focusIn(container: HTMLElement | HTMLElement[], focus: Focus) {
 
     // Try the next one in line
     offset += direction
-  } while (next !== document.activeElement)
+  } while (next !== ownerDocument.activeElement)
 
   // This is a little weird, but let me try and explain: There are a few scenario's
   // in chrome for example where a focused `<a>` tag does not get the default focus

packages/@headlessui-react/src/utils/owner-document.ts

@@ -0,0 +1,14 @@
+import { MutableRefObject } from 'react'
+
+export function getOwnerDocument<T extends Element | MutableRefObject<Element | null>>(
+  element: T | null | undefined
+) {
+  if (element instanceof Element) return element.ownerDocument
+  if (element && element.hasOwnProperty('current')) {
+    if (element.current instanceof Element) {
+      return element.current.ownerDocument
+    }
+  }
+
+  return document
+}

packages/@headlessui-vue/src/components/listbox/listbox.ts

@@ -24,6 +24,7 @@ import { useWindowEvent } from '../../hooks/use-window-event'
 import { useOpenClosed, State, useOpenClosedProvider } from '../../internal/open-closed'
 import { match } from '../../utils/match'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
+import { getOwnerDocument } from '../../utils/owner-document'
 
 enum ListboxStates {
   Open,
@@ -209,7 +210,7 @@ export let Listbox = defineComponent({
 
     useWindowEvent('mousedown', (event) => {
       let target = event.target as HTMLElement
-      let active = document.activeElement
+      let active = getOwnerDocument(event.target as HTMLElement).activeElement
 
       if (listboxState.value !== ListboxStates.Open) return
       if (dom(buttonRef)?.contains(target)) return

packages/@headlessui-vue/src/components/menu/menu.ts

@@ -21,6 +21,7 @@ import { useTreeWalker } from '../../hooks/use-tree-walker'
 import { useOpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
 import { match } from '../../utils/match'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
+import { getOwnerDocument } from '../../utils/owner-document'
 
 enum MenuStates {
   Open,
@@ -162,7 +163,7 @@ export let Menu = defineComponent({
 
     useWindowEvent('mousedown', (event) => {
       let target = event.target as HTMLElement
-      let active = document.activeElement
+      let active = getOwnerDocument(event.target as HTMLElement).activeElement
 
       if (menuState.value !== MenuStates.Open) return
       if (dom(buttonRef)?.contains(target)) return

packages/@headlessui-vue/src/components/popover/popover.ts

@@ -28,6 +28,7 @@ import { dom } from '../../utils/dom'
 import { useWindowEvent } from '../../hooks/use-window-event'
 import { useOpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
+import { getOwnerDocument } from '../../utils/owner-document'
 
 enum PopoverStates {
   Open,
@@ -153,8 +154,8 @@ export let Popover = defineComponent({
     function isFocusWithinPopoverGroup() {
       return (
         groupContext?.isFocusWithinPopoverGroup() ??
-        (dom(button)?.contains(document.activeElement) ||
-          dom(panel)?.contains(document.activeElement))
+        (dom(button)?.contains(getOwnerDocument(button).activeElement) ||
+          dom(panel)?.contains(getOwnerDocument(panel).activeElement))
       )
     }
 
@@ -225,7 +226,7 @@ export let PopoverButton = defineComponent({
       'focus',
       () => {
         previousActiveElementRef.value = activeElementRef.value
-        activeElementRef.value = document.activeElement
+        activeElementRef.value = getOwnerDocument(api.button).activeElement
       },
       true
     )
@@ -269,7 +270,7 @@ export let PopoverButton = defineComponent({
           case Keys.Escape:
             if (api.popoverState.value !== PopoverStates.Open) return closeOthers?.(api.buttonId)
             if (!dom(api.button)) return
-            if (!dom(api.button)?.contains(document.activeElement)) return
+            if (!dom(api.button)?.contains(getOwnerDocument(api.button).activeElement)) return
             event.preventDefault()
             event.stopPropagation()
             api.closePopover()
@@ -463,7 +464,7 @@ export let PopoverPanel = defineComponent({
       if (api.popoverState.value !== PopoverStates.Open) return
       if (!api.panel) return
 
-      let activeElement = document.activeElement as HTMLElement
+      let activeElement = getOwnerDocument(api.panel).activeElement as HTMLElement
       if (dom(api.panel)?.contains(activeElement)) return // Already focused within Dialog
 
       focusIn(dom(api.panel)!, Focus.First)
@@ -475,8 +476,8 @@ export let PopoverPanel = defineComponent({
       if (!dom(api.panel)) return
 
       if (event.key !== Keys.Tab) return
-      if (!document.activeElement) return
-      if (!dom(api.panel)?.contains(document.activeElement)) return
+      if (!getOwnerDocument(api.panel).activeElement) return
+      if (!dom(api.panel)?.contains(getOwnerDocument(api.panel).activeElement)) return
 
       // We will take-over the default tab behaviour so that we have a bit
       // control over what is focused next. It will behave exactly the same,
@@ -516,7 +517,8 @@ export let PopoverPanel = defineComponent({
         if (!focus) return
         if (api.popoverState.value !== PopoverStates.Open) return
         if (!dom(api.panel)) return
-        if (dom(api.panel)?.contains(document.activeElement as HTMLElement)) return
+        if (dom(api.panel)?.contains(getOwnerDocument(api.panel).activeElement as HTMLElement))
+          return
         api.closePopover()
       },
       true
@@ -536,7 +538,7 @@ export let PopoverPanel = defineComponent({
         case Keys.Escape:
           if (api.popoverState.value !== PopoverStates.Open) return
           if (!dom(api.panel)) return
-          if (!dom(api.panel)?.contains(document.activeElement)) return
+          if (!dom(api.panel)?.contains(getOwnerDocument(api.panel).activeElement)) return
           event.preventDefault()
           event.stopPropagation()
           api.closePopover()
@@ -594,7 +596,7 @@ export let PopoverGroup = defineComponent({
     }
 
     function isFocusWithinPopoverGroup() {
-      let element = document.activeElement as HTMLElement
+      let element = getOwnerDocument(groupRef).activeElement as HTMLElement
 
       if (dom(groupRef)?.contains(element)) return true