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

@@ -36,6 +36,7 @@ import { useResolveButtonType } from '../../hooks/use-resolve-button-type'
 import { useLatestValue } from '../../hooks/use-latest-value'
 import { useTreeWalker } from '../../hooks/use-tree-walker'
 import { sortByDomNode } from '../../utils/focus-management'
+import { useOwnerDocument } from 'hooks/use-owner'
 
 enum ComboboxStates {
   Open,
@@ -863,6 +864,7 @@ let Option = forwardRefWithAs(function Option<
   let internalOptionRef = useRef<HTMLLIElement | null>(null)
   let bag = useRef<ComboboxOptionDataRef['current']>({ disabled, value, domRef: internalOptionRef })
   let optionRef = useSyncRefs(ref, internalOptionRef)
+  let ownerDocument = useOwnerDocument(internalOptionRef)
 
   useIsoMorphicEffect(() => {
     bag.current.disabled = disabled
@@ -871,8 +873,8 @@ let Option = forwardRefWithAs(function Option<
     bag.current.value = value
   }, [bag, value])
   useIsoMorphicEffect(() => {
-    bag.current.textValue = document.getElementById(id)?.textContent?.toLowerCase()
-  }, [bag, id])
+    bag.current.textValue = ownerDocument?.getElementById(id)?.textContent?.toLowerCase()
+  }, [bag, id, ownerDocument])
 
   let select = useCallback(() => actions.selectOption(id), [actions, id])
 
@@ -904,10 +906,10 @@ let Option = forwardRefWithAs(function Option<
     if (state.activationTrigger === ActivationTrigger.Pointer) return
     let d = disposables()
     d.requestAnimationFrame(() => {
-      document.getElementById(id)?.scrollIntoView?.({ block: 'nearest' })
+      ownerDocument?.getElementById(id)?.scrollIntoView?.({ block: 'nearest' })
     })
     return d.dispose
-  }, [id, active, state.comboboxState, state.activationTrigger, /* We also want to trigger this when the position of the active item changes so that we can re-trigger the scrollIntoView */ state.activeOptionIndex])
+  }, [id, active, state.comboboxState, ownerDocument, state.activationTrigger, /* We also want to trigger this when the position of the active item changes so that we can re-trigger the scrollIntoView */ state.activeOptionIndex])
 
   let handleClick = useCallback(
     (event: { preventDefault: Function }) => {

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,14 @@ 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'
+import { useOwnerDocument } from '../../hooks/use-owner'
 
 enum DisclosureStates {
   Open,
@@ -155,7 +157,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 reducerBag = useReducer(stateReducer, {
     disclosureState: defaultOpen ? DisclosureStates.Open : DisclosureStates.Closed,
@@ -171,13 +184,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()
@@ -234,6 +248,7 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
   let [state, dispatch] = useDisclosureContext('Disclosure.Button')
   let internalButtonRef = useRef<HTMLButtonElement | null>(null)
   let buttonRef = useSyncRefs(internalButtonRef, ref)
+  let ownerDocument = useOwnerDocument(internalButtonRef)
 
   let panelContext = useDisclosurePanelContext()
   let isWithinPanel = panelContext === null ? false : panelContext === state.panelId
@@ -249,7 +264,7 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
             event.preventDefault()
             event.stopPropagation()
             dispatch({ type: ActionTypes.ToggleDisclosure })
-            document.getElementById(state.buttonId)?.focus()
+            ownerDocument?.getElementById(state.buttonId)?.focus()
             break
         }
       } else {
@@ -263,7 +278,7 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
         }
       }
     },
-    [dispatch, isWithinPanel, state.disclosureState, state.buttonId]
+    [dispatch, isWithinPanel, state.disclosureState, state.buttonId, ownerDocument]
   )
 
   let handleKeyUp = useCallback((event: ReactKeyboardEvent<HTMLButtonElement>) => {
@@ -284,12 +299,12 @@ let Button = forwardRefWithAs(function Button<TTag extends ElementType = typeof
 
       if (isWithinPanel) {
         dispatch({ type: ActionTypes.ToggleDisclosure })
-        document.getElementById(state.buttonId)?.focus()
+        ownerDocument?.getElementById(state.buttonId)?.focus()
       } else {
         dispatch({ type: ActionTypes.ToggleDisclosure })
       }
     },
-    [dispatch, props.disabled, state.buttonId, isWithinPanel]
+    [dispatch, props.disabled, state.buttonId, isWithinPanel, ownerDocument]
   )
 
   let slot = useMemo<ButtonRenderPropArg>(

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

@@ -33,6 +33,8 @@ import { isFocusableElement, FocusableMode, sortByDomNode } from '../../utils/fo
 import { useOpenClosed, State, OpenClosedProvider } 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 ListboxStates {
   Open,
@@ -530,7 +532,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])
@@ -675,7 +677,8 @@ 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 ownerDocument = useOwnerDocument(internalOptionRef)
   let optionRef = useSyncRefs(ref, internalOptionRef)
 
   let bag = useRef<ListboxOptionDataRef['current']>({ disabled, value, domRef: internalOptionRef })
@@ -687,8 +690,8 @@ let Option = forwardRefWithAs(function Option<
     bag.current.value = value
   }, [bag, value])
   useIsoMorphicEffect(() => {
-    bag.current.textValue = document.getElementById(id)?.textContent?.toLowerCase()
-  }, [bag, id])
+    bag.current.textValue = ownerDocument?.getElementById(id)?.textContent?.toLowerCase()
+  }, [bag, id, ownerDocument])
 
   let select = useCallback(() => state.propsRef.current.onChange(value), [state.propsRef, value])
 
@@ -701,19 +704,19 @@ let Option = forwardRefWithAs(function Option<
     if (state.listboxState !== ListboxStates.Open) return
     if (!selected) return
     dispatch({ type: ActionTypes.GoToOption, focus: Focus.Specific, id })
-    document.getElementById(id)?.focus?.()
-  }, [state.listboxState])
+    ownerDocument?.getElementById(id)?.focus?.()
+  }, [state.listboxState, ownerDocument])
 
   useIsoMorphicEffect(() => {
     if (state.listboxState !== ListboxStates.Open) return
     if (!active) return
     if (state.activationTrigger === ActivationTrigger.Pointer) return
     let d = disposables()
     d.requestAnimationFrame(() => {
-      document.getElementById(id)?.scrollIntoView?.({ block: 'nearest' })
+      ownerDocument?.getElementById(id)?.scrollIntoView?.({ block: 'nearest' })
     })
     return d.dispose
-  }, [id, active, state.listboxState, state.activationTrigger, /* We also want to trigger this when the position of the active item changes so that we can re-trigger the scrollIntoView */ state.activeOptionIndex])
+  }, [id, active, state.listboxState, ownerDocument, state.activationTrigger, /* We also want to trigger this when the position of the active item changes so that we can re-trigger the scrollIntoView */ state.activeOptionIndex])
 
   let handleClick = useCallback(
     (event: { preventDefault: Function }) => {

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

@@ -35,6 +35,8 @@ 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 { getOwnerDocument } from '../../utils/owner'
+import { useOwnerDocument } from 'hooks/use-owner'
 
 enum MenuStates {
   Open,
@@ -398,6 +400,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 +418,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,
@@ -454,7 +457,7 @@ let Items = forwardRefWithAs(function Items<TTag extends ElementType = typeof DE
           dispatch({ type: ActionTypes.CloseMenu })
           if (state.activeItemIndex !== null) {
             let { id } = state.items[state.activeItemIndex]
-            document.getElementById(id)?.click()
+            ownerDocument?.getElementById(id)?.click()
           }
           disposables().nextFrame(() => state.buttonRef.current?.focus({ preventScroll: true }))
           break
@@ -501,7 +504,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,19 +574,20 @@ 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)
+  let ownerDocument = getOwnerDocument(internalItemRef.current)
 
   useIsoMorphicEffect(() => {
     if (state.menuState !== MenuStates.Open) return
     if (!active) return
     if (state.activationTrigger === ActivationTrigger.Pointer) return
     let d = disposables()
     d.requestAnimationFrame(() => {
-      document.getElementById(id)?.scrollIntoView?.({ block: 'nearest' })
+      ownerDocument.getElementById(id)?.scrollIntoView?.({ block: 'nearest' })
     })
     return d.dispose
-  }, [id, active, state.menuState, state.activationTrigger, /* We also want to trigger this when the position of the active item changes so that we can re-trigger the scrollIntoView */ state.activeItemIndex])
+  }, [id, active, state.menuState, ownerDocument, state.activationTrigger, /* We also want to trigger this when the position of the active item changes so that we can re-trigger the scrollIntoView */ state.activeItemIndex])
 
   let bag = useRef<MenuItemDataRef['current']>({ disabled, domRef: internalItemRef })
 
@@ -592,8 +596,8 @@ let Item = forwardRefWithAs(function Item<TTag extends ElementType = typeof DEFA
   }, [bag, disabled])
 
   useIsoMorphicEffect(() => {
-    bag.current.textValue = document.getElementById(id)?.textContent?.toLowerCase()
-  }, [bag, id])
+    bag.current.textValue = ownerDocument.getElementById(id)?.textContent?.toLowerCase()
+  }, [bag, id, ownerDocument])
 
   useIsoMorphicEffect(() => {
     dispatch({ type: ActionTypes.RegisterItem, id, dataRef: bag })