From 08790200c17ef700ba68988803f635ac72a8934b Mon Sep 17 00:00:00 2001
From: Bobby Holley <bobbyholley@gmail.com>
Date: Mon, 20 Jul 2015 14:46:58 -0700
Subject: [PATCH] Bug 1161831 - Associate extension URIs with the appropriate
 addon ID. r=billm,sr=bz

---
 caps/nsIAddonPolicyService.idl             |  7 +++++-
 caps/nsScriptSecurityManager.cpp           | 20 ++++++++++++++++
 caps/nsScriptSecurityManager.h             |  3 +++
 toolkit/components/utils/simpleServices.js | 28 ++++++++++++++++++++++
 4 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/caps/nsIAddonPolicyService.idl b/caps/nsIAddonPolicyService.idl
index caad17cb32c..7a9c1f66215 100644
--- a/caps/nsIAddonPolicyService.idl
+++ b/caps/nsIAddonPolicyService.idl
@@ -11,7 +11,7 @@
  * This interface allows the security manager to query custom per-addon security
  * policy.
  */
-[scriptable,uuid(3ec203f8-2bd0-4f4c-8f99-f9f056221231)]
+[scriptable,uuid(8a034ef9-9d14-4c5d-8319-06c1ab574baa)]
 interface nsIAddonPolicyService : nsISupports
 {
   /**
@@ -24,4 +24,9 @@ interface nsIAddonPolicyService : nsISupports
    * Returns true if a given extension:// URI is web-accessible.
    */
   boolean extensionURILoadableByAnyone(in nsIURI aURI);
+
+  /**
+   * Maps an extension URI to the ID of the addon it belongs to.
+   */
+  AString extensionURIToAddonId(in nsIURI aURI);
 };
diff --git a/caps/nsScriptSecurityManager.cpp b/caps/nsScriptSecurityManager.cpp
index 7ab202c36aa..407f2c864ca 100644
--- a/caps/nsScriptSecurityManager.cpp
+++ b/caps/nsScriptSecurityManager.cpp
@@ -360,6 +360,20 @@ nsScriptSecurityManager::GetChannelResultPrincipal(nsIChannel* aChannel,
     return GetChannelURIPrincipal(aChannel, aPrincipal);
 }
 
+nsresult
+nsScriptSecurityManager::MaybeSetAddonIdFromURI(OriginAttributes& aAttrs, nsIURI* aURI)
+{
+  nsAutoCString scheme;
+  nsresult rv = aURI->GetScheme(scheme);
+  NS_ENSURE_SUCCESS(rv, rv);
+  if (scheme.EqualsLiteral("moz-extension") && GetAddonPolicyService()) {
+    rv = GetAddonPolicyService()->ExtensionURIToAddonId(aURI, aAttrs.mAddonId);
+    NS_ENSURE_SUCCESS(rv, rv);
+  }
+
+  return NS_OK;
+}
+
 /* The principal of the URI that this channel is loading. This is never
  * affected by things like sandboxed loads, or loads where we forcefully
  * inherit the principal.  Think of this as the principal of the server
@@ -391,6 +405,8 @@ nsScriptSecurityManager::GetChannelURIPrincipal(nsIChannel* aChannel,
     }
 
     OriginAttributes attrs(UNKNOWN_APP_ID, false);
+    rv = MaybeSetAddonIdFromURI(attrs, uri);
+    NS_ENSURE_SUCCESS(rv, rv);
     nsCOMPtr<nsIPrincipal> prin = BasePrincipal::CreateCodebasePrincipal(uri, attrs);
     prin.forget(aPrincipal);
     return *aPrincipal ? NS_OK : NS_ERROR_FAILURE;
@@ -1097,6 +1113,8 @@ nsScriptSecurityManager::
   OriginAttributes attrs;
   aLoadContext->GetAppId(&attrs.mAppId);
   aLoadContext->GetIsInBrowserElement(&attrs.mInBrowser);
+  nsresult rv = MaybeSetAddonIdFromURI(attrs, aURI);
+  NS_ENSURE_SUCCESS(rv, rv);
   nsCOMPtr<nsIPrincipal> prin = BasePrincipal::CreateCodebasePrincipal(aURI, attrs);
   prin.forget(aPrincipal);
   return *aPrincipal ? NS_OK : NS_ERROR_FAILURE;
@@ -1109,6 +1127,8 @@ nsScriptSecurityManager::GetDocShellCodebasePrincipal(nsIURI* aURI,
 {
   // XXXbholley - Make this more general in bug 1165466.
   OriginAttributes attrs(aDocShell->GetAppId(), aDocShell->GetIsInBrowserElement());
+  nsresult rv = MaybeSetAddonIdFromURI(attrs, aURI);
+  NS_ENSURE_SUCCESS(rv, rv);
   nsCOMPtr<nsIPrincipal> prin = BasePrincipal::CreateCodebasePrincipal(aURI, attrs);
   prin.forget(aPrincipal);
   return *aPrincipal ? NS_OK : NS_ERROR_FAILURE;
diff --git a/caps/nsScriptSecurityManager.h b/caps/nsScriptSecurityManager.h
index 50b9a8a410b..15bf69f5694 100644
--- a/caps/nsScriptSecurityManager.h
+++ b/caps/nsScriptSecurityManager.h
@@ -120,6 +120,9 @@ private:
     inline void
     AddSitesToFileURIWhitelist(const nsCString& aSiteList);
 
+    // If aURI is a moz-extension:// URI, set mAddonId to the associated addon.
+    nsresult MaybeSetAddonIdFromURI(mozilla::OriginAttributes& aAttrs, nsIURI* aURI);
+
     nsCOMPtr<nsIPrincipal> mSystemPrincipal;
     bool mPrefInitialized;
     bool mIsJavaScriptEnabled;
diff --git a/toolkit/components/utils/simpleServices.js b/toolkit/components/utils/simpleServices.js
index 8777da4b7f5..69b727d9ece 100644
--- a/toolkit/components/utils/simpleServices.js
+++ b/toolkit/components/utils/simpleServices.js
@@ -81,6 +81,23 @@ AddonPolicyService.prototype = {
     return cb ? cb(aURI) : false;
   },
 
+  /*
+   * Maps an extension URI to an addon ID.
+   *
+   * @see nsIAddonPolicyService.extensionURIToAddonId
+   */
+  extensionURIToAddonId(aURI) {
+    if (aURI.scheme != "moz-extension") {
+      throw new TypeError("non-extension URI passed");
+    }
+
+    let cb = this.extensionURIToAddonIdCallback;
+    if (!cb) {
+      throw new Error("no callback set to map extension URIs to addon Ids");
+    }
+    return cb(aURI);
+  },
+
   /*
    * Sets the callbacks used in addonMayLoadURI above. Not accessible over
    * XPCOM - callers should use .wrappedJSObject on the service to call it
@@ -99,6 +116,17 @@ AddonPolicyService.prototype = {
     var old = this.extensionURILoadCallback;
     this.extensionURILoadCallback = aCallback;
     return old;
+  },
+
+  /*
+   * Sets the callback used in extensionURIToAddonId above. Not accessible over
+   * XPCOM - callers should use .wrappedJSObject on the service to call it
+   * directly.
+   */
+  setExtensionURIToAddonIdCallback(aCallback) {
+    var old = this.extensionURIToAddonIdCallback;
+    this.extensionURIToAddonIdCallback = aCallback;
+    return old;
   }
 };