Use ownerDocument instead of document (#1158)

* use `ownerDocument` instead of `document`

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

* update changelog

CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Ignore "outside click" on removed elements ([#1193](https://github.com/tailwindlabs/headlessui/pull/1193))
 - Remove `focus()` from Listbox Option ([#1218](https://github.com/tailwindlabs/headlessui/pull/1218))
 - Improve some internal code ([#1221](https://github.com/tailwindlabs/headlessui/pull/1221))
+- Use `ownerDocument` instead of `document` ([#1158](https://github.com/tailwindlabs/headlessui/pull/1158))
 
 ### Added
 
@@ -49,6 +50,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Remove `focus()` from Listbox Option ([#1218](https://github.com/tailwindlabs/headlessui/pull/1218))
 - Improve some internal code ([#1221](https://github.com/tailwindlabs/headlessui/pull/1221))
 - Don’t drop initial character when searching in Combobox ([#1223](https://github.com/tailwindlabs/headlessui/pull/1223))
+- Use `ownerDocument` instead of `document` ([#1158](https://github.com/tailwindlabs/headlessui/pull/1158))
 
 ### Added
 

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

@@ -29,11 +29,13 @@ import { useInertOthers } from '../../hooks/use-inert-others'
 import { Portal } from '../../components/portal/portal'
 import { ForcePortalRoot } from '../../internal/portal-force-root'
 import { Description, useDescriptions } from '../description/description'
-import { useWindowEvent } from '../../hooks/use-window-event'
 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'
+import { useOwnerDocument } from '../../hooks/use-owner'
+import { useEventListener } from '../../hooks/use-event-listener'
 
 enum DialogStates {
   Open,
@@ -133,6 +135,8 @@ let DialogRoot = forwardRefWithAs(function Dialog<
   let internalDialogRef = useRef<HTMLDivElement | null>(null)
   let dialogRef = useSyncRefs(internalDialogRef, ref)
 
+  let ownerDocument = useOwnerDocument(internalDialogRef)
+
   // Validations
   let hasOpen = props.hasOwnProperty('open') || usesOpenClosedState !== null
   let hasOnClose = props.hasOwnProperty('onClose')
@@ -217,7 +221,7 @@ let DialogRoot = forwardRefWithAs(function Dialog<
   })
 
   // Handle `Escape` to close
-  useWindowEvent('keydown', (event) => {
+  useEventListener(ownerDocument?.defaultView, 'keydown', (event) => {
     if (event.key !== Keys.Escape) return
     if (dialogState !== DialogStates.Open) return
     if (hasNestedDialogs) return
@@ -231,17 +235,23 @@ 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)
+    if (!ownerDocument) return
+
+    let documentElement = ownerDocument.documentElement
+    let ownerWindow = ownerDocument.defaultView ?? window
+
+    let overflow = documentElement.style.overflow
+    let paddingRight = documentElement.style.paddingRight
 
-    let scrollbarWidth = window.innerWidth - document.documentElement.clientWidth
+    let scrollbarWidth = ownerWindow.innerWidth - documentElement.clientWidth
 
-    document.documentElement.style.overflow = 'hidden'
-    document.documentElement.style.paddingRight = `${scrollbarWidth}px`
+    documentElement.style.overflow = 'hidden'
+    documentElement.style.paddingRight = `${scrollbarWidth}px`
 
     return () => {
-      document.documentElement.style.overflow = overflow
-      document.documentElement.style.paddingRight = paddingRight
+      documentElement.style.overflow = overflow
+      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,15 @@ let DisclosureRoot = forwardRefWithAs(function Disclosure<
   let close = useCallback(
     (focusableElement?: HTMLElement | MutableRefObject<HTMLElement | null>) => {
       dispatch({ type: ActionTypes.CloseDisclosure })
+      let ownerDocument = getOwnerDocument(internalDisclosureRef)
+      if (!ownerDocument) return
 
       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>) => {