From 38df404ab50d435ee311420695a0adca8b259397 Mon Sep 17 00:00:00 2001
From: Luis Orbaiceta <44276180+luisorbaiceta@users.noreply.github.com>
Date: Wed, 16 Mar 2022 14:55:41 +0100
Subject: [PATCH] Prevent click-outside from losing children refs (#626)

fix(site): prevent click-outside to close children ref
* feat: Add forwardRef for compatibility
* fix(site): remove asChild for dropdown Fragment

Co-authored-by: Dom Sip <dom@vercel.com>
---
 site/components/common/UserNav/UserNav.tsx |  2 +-
 site/lib/click-outside/click-outside.tsx   | 91 ++++++++++++++++------
 2 files changed, 67 insertions(+), 26 deletions(-)

diff --git a/site/components/common/UserNav/UserNav.tsx b/site/components/common/UserNav/UserNav.tsx
index 76cd6c53d..436839588 100644
--- a/site/components/common/UserNav/UserNav.tsx
+++ b/site/components/common/UserNav/UserNav.tsx
@@ -69,7 +69,7 @@ const UserNav: React.FC<{
         {process.env.COMMERCE_CUSTOMERAUTH_ENABLED && (
           <li className={s.item}>
             <Dropdown>
-              <DropdownTrigger asChild>
+              <DropdownTrigger>
                 <button
                   aria-label="Menu"
                   className={s.avatarButton}
diff --git a/site/lib/click-outside/click-outside.tsx b/site/lib/click-outside/click-outside.tsx
index 3c6b0ef28..973345b07 100644
--- a/site/lib/click-outside/click-outside.tsx
+++ b/site/lib/click-outside/click-outside.tsx
@@ -1,42 +1,83 @@
-import React, { useRef, useEffect, MouseEvent } from 'react'
+import React, {
+  useRef,
+  useEffect,
+  MouseEvent,
+  FC,
+  ReactElement,
+  forwardRef,
+  Ref,
+} from 'react'
+import mergeRefs from 'react-merge-refs'
 import hasParent from './has-parent'
 
 interface ClickOutsideProps {
   active: boolean
   onClick: (e?: MouseEvent) => void
-  children: any
+  ref?: Ref<any>
 }
 
-const ClickOutside = ({
-  active = true,
-  onClick,
-  children,
-}: ClickOutsideProps) => {
-  const innerRef = useRef()
+/**
+ * Use forward ref to allow this component to be used with other components like
+ * focus-trap-react, that rely on the same type of ref forwarding to direct children
+ */
+const ClickOutside: FC<ClickOutsideProps> = forwardRef(
+  ({ active = true, onClick, children }, forwardedRef) => {
+    const innerRef = useRef()
 
-  const handleClick = (event: any) => {
-    if (!hasParent(event.target, innerRef?.current)) {
-      if (typeof onClick === 'function') {
+    const child = children ? (React.Children.only(children) as any) : undefined
+
+    if (!child || child.type === React.Fragment) {
+      /**
+       * React Fragments can't be used, as it would not be possible to pass the ref
+       * created here to them.
+       */
+      throw new Error('A valid non Fragment React Children should be provided')
+    }
+
+    if (typeof onClick != 'function') {
+      throw new Error('onClick must be a valid function')
+    }
+
+    useEffect(() => {
+      if (active) {
+        document.addEventListener('mousedown', handleClick)
+        document.addEventListener('touchstart', handleClick)
+      }
+      return () => {
+        if (active) {
+          document.removeEventListener('mousedown', handleClick)
+          document.removeEventListener('touchstart', handleClick)
+        }
+      }
+    })
+
+    const handleClick = (event: any) => {
+      /**
+       * Check if the clicked element is contained by the top level tag provided to the
+       * ClickOutside component, if not, Outside clicked! Fire onClick cb
+       */
+      if (!hasParent(event.target, innerRef?.current)) {
         onClick(event)
       }
     }
-  }
 
-  useEffect(() => {
-    if (active) {
-      document.addEventListener('mousedown', handleClick)
-      document.addEventListener('touchstart', handleClick)
-    }
-
-    return () => {
-      if (active) {
-        document.removeEventListener('mousedown', handleClick)
-        document.removeEventListener('touchstart', handleClick)
+    /**
+     * Preserve the child ref prop if exists and merge it with the one used here and the
+     * proxied by the forwardRef method
+     */
+    const composedRefCallback = (element: ReactElement) => {
+      if (typeof child.ref === 'function') {
+        child.ref(element)
+      } else if (child.ref) {
+        child.ref.current = element
       }
     }
-  })
 
-  return React.cloneElement(children, { ref: innerRef })
-}
+    return React.cloneElement(child, {
+      ref: mergeRefs([composedRefCallback, innerRef, forwardedRef]),
+    })
+  }
+)
 
+ClickOutside.displayName = 'ClickOutside'
 export default ClickOutside