Bug 1751010: Set AUMID correctly for private browsing windows r=mossop,mhowell

This is done differently depending on whether the Private Browsing window is the initial one, or if it's a new window from a running application. In the former case, we're still not setting it up early enough -- the window will initially show up in one icon, and then switch to the other one after we've adjusted the AUMID.

Regardless, the fact that we're doing it for every new window appears to mean that other places that require it (such as GeckoChildProcessHost.cpp) will Do The Right Thing (because they'll pull the cached version that `GetAppUserModelID` finds).

This behind a pref for now both because of the issue noted above, as well as the fact that we're waiting on finalized branding for this.

Differential Revision: https://phabricator.services.mozilla.com/D138194

Author: bhearsum

browser/components/BrowserContentHandler.jsm

@@ -26,18 +26,11 @@ XPCOMUtils.defineLazyModuleGetters(this, {
   ShellService: "resource:///modules/ShellService.jsm",
   UpdatePing: "resource://gre/modules/UpdatePing.jsm",
 });
-XPCOMUtils.defineLazyServiceGetter(
-  this,
-  "WindowsUIUtils",
-  "@mozilla.org/windows-ui-utils;1",
-  "nsIWindowsUIUtils"
-);
-XPCOMUtils.defineLazyServiceGetter(
-  this,
-  "UpdateManager",
-  "@mozilla.org/updates/update-manager;1",
-  "nsIUpdateManager"
-);
+XPCOMUtils.defineLazyServiceGetters(this, {
+  UpdateManager: ["@mozilla.org/updates/update-manager;1", "nsIUpdateManager"],
+  WinTaskbar: ["@mozilla.org/windows-taskbar;1", "nsIWinTaskbar"],
+  WindowsUIUtils: ["@mozilla.org/windows-ui-utils;1", "nsIWindowsUIUtils"],
+});
 
 XPCOMUtils.defineLazyGetter(this, "gSystemPrincipal", () =>
   Services.scriptSecurityManager.getSystemPrincipal()
@@ -48,6 +41,8 @@ XPCOMUtils.defineLazyGlobalGetters(this, [URL]);
 const ONCE_DOMAINS = ["mozilla.org", "firefox.com"];
 const ONCE_PREF = "browser.startup.homepage_override.once";
 
+const PRIVACY_SEGMENTATION_PREF = "browser.privacySegmentation.enabled";
+
 function shouldLoadURI(aURI) {
   if (aURI && !aURI.schemeIs("chrome")) {
     return true;
@@ -276,6 +271,13 @@ function openBrowserWindow(
         win.docShell.QueryInterface(
           Ci.nsILoadContext
         ).usePrivateBrowsing = true;
+        if (Services.prefs.getBoolPref(PRIVACY_SEGMENTATION_PREF)) {
+          // TODO: Changing this after the Window has been painted causes it to
+          // change Taskbar icons if the original one had a different AUMID.
+          // This must stay pref'ed off until this is resolved.
+          // https://bugzilla.mozilla.org/show_bug.cgi?id=1751010
+          WinTaskbar.setGroupIdForWindow(win, WinTaskbar.defaultPrivateGroupId);
+        }
       }
 
       let openTime = win.openTime;

widget/nsWidgetInitData.h

@@ -107,7 +107,8 @@ struct nsWidgetInitData {
         mAlwaysOnTop(false),
         mPIPWindow(false),
         mFissionWindow(false),
-        mResizable(false) {}
+        mResizable(false),
+        mIsPrivate(false) {}
 
   nsWindowType mWindowType;
   nsBorderStyle mBorderStyle;
@@ -137,6 +138,7 @@ struct nsWidgetInitData {
   bool mFissionWindow;
   // True if the window is user-resizable.
   bool mResizable;
+  bool mIsPrivate;
 };
 
 #endif  // nsWidgetInitData_h__

widget/windows/nsWindow.cpp

@@ -86,6 +86,8 @@
 #include <unknwn.h>
 #include <psapi.h>
 #include <rpc.h>
+#include <propvarutil.h>
+#include <propkey.h>
 
 #include "mozilla/Logging.h"
 #include "prtime.h"
@@ -968,6 +970,28 @@ nsresult nsWindow::Create(nsIWidget* aParent, nsNativeWidget aNativeParent,
     return NS_ERROR_FAILURE;
   }
 
+  if (aInitData->mIsPrivate) {
+    if (Preferences::GetBool("browser.privacySegmentation.enabled", false)) {
+      RefPtr<IPropertyStore> pPropStore;
+      if (!FAILED(SHGetPropertyStoreForWindow(mWnd, IID_IPropertyStore,
+                                              getter_AddRefs(pPropStore)))) {
+        PROPVARIANT pv;
+        nsAutoString aumid;
+        // make sure we're using the private browsing AUMID so that taskbar
+        // grouping works properly
+        Unused << NS_WARN_IF(
+            !mozilla::widget::WinTaskbar::GenerateAppUserModelID(aumid, true));
+        if (!FAILED(InitPropVariantFromString(aumid.get(), &pv))) {
+          if (!FAILED(pPropStore->SetValue(PKEY_AppUserModel_ID, pv))) {
+            pPropStore->Commit();
+          }
+
+          PropVariantClear(&pv);
+        }
+      }
+    }
+  }
+
   mDeviceNotifyHandle = InputDeviceUtils::RegisterNotification(mWnd);
 
   // If mDefaultScale is set before mWnd has been set, it will have the scale of

xpfe/appshell/nsAppShellService.cpp

@@ -710,6 +710,10 @@ nsresult nsAppShellService::JustCreateTopWindow(
 
   widgetInitData.mRTL = LocaleService::GetInstance()->IsAppLocaleRTL();
 
+  if (aChromeMask & nsIWebBrowserChrome::CHROME_PRIVATE_WINDOW) {
+    widgetInitData.mIsPrivate = true;
+  }
+
   nsresult rv =
       window->Initialize(parent, center ? aParent : nullptr, aInitialWidth,
                          aInitialHeight, aIsHiddenWindow, widgetInitData);