From 1ead4efab6c67131e480570b578b0ce803204953 Mon Sep 17 00:00:00 2001
From: manishsundriyal <msundriyal27@gmail.com>
Date: Mon, 13 Apr 2020 17:46:58 +0530
Subject: [PATCH] fix:  doesn't unlock the body on unmount

---
 src/useLockBodyScroll.ts | 86 ++++++++++++++++++++++++----------------
 1 file changed, 52 insertions(+), 34 deletions(-)

diff --git a/src/useLockBodyScroll.ts b/src/useLockBodyScroll.ts
index b32a92ea..60133297 100644
--- 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);
 
+      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 (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 });
+          }
+        }
+      };
+
       useEffect(() => {
         const body = getClosestBody(elementRef!.current);
         if (!body) {
           return;
         }
-
-        const bodyInfo = bodies.get(body);
-
         if (locked) {
-          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 });
-          }
+          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);
+        };
+      }, []);
     };