From 22ca6cb21dbcbfe4be8ceb31c3a885f1bd876497 Mon Sep 17 00:00:00 2001
From: baku <baku@mozilla.com>
Date: Tue, 21 Feb 2017 22:49:10 +0100
Subject: [PATCH] Fix for getIdentityService - Issue 227

---
 index.js | 38 ++++++++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/index.js b/index.js
index 87442e7..f32506d 100644
--- a/index.js
+++ b/index.js
@@ -112,7 +112,7 @@ const ContainerService = {
     // need to implement a migration in the future.
     if (installation) {
       let preInstalledIdentities = []; // eslint-disable-line prefer-const
-      ContextualIdentityService.getIdentities().forEach(identity => {
+      ContextualIdentityProxy.getIdentities().forEach(identity => {
         preInstalledIdentities.push(identity.userContextId);
       });
 
@@ -131,12 +131,12 @@ const ContainerService = {
 
       // Maybe rename the Banking container.
       if (prefService.get("privacy.userContext.enabled") !== true) {
-        const identity = ContextualIdentityService.getIdentityFromId(3);
+        const identity = ContextualIdentityProxy.getIdentityFromId(3);
         if (identity && identity.l10nID === "userContextBanking.label") {
-          ContextualIdentityService.update(identity.userContextId,
-                                           "Finance",
-                                           identity.icon,
-                                           identity.color);
+          ContextualIdentityProxy.update(identity.userContextId,
+                                         "Finance",
+                                         identity.icon,
+                                         identity.color);
         }
       }
     }
@@ -238,7 +238,7 @@ const ContainerService = {
     }).sendEvent;
 
     // Begin-Of-Hack
-    function workaroundForCookieManager(method, userContextId) {
+    ContextualIdentityService.workaroundForCookieManager = function(method, userContextId) {
       let identity = method.call(ContextualIdentityService, userContextId);
       if (!identity && userContextId) {
         identity = {
@@ -253,15 +253,15 @@ const ContainerService = {
       return identity;
     }
 
-    const oldGetIdentityFromId = ContextualIdentityService.getIdentityFromId;
+    this._oldGetIdentityFromId = ContextualIdentityService.getIdentityFromId;
     ContextualIdentityService.getIdentityFromId = function(userContextId) {
-      return workaroundForCookieManager(oldGetIdentityFromId, userContextId);
+      return this.workaroundForCookieManager(ContainerService._oldGetIdentityFromId, userContextId);
     };
 
     if ("getPublicIdentityFromId" in ContextualIdentityService) {
-      const oldGetPublicIdentityFromId = ContextualIdentityService.getIdentityFromId;
-      ContextualIdentityService.getIdentityFromId = function(userContextId) {
-        return workaroundForCookieManager(oldGetPublicIdentityFromId, userContextId);
+      this._oldGetPublicIdentityFromId = ContextualIdentityService.getPublicIdentityFromId;
+      ContextualIdentityService.getPublicIdentityFromId = function(userContextId) {
+        return this.workaroundForCookieManager(ContainerService._oldGetPublicIdentityFromId, userContextId);
       };
     }
     // End-Of-Hack
@@ -987,14 +987,24 @@ const ContainerService = {
     this._forgetIdentity();
 
     const preInstalledIdentities = data.preInstalledIdentities;
-    ContextualIdentityService.getIdentities().forEach(identity => {
+    ContextualIdentityProxy.getIdentities().forEach(identity => {
       if (!preInstalledIdentities.includes(identity.userContextId)) {
-        ContextualIdentityService.remove(identity.userContextId);
+        ContextualIdentityProxy.remove(identity.userContextId);
       }
     });
 
     // Let's delete the configuration.
     delete ss.storage.savedConfiguration;
+
+    // Begin-Of-Hack
+    if (this._oldGetIdentityFromId) {
+      ContextualIdentityService.getIdentityFromId = this._oldGetIdentityFromId;
+    }
+
+    if (this._oldGetPublicIdentityFromId) {
+      ContextualIdentityService.getPublicIdentityFromId = this._oldGetPublicIdentityFromId;
+    }
+    // End-Of-Hack
   },
 
   _forgetIdentity(userContextId = 0) {