Bug 836951 - Add mDocShell to nsSecureBrowserUIImpl. Traverse up to the root doc shell. r=bsmith,smaug

docshell/base/nsDocShell.cpp

@@ -2508,6 +2508,7 @@ NS_IMETHODIMP
 nsDocShell::SetSecurityUI(nsISecureBrowserUI *aSecurityUI)
 {
     mSecurityUI = aSecurityUI;
+    mSecurityUI->SetDocShell(this);
     return NS_OK;
 }
 

dom/ipc/TabParent.cpp

@@ -30,6 +30,7 @@
 #include "nsFocusManager.h"
 #include "nsFrameLoader.h"
 #include "nsIContent.h"
+#include "nsIDocShell.h"
 #include "nsIDOMApplicationRegistry.h"
 #include "nsIDOMElement.h"
 #include "nsIDOMEvent.h"
@@ -501,6 +502,14 @@ TabParent::GetState(uint32_t *aState)
   return NS_OK;
 }
 
+NS_IMETHODIMP
+TabParent::SetDocShell(nsIDocShell *aDocShell)
+{
+  NS_ENSURE_ARG(aDocShell);
+  NS_WARNING("No mDocShell member in TabParent so there is no docShell to set");
+  return NS_OK;
+}
+
 NS_IMETHODIMP
 TabParent::GetTooltipText(nsAString & aTooltipText)
 {

netwerk/base/public/nsISecureBrowserUI.idl

@@ -8,11 +8,13 @@
 
 interface nsIDOMWindow;
 interface nsIDOMElement;
+interface nsIDocShell;
 
-[scriptable, uuid(081e31e0-a144-11d3-8c7c-00609792278c)]
+[scriptable, uuid(e97e5688-add2-4a1d-acae-396d7702e382)]
 interface nsISecureBrowserUI : nsISupports
 {
     void init(in nsIDOMWindow window);
+    void setDocShell(in nsIDocShell docShell);
 
     readonly attribute unsigned long state;
     readonly attribute AString tooltipText;

security/manager/boot/src/nsSecureBrowserUIImpl.cpp

@@ -19,6 +19,7 @@
 #include "nsIObserverService.h"
 #include "nsCURILoader.h"
 #include "nsIDocShell.h"
+#include "nsIDocShellTreeItem.h"
 #include "nsIDocument.h"
 #include "nsIPrincipal.h"
 #include "nsIDOMElement.h"
@@ -286,12 +287,21 @@ nsSecureBrowserUIImpl::MapInternalToExternalState(uint32_t* aState, lockIconStat
   if (ev && (*aState & STATE_IS_SECURE))
     *aState |= nsIWebProgressListener::STATE_IDENTITY_EV_TOPLEVEL;
 
-  nsCOMPtr<nsPIDOMWindow> piwin = do_QueryReferent(mWindow);
-  if (!piwin)
+  nsCOMPtr<nsIDocShell> docShell = do_QueryReferent(mDocShell);
+  if (!docShell)
     return NS_OK;
 
-  nsIDocShell* docShell = piwin->GetDocShell();
-  MOZ_ASSERT(docShell);
+  int32_t docShellType;
+  // For content docShell's, the mixed content security state is set on the root docShell.
+  if (NS_SUCCEEDED(docShell->GetItemType(&docShellType)) && docShellType == nsIDocShellTreeItem::typeContent) {
+    nsCOMPtr<nsIDocShellTreeItem> docShellTreeItem(do_QueryInterface(docShell));
+    nsCOMPtr<nsIDocShellTreeItem> sameTypeRoot;
+    docShellTreeItem->GetSameTypeRootTreeItem(getter_AddRefs(sameTypeRoot));
+    NS_ASSERTION(sameTypeRoot, "No document shell root tree item from document shell tree item!");
+    docShell = do_QueryInterface(sameTypeRoot);
+    if (!docShell)
+      return NS_OK;
+  }
 
   // Has a Mixed Content Load initiated in nsMixedContentBlocker?
   // If so, the state should be broken; overriding the previous state
@@ -319,6 +329,14 @@ nsSecureBrowserUIImpl::MapInternalToExternalState(uint32_t* aState, lockIconStat
   return NS_OK;
 }
 
+NS_IMETHODIMP
+nsSecureBrowserUIImpl::SetDocShell(nsIDocShell *aDocShell)
+{
+  nsresult rv;
+  mDocShell = do_GetWeakReference(aDocShell, &rv);
+  return rv;
+}
+
 NS_IMETHODIMP
 nsSecureBrowserUIImpl::GetTooltipText(nsAString& aText)
 {

security/manager/boot/src/nsSecureBrowserUIImpl.h

@@ -17,6 +17,7 @@
 #include "nsIStringBundle.h"
 #include "nsISecureBrowserUI.h"
 #include "nsIDocShell.h"
+#include "nsIDocShellTreeItem.h"
 #include "nsIWebProgressListener.h"
 #include "nsIFormSubmitObserver.h"
 #include "nsIURI.h"
@@ -66,6 +67,7 @@ protected:
   mozilla::ReentrantMonitor mReentrantMonitor;
   
   nsWeakPtr mWindow;
+  nsWeakPtr mDocShell;
   nsCOMPtr<nsINetUtil> mIOService;
   nsCOMPtr<nsIStringBundle> mStringBundle;
   nsCOMPtr<nsIURI> mCurrentURI;