2 files changed, 33 insertions, 36 deletions

diff --git a/browser/base/content/browser-addons.js b/browser/base/content/browser-addons.js
index 1f81d1fb0..1d881536a 100644
--- a/browser/base/content/browser-addons.js
+++ b/browser/base/content/browser-addons.js
@@ -459,15 +459,15 @@ var LightWeightThemeWebInstaller = {
 
     switch (message.name) {
       case "LightWeightThemeWebInstaller:Install": {
-        this._installRequest(data.themeData, data.baseURI);
+        this._installRequest(data.themeData, data.principal, data.baseURI);
         break;
       }
       case "LightWeightThemeWebInstaller:Preview": {
-        this._preview(data.themeData, data.baseURI);
+        this._preview(data.themeData, data.principal, data.baseURI);
         break;
       }
       case "LightWeightThemeWebInstaller:ResetPreview": {
-        this._resetPreview(data && data.baseURI);
+        this._resetPreview(data && data.principal);
         break;
       }
     }
@@ -489,33 +489,24 @@ var LightWeightThemeWebInstaller = {
     return this._manager = temp.LightweightThemeManager;
   },
 
-  _installRequest: function (dataString, baseURI) {
+  _installRequest(dataString, principal, baseURI) {
+    // Don't allow installing off null principals.
+    if (!principal.URI) {
+      return;
+    }
+
     let data = this._manager.parseTheme(dataString, baseURI);
 
     if (!data) {
       return;
     }
 
-    let uri = makeURI(baseURI);
-
     // A notification bar with the option to undo is normally shown after a
     // theme is installed.  But the discovery pane served from the url(s)
     // below has its own toggle switch for quick undos, so don't show the
     // notification in that case.
-    let notify = uri.prePath != "https://discovery.addons.mozilla.org";
-    if (notify) {
-      try {
-        if (Services.prefs.getBoolPref("extensions.webapi.testing")
-            && (uri.prePath == "https://discovery.addons.allizom.org"
-                || uri.prePath == "https://discovery.addons-dev.allizom.org")) {
-          notify = false;
-        }
-      } catch (e) {
-        // getBoolPref() throws if the testing pref isn't set.  ignore it.
-      }
-    }
-
-    if (this._isAllowed(baseURI)) {
+    let notify = this._shouldShowUndoPrompt(principal);
+    if (this._isAllowed(principal)) {
       this._install(data, notify);
       return;
     }
@@ -526,7 +517,7 @@ var LightWeightThemeWebInstaller = {
       gNavigatorBundle.getString("lwthemeInstallRequest.allowButton.accesskey");
     let message =
       gNavigatorBundle.getFormattedString("lwthemeInstallRequest.message",
-                                          [uri.host]);
+                                          [principal.URI.host]);
     let buttons = [{
       label: allowButtonText,
       accessKey: allowButtonAccesskey,
@@ -629,8 +620,8 @@ var LightWeightThemeWebInstaller = {
       });
   },
 
-  _preview: function (dataString, baseURI) {
-    if (!this._isAllowed(baseURI))
+  _preview(dataString, principal, baseURI) {
+    if (!this._isAllowed(principal))
       return;
 
     let data = this._manager.parseTheme(dataString, baseURI);
@@ -642,29 +633,33 @@ var LightWeightThemeWebInstaller = {
     this._manager.previewTheme(data);
   },
 
-  _resetPreview: function (baseURI) {
-    if (baseURI && !this._isAllowed(baseURI))
+  _resetPreview(principal) {
+    if (!this._isAllowed(principal))
       return;
     gBrowser.tabContainer.removeEventListener("TabSelect", this, false);
     this._manager.resetPreview();
   },
 
-  _isAllowed: function (srcURIString) {
-    let uri;
-    try {
-      uri = makeURI(srcURIString);
-    }
-    catch (e) {
-      // makeURI fails if srcURIString is a nonsense URI
+  _isAllowed(principal) {
+    if (!principal || !principal.URI || !principal.URI.schemeIs("https")) {
       return false;
     }
 
-    if (!uri.schemeIs("https")) {
+    let pm = Services.perms;
+    return pm.testPermission(principal.URI, "install") == pm.ALLOW_ACTION;
+  },
+
+  _shouldShowUndoPrompt(principal) {
+    if (!principal || !principal.URI) {
+      return true;
+    }
+
+    let prePath = principal.URI.prePath;
+    if (prePath == "https://addons.palemoon.org") {
       return false;
     }
 
-    let pm = Services.perms;
-    return pm.testPermission(uri, "install") == pm.ALLOW_ACTION;
+    return true;
   }
 };
 
diff --git a/browser/base/content/content.js b/browser/base/content/content.js
index 46e9b45d6..496e0d111 100644
--- a/browser/base/content/content.js
+++ b/browser/base/content/content.js
@@ -896,6 +896,7 @@ var LightWeightThemeWebInstallListener = {
       case "InstallBrowserTheme": {
         sendAsyncMessage("LightWeightThemeWebInstaller:Install", {
           baseURI: event.target.baseURI,
+          principal: event.target.nodePrincipal,
           themeData: event.target.getAttribute("data-browsertheme"),
         });
         break;
@@ -903,6 +904,7 @@ var LightWeightThemeWebInstallListener = {
       case "PreviewBrowserTheme": {
         sendAsyncMessage("LightWeightThemeWebInstaller:Preview", {
           baseURI: event.target.baseURI,
+          principal: event.target.nodePrincipal,
           themeData: event.target.getAttribute("data-browsertheme"),
         });
         this._previewWindow = event.target.ownerGlobal;
@@ -917,7 +919,7 @@ var LightWeightThemeWebInstallListener = {
       case "ResetBrowserThemePreview": {
         if (this._previewWindow) {
           sendAsyncMessage("LightWeightThemeWebInstaller:ResetPreview",
-                           {baseURI: event.target.baseURI});
+                           {principal: event.target.nodePrincipal});
           this._resetPreviewWindow();
         }
         break;