Bug 498648 - Start private browsing while editing a message, cancel, doesn't cancel private browsing; r=bz,jst,ehsan

browser/components/privatebrowsing/src/nsPrivateBrowsingService.js

@@ -124,6 +124,9 @@ PrivateBrowsingService.prototype = {
   // List of view source window URIs for restoring later
   _viewSrcURLs: [],
 
+  // List of nsIXULWindows we are going to be closing during the transition
+  _windowsToClose: [],
+
   // XPCOM registration
   classDescription: "PrivateBrowsing Service",
   contractID: "@mozilla.org/privatebrowsing;1",
@@ -208,12 +211,20 @@ PrivateBrowsingService.prototype = {
           // just in case the only remaining window after setBrowserState is different.
           // it probably shouldn't be with the current sessionstore impl, but we shouldn't
           // rely on behaviour the API doesn't guarantee
-          let browser = this._getBrowserWindow().gBrowser;
+          browserWindow = this._getBrowserWindow();
+          let browser = browserWindow.gBrowser;
 
           // this ensures a clean slate from which to transition into or out of
           // private browsing
           browser.addTab();
+          browser.getBrowserForTab(browser.tabContainer.firstChild).stop();
           browser.removeTab(browser.tabContainer.firstChild);
+          browserWindow.getInterface(Ci.nsIWebNavigation)
+                       .QueryInterface(Ci.nsIDocShellTreeItem)
+                       .treeOwner
+                       .QueryInterface(Ci.nsIInterfaceRequestor)
+                       .getInterface(Ci.nsIXULWindow)
+                       .docShell.contentViewer.resetCloseWindow();
         }
       }
     }
@@ -298,6 +309,20 @@ PrivateBrowsingService.prototype = {
            getMostRecentWindow("navigator:browser");
   },
 
+  _ensureCanCloseWindows: function PBS__ensureCanCloseWindows() {
+    let windowMediator = Cc["@mozilla.org/appshell/window-mediator;1"].
+                         getService(Ci.nsIWindowMediator);
+    let windowsEnum = windowMediator.getXULWindowEnumerator("navigator:browser");
+
+    while (windowsEnum.hasMoreElements()) {
+      let win = windowsEnum.getNext().QueryInterface(Ci.nsIXULWindow);
+      if (win.docShell.contentViewer.permitUnload(true))
+        this._windowsToClose.push(win);
+      else
+        throw Cr.NS_ERROR_ABORT;
+    }
+  },
+
   _closePageInfoWindows: function PBS__closePageInfoWindows() {
     let pageInfoEnum = Cc["@mozilla.org/appshell/window-mediator;1"].
                        getService(Ci.nsIWindowMediator).
@@ -401,6 +426,8 @@ PrivateBrowsingService.prototype = {
             return;
         }
 
+        this._ensureCanCloseWindows();
+
         this._autoStarted = this._prefs.getBoolPref("browser.privatebrowsing.autostart");
         this._inPrivateBrowsing = val != false;
 
@@ -422,9 +449,16 @@ PrivateBrowsingService.prototype = {
         this._onAfterPrivateBrowsingModeChange();
       }
     } catch (ex) {
-      Cu.reportError("Exception thrown while processing the " +
+      // We aborted the transition to/from private browsing, we must restore the
-        "private browsing mode change request: " + ex.toString());
+      // beforeunload handling on all the windows for which we switched it off.
+      for (let i = 0; i < this._windowsToClose.length; i++)
+        this._windowsToClose[i].docShell.contentViewer.resetCloseWindow();
+      // We don't log an error when the transition is canceled from beforeunload
+      if (ex != Cr.NS_ERROR_ABORT)
+        Cu.reportError("Exception thrown while processing the " +
+          "private browsing mode change request: " + ex.toString());
     } finally {
+      this._windowsToClose = [];
       this._alreadyChangingMode = false;
     }
   },

content/html/document/src/nsHTMLDocument.cpp

@@ -1870,7 +1870,7 @@ nsHTMLDocument::OpenCommon(const nsACString& aContentType, PRBool aReplace)
 
     if (cv) {
       PRBool okToUnload;
-      rv = cv->PermitUnload(&okToUnload);
+      rv = cv->PermitUnload(PR_FALSE, &okToUnload);
 
       if (NS_SUCCEEDED(rv) && !okToUnload) {
         // We don't want to unload, so stop here, but don't throw an

docshell/base/nsDocShell.cpp

@@ -6149,7 +6149,7 @@ nsDocShell::CreateAboutBlankContentViewer(nsIPrincipal* aPrincipal,
     // in the current document.
 
     PRBool okToUnload;
-    rv = mContentViewer->PermitUnload(&okToUnload);
+    rv = mContentViewer->PermitUnload(PR_FALSE, &okToUnload);
 
     if (NS_SUCCEEDED(rv) && !okToUnload) {
       // The user chose not to unload the page, interrupt the load.
@@ -7939,7 +7939,7 @@ nsDocShell::InternalLoad(nsIURI * aURI,
     // protocol handler deals with this for javascript: URLs.
     if (!bIsJavascript && mContentViewer) {
         PRBool okToUnload;
-        rv = mContentViewer->PermitUnload(&okToUnload);
+        rv = mContentViewer->PermitUnload(PR_FALSE, &okToUnload);
 
         if (NS_SUCCEEDED(rv) && !okToUnload) {
             // The user chose not to unload the page, interrupt the

docshell/base/nsIContentViewer.idl

@@ -13,7 +13,7 @@ struct nsIntRect;
 [ptr] native nsIWidgetPtr(nsIWidget);
 [ref] native nsIntRectRef(nsIntRect);
 
-[scriptable, uuid(c9aba5da-7d8b-46a8-87cd-9ab7e16480b8)]
+[scriptable, uuid(08665a60-b398-11de-8a39-0800200c9a66)]
 interface nsIContentViewer : nsISupports
 {
 
@@ -24,7 +24,27 @@ interface nsIContentViewer : nsISupports
 
   void loadStart(in nsISupports aDoc);
   void loadComplete(in unsigned long aStatus);
-  boolean permitUnload();
+
+  /**
+   * Checks if the document wants to prevent unloading by firing beforeunload on
+   * the document, and if it does, prompts the user. The result is returned.
+   *
+   * @param aCallerClosesWindow indicates that the current caller will close the
+   *        window. If the method returns true, all subsequent calls will be
+   *        ignored.
+   */
+  boolean permitUnload([optional] in boolean aCallerClosesWindow);
+
+  /**
+   * Works in tandem with permitUnload, if the caller decides not to close the
+   * window it indicated it will, it is the caller's responsibility to reset
+   * that with this method.
+   *
+   * @Note this method is only meant to be called on documents for which the
+   *  caller has indicated that it will close the window. If that is not the case
+   *  the behavior of this method is undefined.
+   */
+  void resetCloseWindow();
   void pageHide(in boolean isUnload);
 
   /**

dom/base/nsGlobalWindow.cpp

@@ -5540,7 +5540,7 @@ nsGlobalWindow::Close()
   if (!mInClose && !mIsClosed && cv) {
     PRBool canClose;
 
-    rv = cv->PermitUnload(&canClose);
+    rv = cv->PermitUnload(PR_FALSE, &canClose);
     if (NS_SUCCEEDED(rv) && !canClose)
       return NS_OK;
 

dom/src/jsurl/nsJSProtocolHandler.cpp

@@ -784,7 +784,7 @@ nsJSChannel::EvaluateScript()
             if (cv) {
                 PRBool okToUnload;
 
-                if (NS_SUCCEEDED(cv->PermitUnload(&okToUnload)) &&
+                if (NS_SUCCEEDED(cv->PermitUnload(PR_FALSE, &okToUnload)) &&
                     !okToUnload) {
                     // The user didn't want to unload the current
                     // page, translate this into an undefined

layout/base/nsDocumentViewer.cpp

@@ -497,6 +497,7 @@ protected:
   nsCString mPrevDocCharacterSet;
   
   PRPackedBool mIsPageMode;
+  PRPackedBool mCallerIsClosingWindow;
 
 };
 
@@ -528,6 +529,7 @@ void DocumentViewerImpl::PrepareToStartLoad()
   mStopped          = PR_FALSE;
   mLoaded           = PR_FALSE;
   mDeferredWindowClose = PR_FALSE;
+  mCallerIsClosingWindow = PR_FALSE;
 
 #ifdef NS_PRINTING
   mPrintIsPending        = PR_FALSE;
@@ -1083,11 +1085,11 @@ DocumentViewerImpl::LoadComplete(nsresult aStatus)
 }
 
 NS_IMETHODIMP
-DocumentViewerImpl::PermitUnload(PRBool *aPermitUnload)
+DocumentViewerImpl::PermitUnload(PRBool aCallerClosesWindow, PRBool *aPermitUnload)
 {
   *aPermitUnload = PR_TRUE;
 
-  if (!mDocument || mInPermitUnload) {
+  if (!mDocument || mInPermitUnload || mCallerIsClosingWindow) {
     return NS_OK;
   }
 
@@ -1193,15 +1195,47 @@ DocumentViewerImpl::PermitUnload(PRBool *aPermitUnload)
 
       if (docShell) {
         nsCOMPtr<nsIContentViewer> cv;
-      docShell->GetContentViewer(getter_AddRefs(cv));
+        docShell->GetContentViewer(getter_AddRefs(cv));
 
-      if (cv) {
+        if (cv) {
-        cv->PermitUnload(aPermitUnload);
+          cv->PermitUnload(aCallerClosesWindow, aPermitUnload);
         }
       }
     }
   }
 
+  if (aCallerClosesWindow && *aPermitUnload)
+    mCallerIsClosingWindow = PR_TRUE;
+
+  return NS_OK;
+}
+
+NS_IMETHODIMP
+DocumentViewerImpl::ResetCloseWindow()
+{
+  mCallerIsClosingWindow = PR_FALSE;
+
+  nsCOMPtr<nsIDocShellTreeNode> docShellNode(do_QueryReferent(mContainer));
+  if (docShellNode) {
+    PRInt32 childCount;
+    docShellNode->GetChildCount(&childCount);
+
+    for (PRInt32 i = 0; i < childCount; ++i) {
+      nsCOMPtr<nsIDocShellTreeItem> item;
+      docShellNode->GetChildAt(i, getter_AddRefs(item));
+
+      nsCOMPtr<nsIDocShell> docShell(do_QueryInterface(item));
+
+      if (docShell) {
+        nsCOMPtr<nsIContentViewer> cv;
+        docShell->GetContentViewer(getter_AddRefs(cv));
+
+        if (cv) {
+          cv->ResetCloseWindow();
+        }
+      }
+    }
+  }
   return NS_OK;
 }