fix: doesn't unlock the body on unmount

streamich · Apr 13, 2020 · 1ead4ef · 1ead4ef
1 parent 2665eb5
commit 1ead4ef
Showing 1 changed file with 48 additions and 30 deletions.
diff --git a/src/useLockBodyScroll.ts b/src/useLockBodyScroll.ts
@@ -48,47 +48,65 @@ export default !doc
   : function useLockBody(locked: boolean = true, elementRef?: RefObject<HTMLElement>) {
       elementRef = elementRef || useRef(doc!.body);
 
-      useEffect(() => {
-        const body = getClosestBody(elementRef!.current);
-        if (!body) {
-          return;
+      const lock = body => {
+        const bodyInfo = bodies.get(body);
+        if (!bodyInfo) {
+          bodies.set(body, { counter: 1, initialOverflow: body.style.overflow });
+          if (isIosDevice) {
+            if (!documentListenerAdded) {
+              document.addEventListener('touchmove', preventDefault, { passive: false });
+
+              documentListenerAdded = true;
+            }
+          } else {
+            body.style.overflow = 'hidden';
+          }
+        } else {
+          bodies.set(body, { counter: bodyInfo.counter + 1, initialOverflow: bodyInfo.initialOverflow });
         }
+      };
 
+      const unlock = body => {
         const bodyInfo = bodies.get(body);
-
-        if (locked) {
-          if (!bodyInfo) {
-            bodies.set(body, { counter: 1, initialOverflow: body.style.overflow });
+        if (bodyInfo) {
+          if (bodyInfo.counter === 1) {
+            bodies.delete(body);
             if (isIosDevice) {
-              if (!documentListenerAdded) {
-                document.addEventListener('touchmove', preventDefault, { passive: false });
+              body.ontouchmove = null;
 
-                documentListenerAdded = true;
+              if (documentListenerAdded) {
+                document.removeEventListener('touchmove', preventDefault);
+                documentListenerAdded = false;
               }
             } else {
-              body.style.overflow = 'hidden';
+              body.style.overflow = bodyInfo.initialOverflow;
             }
           } else {
-            bodies.set(body, { counter: bodyInfo.counter + 1, initialOverflow: bodyInfo.initialOverflow });
+            bodies.set(body, { counter: bodyInfo.counter - 1, initialOverflow: bodyInfo.initialOverflow });
           }
+        }
+      };
+
+      useEffect(() => {
+        const body = getClosestBody(elementRef!.current);
+        if (!body) {
+          return;
+        }
+        if (locked) {
+          lock(body);
         } else {
-          if (bodyInfo) {
-            if (bodyInfo.counter === 1) {
-              bodies.delete(body);
-              if (isIosDevice) {
-                body.ontouchmove = null;
-
-                if (documentListenerAdded) {
-                  document.removeEventListener('touchmove', preventDefault);
-                  documentListenerAdded = false;
-                }
-              } else {
-                body.style.overflow = bodyInfo.initialOverflow;
-              }
-            } else {
-              bodies.set(body, { counter: bodyInfo.counter - 1, initialOverflow: bodyInfo.initialOverflow });
-            }
-          }
+          unlock(body);
         }
       }, [locked, elementRef.current]);
+
+      // clean up, on un-mount
+      useEffect(() => {
+        const body = getClosestBody(elementRef!.current);
+        if (!body) {
+          return;
+        }
+        return () => {
+          unlock(body);
+        };
+      }, []);
     };