Additional tweak to make session history a little safer. Bug 305181, r=bryner, sr=darin

docshell/base/nsDocShell.cpp

@@ -2345,6 +2345,37 @@ nsDocShell::AddChild(nsIDocShellTreeItem * aChild)
     return NS_OK;
 }
 
+/* static */ void
+nsDocShell::RemoveChildsSHEntriesFrom(nsISHEntry *aParentEntry,
+                                      nsDocShell *aChildShell)
+{
+    NS_PRECONDITION(aChildShell, "Must have child shell!");
+    
+    nsCOMPtr<nsISHContainer> container(do_QueryInterface(aParentEntry));
+    if (!container)
+        return;
+
+    PRInt32 childCount;
+    container->GetChildCount(&childCount);
+    // Iterate backwards so removals are ok
+    for (PRInt32 i = childCount - 1; i >= 0; --i) {
+        nsCOMPtr<nsISHEntry> childEntry;
+        container->GetChildAt(i, getter_AddRefs(childEntry));
+        if (!childEntry) {
+            // childEntry can be null for valid reasons, for example if the
+            // docshell at index i never loaded anything useful.
+            continue;
+        }
+
+        if (!aChildShell->HasHistoryEntry(childEntry)) {
+            // Not an SHEntry for aChildShell
+            continue;
+        }
+
+        container->RemoveChild(childEntry);
+    }
+}
+
 NS_IMETHODIMP
 nsDocShell::RemoveChild(nsIDocShellTreeItem * aChild)
 {
@@ -2358,6 +2389,11 @@ nsDocShell::RemoveChild(nsIDocShellTreeItem * aChild)
     
     aChild->SetTreeOwner(nsnull);
 
+    // Make sure to remove the child's SHEntry from our SHEntry's child list
+    nsDocShell* childAsDocshell = NS_STATIC_CAST(nsDocShell*, aChild);
+    RemoveChildsSHEntriesFrom(mOSHE, childAsDocshell);
+    RemoveChildsSHEntriesFrom(mLSHE, childAsDocshell);    
+
     return nsDocLoader::AddDocLoaderAsChildOfRoot(childAsDocLoader);
 }
 

docshell/base/nsDocShell.h

@@ -334,6 +334,11 @@ protected:
                                        WalkHistoryEntriesFunc aCallback,
                                        void *aData);
 
+    // For each child of aParentEntry, check whether it belongs to aChildShell.
+    // If it does, remove it from the child list of aParentEntry.
+    static void RemoveChildsSHEntriesFrom(nsISHEntry *aParentEntry,
+                                          nsDocShell *aChildShell);
+
     // overridden from nsDocLoader, this provides more information than the
     // normal OnStateChange with flags STATE_REDIRECTING
     virtual void OnRedirectStateChange(nsIChannel* aOldChannel,