From 273c3c56903a84d05a51a6f39b62e7fdd760342b Mon Sep 17 00:00:00 2001
From: Will Bamberg <wbamberg@mozilla.com>
Date: Thu, 6 Jul 2017 19:31:10 -0700
Subject: [PATCH] Improve content script injection

---
 beastify/content_scripts/beastify.js |  2 --
 beastify/popup/choose_beast.js       | 43 +++++++++++++++++++++-------
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/beastify/content_scripts/beastify.js b/beastify/content_scripts/beastify.js
index 88f3f71..202d84b 100644
--- a/beastify/content_scripts/beastify.js
+++ b/beastify/content_scripts/beastify.js
@@ -2,12 +2,10 @@
 beastify():
 * removes every node in the document.body,
 * then inserts the chosen beast
-* then removes itself as a listener 
 */
 function beastify(request, sender, sendResponse) {
   removeEverything();
   insertBeast(request.beastURL);
-  browser.runtime.onMessage.removeListener(beastify);
 }
 
 /*
diff --git a/beastify/popup/choose_beast.js b/beastify/popup/choose_beast.js
index d70225c..cbb2306 100644
--- a/beastify/popup/choose_beast.js
+++ b/beastify/popup/choose_beast.js
@@ -16,33 +16,56 @@ function beastNameToURL(beastName) {
 Listen for clicks in the popup.
 
 If the click is on one of the beasts:
-  Inject the "beastify.js" content script in the active tab.
+  (1) inject the test script into the active tab.
+  This checks whether the tab's scope includes a function called "beastify()".
 
-  Then get the active tab and send "beastify.js" a message
-  containing the URL to the chosen beast's image.
+  (2) if the tab doesn't already have a function called "beastify()",
+  then inject the beastify.js content script. Otherwise, do nothing.
+  This should ensure that beastify.js is only injected once, even if
+  the user clicks a beast again.
+
+  (3) send a message to "beastify.js" that contains the URL to
+  the chosen beast's image.
 
 If it's on a button wich contains class "clear":
   Reload the page.
   Close the popup. This is needed, as the content script malfunctions after page reloads.
 */
 document.addEventListener("click", (e) => {
-  if (e.target.classList.contains("beast")) {
+
+  function injectTestScript() {
+    return browser.tabs.executeScript({
+      code: "typeof beastify === 'function';",
+    });
+  }
+
+  function injectBeastify(testResults) {
+    if (!testResults || testResults[0] !== true) {
+      return browser.tabs.executeScript({
+        file: "/content_scripts/beastify.js"
+      });
+    } else {
+      return Promise.resolve(true);
+    }
+  }
+
+  function messageBeastify() {
     var chosenBeast = e.target.textContent;
     var chosenBeastURL = beastNameToURL(chosenBeast);
-
-    browser.tabs.executeScript(null, {
-      file: "/content_scripts/beastify.js"
-    });
-
     var gettingActiveTab = browser.tabs.query({active: true, currentWindow: true});
     gettingActiveTab.then((tabs) => {
       browser.tabs.sendMessage(tabs[0].id, {beastURL: chosenBeastURL});
     });
   }
+
+  if (e.target.classList.contains("beast")) {
+    injectTestScript()
+      .then(injectBeastify)
+      .then(messageBeastify);
+  }
   else if (e.target.classList.contains("clear")) {
     browser.tabs.reload();
     window.close();
-
     return;
   }
 });