fix Portal cleanup issue (#266)

When the last portal is unmounted we will
document.body.removeChild(target), but this crashes when something
tampered with it. This is reproduceable in the tests. Instead we now
ensure that the portal root always exists when mounting. We will also
ensure to target.parentElement.removeChild which will ensure that we can
reference the parent correctly.

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

@@ -6,7 +6,7 @@ import { Portal } from './portal'
 import { click } from '../../test-utils/interactions'
 
 function getPortalRoot() {
-  return document.getElementById('headlessui-portal-root')
+  return document.getElementById('headlessui-portal-root')!
 }
 
 beforeEach(() => {
@@ -140,3 +140,136 @@ it('should cleanup the Portal root when the last Portal is unmounted', async ()
   expect(getPortalRoot()).not.toBe(null)
   expect(getPortalRoot().childNodes).toHaveLength(1)
 })
+
+it('should be possible to render multiple portals at the same time', async () => {
+  expect(getPortalRoot()).toBe(null)
+
+  function Example() {
+    let [renderA, setRenderA] = useState(true)
+    let [renderB, setRenderB] = useState(true)
+    let [renderC, setRenderC] = useState(true)
+
+    return (
+      <main id="parent">
+        <button id="a" onClick={() => setRenderA(v => !v)}>
+          Toggle A
+        </button>
+        <button id="b" onClick={() => setRenderB(v => !v)}>
+          Toggle B
+        </button>
+        <button id="c" onClick={() => setRenderC(v => !v)}>
+          Toggle C
+        </button>
+
+        <button
+          id="double"
+          onClick={() => {
+            setRenderA(v => !v)
+            setRenderB(v => !v)
+          }}
+        >
+          Toggle A & B{' '}
+        </button>
+
+        {renderA && (
+          <Portal>
+            <p id="content1">Contents 1 ...</p>
+          </Portal>
+        )}
+
+        {renderB && (
+          <Portal>
+            <p id="content2">Contents 2 ...</p>
+          </Portal>
+        )}
+
+        {renderC && (
+          <Portal>
+            <p id="content3">Contents 3 ...</p>
+          </Portal>
+        )}
+      </main>
+    )
+  }
+
+  render(<Example />)
+
+  expect(getPortalRoot()).not.toBe(null)
+  expect(getPortalRoot().childNodes).toHaveLength(3)
+
+  // Remove Portal 1
+  await click(document.getElementById('a'))
+  expect(getPortalRoot().childNodes).toHaveLength(2)
+
+  // Remove Portal 2
+  await click(document.getElementById('b'))
+  expect(getPortalRoot().childNodes).toHaveLength(1)
+
+  // Re-add Portal 1
+  await click(document.getElementById('a'))
+  expect(getPortalRoot().childNodes).toHaveLength(2)
+
+  // Remove Portal 3
+  await click(document.getElementById('c'))
+  expect(getPortalRoot().childNodes).toHaveLength(1)
+
+  // Remove Portal 1
+  await click(document.getElementById('a'))
+  expect(getPortalRoot()).toBe(null)
+
+  // Render A and B at the same time!
+  await click(document.getElementById('double'))
+  expect(getPortalRoot().childNodes).toHaveLength(2)
+})
+
+it('should be possible to tamper with the modal root and restore correctly', async () => {
+  expect(getPortalRoot()).toBe(null)
+
+  function Example() {
+    let [renderA, setRenderA] = useState(true)
+    let [renderB, setRenderB] = useState(true)
+
+    return (
+      <main id="parent">
+        <button id="a" onClick={() => setRenderA(v => !v)}>
+          Toggle A
+        </button>
+        <button id="b" onClick={() => setRenderB(v => !v)}>
+          Toggle B
+        </button>
+
+        {renderA && (
+          <Portal>
+            <p id="content1">Contents 1 ...</p>
+          </Portal>
+        )}
+
+        {renderB && (
+          <Portal>
+            <p id="content2">Contents 2 ...</p>
+          </Portal>
+        )}
+      </main>
+    )
+  }
+
+  render(<Example />)
+
+  expect(getPortalRoot()).not.toBe(null)
+
+  // Tamper tamper
+  document.body.removeChild(document.getElementById('headlessui-portal-root')!)
+
+  // Hide Portal 1 and 2
+  await click(document.getElementById('a'))
+  await click(document.getElementById('b'))
+
+  expect(getPortalRoot()).toBe(null)
+
+  // Re-show Portal 1 and 2
+  await click(document.getElementById('a'))
+  await click(document.getElementById('b'))
+
+  expect(getPortalRoot()).not.toBe(null)
+  expect(getPortalRoot().childNodes).toHaveLength(2)
+})

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

@@ -13,6 +13,16 @@ import { render } from '../../utils/render'
 import { useIsoMorphicEffect } from '../../hooks/use-iso-morphic-effect'
 import { StackProvider, useElemenStack } from '../../internal/stack-context'
 
+function resolvePortalRoot() {
+  if (typeof window === 'undefined') return null
+  let existingRoot = document.getElementById('headlessui-portal-root')
+  if (existingRoot) return existingRoot
+
+  let root = document.createElement('div')
+  root.setAttribute('id', 'headlessui-portal-root')
+  return document.body.appendChild(root)
+}
+
 // ---
 
 let DEFAULT_PORTAL_TAG = Fragment
@@ -21,15 +31,7 @@ interface PortalRenderPropArg {}
 export function Portal<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
   props: Props<TTag, PortalRenderPropArg>
 ) {
-  let [target] = useState(() => {
-    if (typeof window === 'undefined') return null
-    let existingRoot = document.getElementById('headlessui-portal-root')
-    if (existingRoot) return existingRoot
-
-    let root = document.createElement('div')
-    root.setAttribute('id', 'headlessui-portal-root')
-    return document.body.appendChild(root)
-  })
+  let [target, setTarget] = useState(resolvePortalRoot)
   let [element] = useState<HTMLDivElement | null>(() =>
     typeof window === 'undefined' ? null : document.createElement('div')
   )
@@ -37,7 +39,7 @@ export function Portal<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
   useElemenStack(element)
 
   useIsoMorphicEffect(() => {
-    if (!target) return
+    if (!target) return setTarget(resolvePortalRoot())
     if (!element) return
 
     target.appendChild(element)
@@ -47,7 +49,7 @@ export function Portal<TTag extends ElementType = typeof DEFAULT_PORTAL_TAG>(
       if (!element) return
 
       target.removeChild(element)
-      if (target.childNodes.length <= 0) document.body.removeChild(target)
+      if (target.childNodes.length <= 0) target.parentElement?.removeChild(target)
     }
   }, [target, element])