code review re. sandbox shutdown

2 files changed, 49 insertions, 37 deletions

diff --git a/platform/firefox/frameModule.js b/platform/firefox/frameModule.js
index 07e440e..79c3c33 100644
--- a/platform/firefox/frameModule.js
+++ b/platform/firefox/frameModule.js
@@ -211,11 +211,12 @@ const contentObserver = {
         return this.ACCEPT;
     },
 
-    initContentScripts: function(win, sandbox) {
+    initContentScripts: function(win, create) {
         let messager = getMessageManager(win);
         let sandboxId = hostName + ':sb:' + this.uniqueSandboxId++;
+        let sandbox;
 
-        if ( sandbox ) {
+        if ( create ) {
             let sandboxName = [
                 win.location.href.slice(0, 100),
                 win.document.title.slice(0, 100)
@@ -236,6 +237,23 @@ const contentObserver = {
                     // I've seen this happens, need to investigate why.
                 }
             };
+
+            // The goal is to have content scripts removed from web pages. This
+            // helps remove traces of uBlock from memory when disabling/removing
+            // the addon.
+            // For example, this takes care of:
+            //   https://github.com/gorhill/uBlock/commit/ea4faff383789053f423498c1f1165c403fde7c7#commitcomment-11964137
+            //   > "gets the whole selected tab flashing"
+            sandbox.outerShutdown = function() {
+                sandbox.removeMessageListener();
+                sandbox.addMessageListener =
+                sandbox.injectScript =
+                sandbox.removeMessageListener =
+                sandbox.sendAsyncMessage =
+                sandbox.outerShutdown = function(){};
+                sandbox.vAPI = {};
+                messager = null;
+            };
         }
         else {
             sandbox = win;
@@ -283,22 +301,6 @@ const contentObserver = {
             sandbox._messageListener_ = null;
         };
 
-        // The goal is to have content scripts removed from web pages. This
-        // helps remove traces of uBlock from memory when disabling/removing
-        // the addon.
-        // For example, this takes care of:
-        //   https://github.com/gorhill/uBlock/commit/ea4faff383789053f423498c1f1165c403fde7c7#commitcomment-11964137
-        //   > "gets the whole selected tab flashing"
-        sandbox.shutdownSandbox = function() {
-            sandbox.removeMessageListener();
-            sandbox.addMessageListener =
-            sandbox.injectScript =
-            sandbox.removeMessageListener =
-            sandbox.sendAsyncMessage =
-            sandbox.shutdownSandbox = null;
-            messager = null;
-        };
-
         return sandbox;
     },
 
diff --git a/platform/firefox/vapi-client.js b/platform/firefox/vapi-client.js
index 0bdf617..a69469b 100644
--- a/platform/firefox/vapi-client.js
+++ b/platform/firefox/vapi-client.js
@@ -19,7 +19,7 @@
     Home: https://github.com/gorhill/uBlock
 */
 
-/* global addMessageListener, removeMessageListener, sendAsyncMessage, shutdownSandbox */
+/* global addMessageListener, removeMessageListener, sendAsyncMessage, outerShutdown */
 
 // For non background pages
 
@@ -78,27 +78,24 @@ vAPI.messaging = {
 
     builtinChannelHandler: function(msg) {
         if ( msg.cmd === 'injectScript' ) {
+            // injectScript is not always present.
+            // - See contentObserver.initContentScripts in frameModule.js
+            if ( typeof self.injectScript !== 'function' )  {
+                return;
+            }
             var details = msg.details;
             if ( !details.allFrames && window !== window.top ) {
                 return;
             }
-            // TODO: investigate why this happens, and if this happens
-            // legitimately (content scripts not injected I suspect, so
-            // that would make this legitimate).
-            // Case: open popup UI from icon in uBlock's logger
-            if ( typeof self.injectScript === 'function' )  {
-                self.injectScript(details.file);
-            }
+            self.injectScript(details.file);
             return;
         }
         if ( msg.cmd === 'shutdownSandbox' ) {
             vAPI.shutdown.exec();
-            vAPI.messaging.close();
-            window.removeEventListener('pagehide', vAPI.messaging.toggleListener, true);
-            window.removeEventListener('pageshow', vAPI.messaging.toggleListener, true);
-            vAPI.messaging = null;
-            vAPI = {};
-            shutdownSandbox();
+            vAPI.messaging.shutdown();
+            if ( typeof self.outerShutdown === 'function' ) {
+                outerShutdown();
+            }
             return;
         }
     },
@@ -109,6 +106,13 @@ vAPI.messaging = {
         window.addEventListener('pageshow', this.toggleListener, true);
     },
 
+    shutdown: function() {
+        this.close();
+        window.removeEventListener('pagehide', this.toggleListener, true);
+        window.removeEventListener('pageshow', this.toggleListener, true);
+        vAPI.messaging = null;
+    },
+
     close: function() {
         window.removeEventListener('pagehide', this.toggleListener, true);
         window.removeEventListener('pageshow', this.toggleListener, true);
@@ -146,18 +150,24 @@ vAPI.messaging = {
     },
 
     toggleListener: function({type, persisted}) {
-        if ( !vAPI.messaging ) {
+        var me = vAPI.messaging;
+        if ( !me ) {
             return;
         }
 
         if ( type === 'pagehide' ) {
-            vAPI.messaging.disconnect();
+            if ( !persisted ) {
+                vAPI.shutdown.exec();
+                vAPI.messaging.shutdown();
+                if ( typeof self.outerShutdown === 'function' ) {
+                    outerShutdown();
+                }
+            }
+            me.disconnect();
             return;
         }
 
-        if ( persisted ) {
-            vAPI.messaging.connect();
-        }
+        me.connect();
     }
 };