From e3cbcc43f982ad3eb0775f1b3cf2425df13a0456 Mon Sep 17 00:00:00 2001
From: Robin Malfait <malfait.robin@gmail.com>
Date: Mon, 1 Mar 2021 13:13:22 +0100
Subject: [PATCH] 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.
---
 .../src/components/portal/portal.test.tsx     | 135 +++++++++++++++++-
 .../src/components/portal/portal.tsx          |  24 ++--
 2 files changed, 147 insertions(+), 12 deletions(-)

diff --git a/packages/@headlessui-react/src/components/portal/portal.test.tsx b/packages/@headlessui-react/src/components/portal/portal.test.tsx
index e504bd0..db4c9c6 100644
--- a/packages/@headlessui-react/src/components/portal/portal.test.tsx
+++ b/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)
+})
diff --git a/packages/@headlessui-react/src/components/portal/portal.tsx b/packages/@headlessui-react/src/components/portal/portal.tsx
index ede249a..8b4aa96 100644
--- a/packages/@headlessui-react/src/components/portal/portal.tsx
+++ b/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])