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/dialog/dialog.tsx

@@ -34,6 +34,7 @@ import { useOpenClosed, State } from '../../internal/open-closed'
 import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { StackProvider, StackMessage } from '../../internal/stack-context'
 import { useOutsideClick } from '../../hooks/use-outside-click'
+import { getOwnerDocument } from '../../utils/owner'
 
 enum DialogStates {
   Open,
@@ -231,17 +232,20 @@ let DialogRoot = forwardRefWithAs(function Dialog<
     if (dialogState !== DialogStates.Open) return
     if (hasParentDialog) return
 
-    let overflow = document.documentElement.style.overflow
-    let paddingRight = document.documentElement.style.paddingRight
+    let ownerDocument = getOwnerDocument(internalDialogRef.current)
+    let ownerWindow = ownerDocument.defaultView ?? window
 
-    let scrollbarWidth = window.innerWidth - document.documentElement.clientWidth
+    let overflow = ownerDocument.documentElement.style.overflow
+    let paddingRight = ownerDocument.documentElement.style.paddingRight
 
-    document.documentElement.style.overflow = 'hidden'
-    document.documentElement.style.paddingRight = `${scrollbarWidth}px`
+    let scrollbarWidth = ownerWindow.innerWidth - ownerDocument.documentElement.clientWidth
+
+    ownerDocument.documentElement.style.overflow = 'hidden'
+    ownerDocument.documentElement.style.paddingRight = `${scrollbarWidth}px`
 
     return () => {
-      document.documentElement.style.overflow = overflow
-      document.documentElement.style.paddingRight = paddingRight
+      ownerDocument.documentElement.style.overflow = overflow
+      ownerDocument.documentElement.style.paddingRight = paddingRight
     }
   }, [dialogState, hasParentDialog])
 

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

@@ -22,12 +22,13 @@ import React, {
 import { Props } from '../../types'
 import { match } from '../../utils/match'
 import { forwardRefWithAs, render, Features, PropsForFeatures } from '../../utils/render'
-import { useSyncRefs } from '../../hooks/use-sync-refs'
+import { optionalRef, useSyncRefs } from '../../hooks/use-sync-refs'
 import { useId } from '../../hooks/use-id'
 import { Keys } from '../keyboard'
 import { isDisabledReactIssue7711 } from '../../utils/bugs'
 import { OpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
+import { getOwnerDocument } from '../../utils/owner'
 
 enum DisclosureStates {
   Open,
@@ -158,7 +159,18 @@ let DisclosureRoot = forwardRefWithAs(function Disclosure<
   let { defaultOpen = false, ...passthroughProps } = props
   let buttonId = `headlessui-disclosure-button-${useId()}`
   let panelId = `headlessui-disclosure-panel-${useId()}`
-  let disclosureRef = useSyncRefs(ref)
+  let internalDisclosureRef = useRef<HTMLElement | null>(null)
+  let disclosureRef = useSyncRefs(
+    ref,
+    optionalRef(
+      (ref) => {
+        internalDisclosureRef.current = ref as unknown as HTMLElement | null
+      },
+      props.as === undefined ||
+        // @ts-expect-error The `as` prop _can_ be a Fragment
+        props.as === React.Fragment
+    )
+  )
 
   let panelRef = useRef<StateDefinition['panelRef']['current']>(null)
   let buttonRef = useRef<StateDefinition['buttonRef']['current']>(null)
@@ -179,13 +191,14 @@ let DisclosureRoot = forwardRefWithAs(function Disclosure<
   let close = useCallback(
     (focusableElement?: HTMLElement | MutableRefObject<HTMLElement | null>) => {
       dispatch({ type: ActionTypes.CloseDisclosure })
+      let ownerDocument = getOwnerDocument(internalDisclosureRef)
 
       let restoreElement = (() => {
-        if (!focusableElement) return document.getElementById(buttonId)
+        if (!focusableElement) return ownerDocument.getElementById(buttonId)
         if (focusableElement instanceof HTMLElement) return focusableElement
         if (focusableElement.current instanceof HTMLElement) return focusableElement.current
 
-        return document.getElementById(buttonId)
+        return ownerDocument.getElementById(buttonId)
       })()
 
       restoreElement?.focus()

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

@@ -36,6 +36,7 @@ import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
 import { useOutsideClick } from '../../hooks/use-outside-click'
 import { VisuallyHidden } from '../../internal/visually-hidden'
 import { objectToFormEntries } from '../../utils/form'
+import { getOwnerDocument } from '../../utils/owner'
 
 enum ListboxStates {
   Open,
@@ -559,7 +560,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])
@@ -704,7 +705,7 @@ let Option = forwardRefWithAs(function Option<
   let active =
     state.activeOptionIndex !== null ? state.options[state.activeOptionIndex].id === id : false
   let selected = state.propsRef.current.value === value
-  let internalOptionRef = useRef<HTMLElement | null>(null)
+  let internalOptionRef = useRef<HTMLLIElement | null>(null)
   let optionRef = useSyncRefs(ref, internalOptionRef)
 
   useIsoMorphicEffect(() => {

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

@@ -35,6 +35,7 @@ import { useOutsideClick } from '../../hooks/use-outside-click'
 import { useTreeWalker } from '../../hooks/use-tree-walker'
 import { useOpenClosed, State, OpenClosedProvider } from '../../internal/open-closed'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
+import { useOwnerDocument } from '../../hooks/use-owner'
 
 enum MenuStates {
   Open,
@@ -398,6 +399,7 @@ let Items = forwardRefWithAs(function Items<TTag extends ElementType = typeof DE
 ) {
   let [state, dispatch] = useMenuContext('Menu.Items')
   let itemsRef = useSyncRefs(state.itemsRef, ref)
+  let ownerDocument = useOwnerDocument(state.itemsRef)
 
   let id = `headlessui-menu-items-${useId()}`
   let searchDisposables = useDisposables()
@@ -415,10 +417,10 @@ 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 === ownerDocument?.activeElement) return
 
     container.focus({ preventScroll: true })
-  }, [state.menuState, state.itemsRef])
+  }, [state.menuState, state.itemsRef, ownerDocument])
 
   useTreeWalker({
     container: state.itemsRef.current,
@@ -501,7 +503,7 @@ let Items = forwardRefWithAs(function Items<TTag extends ElementType = typeof DE
           break
       }
     },
-    [dispatch, searchDisposables, state]
+    [dispatch, searchDisposables, state, ownerDocument]
   )
 
   let handleKeyUp = useCallback((event: ReactKeyboardEvent<HTMLButtonElement>) => {
@@ -571,7 +573,7 @@ let Item = forwardRefWithAs(function Item<TTag extends ElementType = typeof DEFA
   let [state, dispatch] = useMenuContext('Menu.Item')
   let id = `headlessui-menu-item-${useId()}`
   let active = state.activeItemIndex !== null ? state.items[state.activeItemIndex].id === id : false
-  let internalItemRef = useRef<HTMLElement | null>(null)
+  let internalItemRef = useRef<HTMLElement>(null)
   let itemRef = useSyncRefs(ref, internalItemRef)
 
   useIsoMorphicEffect(() => {

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

@@ -37,6 +37,8 @@ import { useWindowEvent } from '../../hooks/use-window-event'
 import { OpenClosedProvider, State, useOpenClosed } from '../../internal/open-closed'
 import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
 import { useOutsideClick } from '../../hooks/use-outside-click'
+import { getOwnerDocument } from '../../utils/owner'
+import { useOwnerDocument } from '../../hooks/use-owner'
 
 enum PopoverStates {
   Open,
@@ -198,7 +200,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])
 
@@ -308,17 +311,16 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
     isWithinPanel ? null : (button) => dispatch({ type: ActionTypes.SetButton, button })
   )
   let withinPanelButtonRef = useSyncRefs(internalButtonRef, ref)
+  let ownerDocument = useOwnerDocument(internalButtonRef)
 
   // TODO: Revisit when handling Tab/Shift+Tab when using Portal's
   let activeElementRef = useRef<Element | null>(null)
-  let previousActiveElementRef = useRef<Element | null>(
-    typeof window === 'undefined' ? null : document.activeElement
-  )
+  let previousActiveElementRef = useRef<Element | null>(null)
   useWindowEvent(
     'focus',
     () => {
       previousActiveElementRef.current = activeElementRef.current
-      activeElementRef.current = document.activeElement
+      activeElementRef.current = ownerDocument.activeElement as HTMLElement
     },
     true
   )
@@ -350,7 +352,7 @@ 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(ownerDocument.activeElement)) return
             event.preventDefault()
             event.stopPropagation()
             dispatch({ type: ActionTypes.ClosePopover })
@@ -369,7 +371,7 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
               if (state.panel.contains(previousActiveElementRef.current)) return
 
               // Check if the last focused element is *after* the button in the DOM
-              let focusableElements = getFocusableElements()
+              let focusableElements = getFocusableElements(ownerDocument.body)
               let previousIdx = focusableElements.indexOf(
                 previousActiveElementRef.current as HTMLElement
               )
@@ -425,7 +427,7 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
           if (state.panel.contains(previousActiveElementRef.current)) return
 
           // Check if the last focused element is *after* the button in the DOM
-          let focusableElements = getFocusableElements()
+          let focusableElements = getFocusableElements(ownerDocument.body)
           let previousIdx = focusableElements.indexOf(
             previousActiveElementRef.current as HTMLElement
           )
@@ -586,6 +588,7 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
   let panelRef = useSyncRefs(internalPanelRef, ref, (panel) => {
     dispatch({ type: ActionTypes.SetPanel, panel })
   })
+  let ownerDocument = useOwnerDocument(internalPanelRef)
 
   let usesOpenClosedState = useOpenClosed()
   let visible = (() => {
@@ -602,7 +605,7 @@ 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(ownerDocument.activeElement)) return
           event.preventDefault()
           event.stopPropagation()
           dispatch({ type: ActionTypes.ClosePopover })
@@ -631,7 +634,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 = ownerDocument.activeElement as HTMLElement
     if (internalPanelRef.current.contains(activeElement)) return // Already focused within Dialog
 
     focusIn(internalPanelRef.current, Focus.First)
@@ -642,9 +645,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 (!ownerDocument.activeElement) return
     if (!internalPanelRef.current) return
-    if (!internalPanelRef.current.contains(document.activeElement)) return
+    if (!internalPanelRef.current.contains(ownerDocument.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,
@@ -659,7 +662,7 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
     } else if (result === FocusResult.Overflow) {
       if (!state.button) return
 
-      let elements = getFocusableElements()
+      let elements = getFocusableElements(ownerDocument.body)
       let buttonIdx = elements.indexOf(state.button)
 
       let nextElements = elements
@@ -672,7 +675,7 @@ let Panel = forwardRefWithAs(function Panel<TTag extends ElementType = typeof DE
       // focusable). Therefore we will try and focus the very first item in
       // the document.body.
       if (focusIn(nextElements, Focus.First) === FocusResult.Error) {
-        focusIn(document.body, Focus.First)
+        focusIn(ownerDocument.body, Focus.First)
       }
     }
   })
@@ -685,7 +688,7 @@ 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(ownerDocument.activeElement)) return
       dispatch({ type: ActionTypes.ClosePopover })
     },
     true
@@ -753,15 +756,16 @@ let Group = forwardRefWithAs(function Group<TTag extends ElementType = typeof DE
   )
 
   let isFocusWithinPopoverGroup = useCallback(() => {
-    let element = document.activeElement as HTMLElement
+    let ownerDocument = getOwnerDocument(internalGroupRef)
+    let element = ownerDocument.activeElement
 
     if (internalGroupRef.current?.contains(element)) return true
 
     // Check if the focus is in one of the button or panel elements. This is important in case you are rendering inside a Portal.
     return popovers.some((bag) => {
       return (
-        document.getElementById(bag.buttonId)?.contains(element) ||
-        document.getElementById(bag.panelId)?.contains(element)
+        ownerDocument.getElementById(bag.buttonId)?.contains(element) ||
+        ownerDocument.getElementById(bag.panelId)?.contains(element)
       )
     })
   }, [internalGroupRef, popovers])

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

@@ -18,30 +18,36 @@ import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { usePortalRoot } from '../../internal/portal-force-root'
 import { useServerHandoffComplete } from '../../hooks/use-server-handoff-complete'
 import { useSyncRefs } from '../../hooks/use-sync-refs'
+import { useOwnerDocument } from '../../hooks/use-owner'
 
-function usePortalTarget(): HTMLElement | null {
+function usePortalTarget(ref: MutableRefObject<HTMLElement>): HTMLElement | null {
   let forceInRoot = usePortalRoot()
   let groupTarget = useContext(PortalGroupContext)
+
+  let ownerDocument = useOwnerDocument(ref)
+
   let [target, setTarget] = useState(() => {
     // Group context is used, but still null
     if (!forceInRoot && groupTarget !== null) return null
 
     // No group context is used, let's create a default portal root
     if (typeof window === 'undefined') return null
-    let existingRoot = document.getElementById('headlessui-portal-root')
+    let existingRoot = ownerDocument?.getElementById('headlessui-portal-root')
     if (existingRoot) return existingRoot
 
-    let root = document.createElement('div')
+    if (ownerDocument === null) return null
+
+    let root = ownerDocument.createElement('div')
     root.setAttribute('id', 'headlessui-portal-root')
-    return document.body.appendChild(root)
+    return ownerDocument.body.appendChild(root)
   })
 
   // Ensure the portal root is always in the DOM
   useEffect(() => {
     if (target === null) return
 
-    if (!document.body.contains(target)) {
-      document.body.appendChild(target)
+    if (!ownerDocument?.body.contains(target)) {
+      ownerDocument?.body.appendChild(target)
     }
   }, [target])
 
@@ -63,11 +69,14 @@ let PortalRoot = forwardRefWithAs(function Portal<
   TTag extends ElementType = typeof DEFAULT_PORTAL_TAG
 >(props: Props<TTag, PortalRenderPropArg>, ref: Ref<HTMLElement>) {
   let passthroughProps = props
-  let target = usePortalTarget()
+  let portalRef = useSyncRefs(ref)
+  // @ts-expect-error figure out the correct type here,
+  let ownerDocument = useOwnerDocument(ref)
+  // @ts-expect-error figure out the correct type here,
+  let target = usePortalTarget(ref)
   let [element] = useState<HTMLDivElement | null>(() =>
-    typeof window === 'undefined' ? null : document.createElement('div')
+    typeof window === 'undefined' ? null : ownerDocument?.createElement('div') ?? null
   )
-  let portalRef = useSyncRefs(ref)
 
   let ready = useServerHandoffComplete()